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?
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.
# 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
