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?