Calculating partial trace of array in NumPy

Question

A simulation I'm doing requires me to calculate the partial trace of a large density matrix. I am trying to calculate it using tools from numpy, but my code seems to be having some problems. For background, let me explain the arrays I am interested in a little more, and the way I'm defining the partial trace. Then, I will give the code I have and the errors I am getting.

First, the partial trace. If I have a tensor product of vector spaces

$$ V = \prod_i V_i$$

and a linear operator $T: V \to V$, then given a basis I can store all the information about $T$ as a multidimensional array $T_{out_1,..,out_n,in_1,..,in_n}$ where $T_{out_1,..,out_n,in_1,..,in_n}$ is the $v_{out_1} \otimes v_{out_2}\otimes \dots \otimes v_{out_n}$ component of $T(v_{in_1} \otimes \dots \otimes v_{in_n} )$. (Compare this to a matrix in a basis where $M_{ij}$ is the $i^{th}$ component of $M(v_j)$.)

The partial trace over some of these indices is a new operator given as a multidimensional array defined by $\widetilde{T}_{kept_{out},kept_{in}} = \sum_{traced} T_{keep_{out},traced, keep_{in},traced}$, where $kept_{in}$, $kept_{out}$, and $traced$ are all multi indices referring to a subset of basis states. The sum is taken over all combinations of indices in the traced set.

My code for computing this in numpy is:

def trace_index(array,i):
"""
Given an array and an index i, traces out the index.
"""
    n = len(array.shape)
    tot = list(range(n))
    tot[n//2+i-1] = i-1
    return np.einsum(array,list(tot))

def partial_trace(array,indices):
"""
Given an array and a list of indices to trace out, traces out the indices
"""

    in_sorted = sorted(indices,reverse=True) 
    cur_trace = array
    for i in in_sorted:
        cur_trace = trace_index(cur_trace,i)
    return cur_trace

I trace them in the descending order eg. 5,4,... because then I can apply trace_index to the indices one at a time. If I trace index 5 and then index 4, index 4 is still index 4 after tracing index 5. If I do it the other way, after I traced index 4, there is no index 5.

This code seems to work well for small cases, but for larger ones I

ValueError: invalid subscript '}' in einstein sum subscripts string, subscripts must be letters

So, my question is this: Is there a better way to do what I am trying to do than what I am doing?

Answer 1

I didn't follow exactly your notation so I can't say for sure what this would look like for your example, but two suggestions:

• if you really want your life made simple, check out qutip, the Quantum Toolbox in Python. It has classes specifically for quantum operators and state vectors, supports tensor product of operators in different spaces and keeps track of the original shape; as a result, it can quite easily carry out partial traces as well.
• if you want to go on with numpy (and you may have reason to: in at least one occasion in the past I have found out that qutip mysteriously killed performance in a computation when changing basis to an operator), then look into .reshape. Specifically, if your tensor space is of size $n\otimes m$, you'll want to use A.reshape((n, m, n, m)) for an operator A. This should preserve order and split the different parts along different axes, which then you can sum or trace along as you need.

Answer 2

I don't have enough reputation to comment, but I wanted to give a basic complement of gIS's answer.

Simply put, if we have the decomposition over $V=V_1\otimes V_2$, which we represent with, say, the reshaping [2,4,2,4] why do we einsum over the first and third indices for trace over $V_1$ and second and fourth for trace over $V_2$? Maybe this is obvious for most of you, but it wasn't for me so here is how I reasoned about it, in case someone else also has doubts about how to connect reshape with the tensor product.

Take a matrix $M$ which is as a linear function over $V$, $M:V\to V$. Then given a basis $\{\vert a \rangle \}$ we get the coefficients of the matrix as $\langle a \vert M \vert b \rangle=M_{ab}$. The full matrix can be written as

$$ M = \sum_{ab} M_{ab} \vert a \rangle \otimes \langle b \vert = \sum_{ab} M_{ab} \vert a \rangle \langle b \vert$$

where $\vert a \rangle$ is a vector on $V$ (basically a column vector) and $\langle b \vert$ is a member of the dual $V^*$, basically a row vector. We omit the tensor product where it's ovbious.

Now we insert the decomposition $V=V_1 \otimes V_2$ as ``splitting'' the basis vectors $\vert a \rangle \to \vert i \rangle \otimes \vert j \rangle$, $\langle b \vert \to \langle k \vert \otimes \langle l \vert$ where $0< i,k<\text{dim}V_1$,$0<j,l<\text{dim}V_2$. We also substitute the index a with the pair i,j and b with k,l This gives

$$ M = \sum_{ijkl} M_{i,j,k,l} \vert i \rangle_{V_1} \otimes \vert j \rangle_{V_2} \otimes \langle k \vert_{V_1} \otimes \langle l \vert_{V_2} = \sum_{ijkl} M_{i,j,k,l} \vert i \rangle \langle k \vert \otimes \vert j \rangle \langle l \vert $$

It's obvious from this that if we want to trace over $V_1$ we need to sum over the first and third, and if we want to trace out $V_2$ then we sum over the second and fourth index.

Answer 3

If you use qutip, partial trace operations are already built-in. For example, here is how you can compute the partial trace of a random density matrix over three qubits (that is, an hermitian, trace-1 matrix living in a tensor product space of dimensions $(2, 2, 2)$), and then trace out the last space:

import qutip
dm = qutip.rand_dm_hs(8, dims=[[2] * 3] * 2)
dm.ptrace([0, 1])

On the other hand, there are circumstances in which you may not want to use qutip. Say therefore that you have a matrix stored in a numpy array, which represents a tensor in a tensor product space width dimensions $(d_1,...,d_n)$. The matrix has therefore shape $(d_1\cdots d_n, d_1\cdots d_n)$.

A convenient way to compute the partial trace is, how the other answer suggests, to first reshape the matrix and then to sum the appropriate axes.

For example, suppose you have a matrix representing an object in a space $V= V_1\otimes V_2$, with $V_1$ of dimension $2$ and $V_2$ of dimension $4$. Here is how partial trace the first space:

import numpy as np
import qutip
# generate test matrix (using qutip for convenience)
dm = qutip.rand_dm_hs(8, dims=[[2, 4]] * 2).full()
# reshape to do the partial trace easily using np.einsum
reshaped_dm = dm.reshape([2, 4, 2, 4])
# compute the partial trace
reduced_dm = np.einsum('ijik->jk', reshaped_dm)

while to partial trace with respect to the second space you just change the last line to reduced_dm = np.einsum('jiki->jk', reshaped_dm).

Here is the same thing as above, including also a consistency check with the partial trace given by qutip:

import numpy as np
import qutip
# generate test matrix (using qutip for convenience)
dm = qutip.rand_dm_hs(8, dims=[[2, 4]] * 2).full()
# reshape to do the partial trace easily using np.einsum
reshaped_dm = dm.reshape([2, 4, 2, 4])
# partial trace the second space
reduced_dm = np.einsum('jiki->jk', reshaped_dm)
# check results with qutip
qutip_dm = qutip.Qobj(dm, dims=[[2, 4]] * 2)
reduced_dm_via_qutip = qutip_dm.ptrace([0]).full()
# check consistency of results
np.allclose(reduced_dm, reduced_dm_via_qutip)