I want to numerically compute something of the form $x=\langle\psi| A\otimes I |\psi\rangle$ where $I$ is the identity.

Of course I can construct $A\otimes I$ and take the inner product, but I would like to save memory and compute $x$ without having to construct this large operator.

What would be an efficient way to compute $x$ using simple numpy operations (reshape, einsum, trace etc) and without constructing the whole $ A\otimes I$ operator?


1 Answer 1


I don't know if it's the fastest, but you can use the following subroutine for applying an operation to a subsystem of your state. This code will take an operation that's the correct shape for acting on a subsystem of $|\psi\rangle$ and apply it to that subsystem without having to tensor up to the size of the state space that $|\psi\rangle$ lives in:

def mul_on_subsystem(op, state, target_subsystem):
    # `state` has a tensor shape, e.g. (dim_A, dim_B, dim_C)

    # make sure `op` is correctly shaped to be applied to the target subsystem!
    target_dim = state.shape[target_subsystem]
    assert op.shape == (target_dim, target_dim)

    # (1) tensordot op onto the target axis of the state
    # this applies axis 1 of `op` to the target axis of `state`
    out = np.tensordot(op, state, axes=([1], [target_subsystem])) 

    # (2) rearrange the axes that get scrambled by tensordot
    out = np.moveaxis(out, 0, target_subsystem)
    return out

Once you've applied $A$ to $|\psi\rangle$ in the way you like, you can compute the inner product in the obvious way,

np.dot(state.conj().flatten(), mul_on_subsystem(local_operation, state, target_subsystem).flatten())

This technique also generalizes nicely to density operators, albeit with some very annoying shuffling of axes.

Timing/validation/sample code:

# declare your subsystem sizes; these are "2" if its qubits...
dim_a = 25 # axis 0
dim_b = 25 # axis 1
dim_c = 35 # axis 2

# reshape your state to get it into tensor form (dim_a, dim_b, dim_c)
state = np.random.rand(dim_a, dim_b, dim_c)
state = state / np.linalg.norm(state)

operation_A = np.random.rand(dim_a, dim_a)
operation_B = np.random.rand(dim_b, dim_b)
operation_C = np.random.rand(dim_c, dim_c)

# being generous and not including the time to construct the tensored operation
AII = np.kron(operation_A, np.kron(np.eye(dim_b), np.eye(dim_c)))
IBI = np.kron(np.eye(dim_a), np.kron(operation_B, np.eye(dim_c)))
IIC = np.kron(np.eye(dim_a), np.kron(np.eye(dim_b), operation_C))

for i in range(3):
    # using mul_on_subsystem
    local_operation = [operation_A, operation_B, operation_C][i]
    t0 = time.time()
    fast = mul_on_subsystem(local_operation, state, i)
    time_fast = time.time() - t0

    # using tensored-up matrices
    tensored_operation = [AII, IBI, IIC][i]
    t0 = time.time()
    slow = np.dot(tensored_operation, state.flatten())
    time_slow = time.time() - t0

    system = ["A", "B", "C"][i]
    print(f"system {system} fast: {time_fast:.5f} s, slow: {time_slow:.5f} s")
    np.testing.assert_allclose(fast.flatten(), slow)

Output on my laptop

system A fast: 0.02700 s, slow: 0.47000 s
system B fast: 0.00100 s, slow: 0.21201 s
system C fast: 0.00000 s, slow: 0.32400 s

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.