I have a python implementation of calculating the partial trace over select dimensions.
def partial_trace(rho, keep, dims): """Calculate the partial trace Parameters ---------- ρ : 2D array Matrix to trace keep : array An array of indices of the spaces to keep after being traced. For instance, if the space is A x B x C x D and we want to trace out B and D, keep = [0,2] dims : array An array of the dimensions of each space. For instance, if the space is A x B x C x D, dims = [dim_A, dim_B, dim_C, dim_D] Returns ------- ρ_a : 2D array Traced matrix """ dims = np.asarray(dims) N = dims.size # Indices to trace rest = np.delete(np.arange(N), keep) # Reshape into tensor rho_a = rho.reshape(np.tile(dims, 2)) # Trace indices for i,d in enumerate(rest[::-1]): rho_a = np.trace(rho_a, axis1=d, axis2=N-i+d) # Reshape into matrix N = np.prod(dims[keep]) return rho_a.reshape(N,N)
I would like to calculate the partial trace of $|u\rangle \langle u|$ but I run into an insufficient memory error when I try to construct the outer product. Instead, I would like to directly calculate the partial trace from $|u \rangle$.
How would I do this in Python?