1
$\begingroup$

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?

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

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)
|cite|improve this answer
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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