# Calculate partial trace of an outer product in Python?

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?

• Are the vectors $|u\rangle$ in $R^n$? Aug 15 '18 at 20:01
• The elements are complex Aug 15 '18 at 20:58
• So, $\mathbb{C}^n$? Aug 15 '18 at 21:13
• Yes, $\mathbb{C}^n$. Aug 15 '18 at 21:26
• Isn't in that case: $\mathrm{Tr}(|u\rangle \langle u| ) = \langle u| u\rangle$? Aug 15 '18 at 21:27

I modified my implementation of the partial trace to use einsum.

def partial_trace(rho, keep, dims, optimize=False):
"""Calculate the partial trace

ρ_a = Tr_b(ρ)

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
"""
keep = np.asarray(keep)
dims = np.asarray(dims)
Ndim = dims.size
Nkeep = np.prod(dims[keep])

idx1 = [i for i in range(Ndim)]
idx2 = [Ndim+i if i in keep else i for i in range(Ndim)]
rho_a = rho.reshape(np.tile(dims,2))
rho_a = np.einsum(rho_a, idx1+idx2, optimize=optimize)
return rho_a.reshape(Nkeep, Nkeep)

This can be modified to take a vector as input.

def ptrace_outer(u, keep, dims, optimize=False):
"""Calculate the partial trace of an outer product

ρ_a = Tr_b(|u><u|)

Parameters
----------
u : array
Vector to use for outer product
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
"""
keep = np.asarray(keep)
dims = np.asarray(dims)
Ndim = dims.size
Nkeep = np.prod(dims[keep])

idx1 = [i for i in range(Ndim)]
idx2 = [Ndim+i if i in keep else i for i in range(Ndim)]
u = u.reshape(dims)
rho_a = np.einsum(u, idx1, u.conj(), idx2, optimize=optimize)
return rho_a.reshape(Nkeep, Nkeep)