Why is trace not preserved in tensordot

Question

I am trying to take a tensor product of three density matrices and express it in the product basis. Each of these matrices have trace 1 and theoretically, the product matrix should as well. But doing this in numpy seems to have some unintended effects. Even reshaping the intermediary array to a 2d form gives the same answer.

In [31]: rho1
Out[31]: 
array([[0. , 0. , 0. , 0. , 0. ],
       [0. , 0.1, 0. , 0. , 0. ],
       [0. , 0. , 0.2, 0. , 0. ],
       [0. , 0. , 0. , 0.3, 0. ],
       [0. , 0. , 0. , 0. , 0.4]])

In [32]: np.trace(rho1)
Out[32]: 1.0

In [33]: rho2
Out[33]: 
array([[0.2, 0. , 0. , 0. , 0. ],
       [0. , 0.2, 0. , 0. , 0. ],
       [0. , 0. , 0.2, 0. , 0. ],
       [0. , 0. , 0. , 0.2, 0. ],
       [0. , 0. , 0. , 0. , 0.2]])

In [34]: np.trace(rho2)
Out[34]: 1.0

In [35]: rho3
Out[35]: 
array([[0.5, 0. , 0. , 0. , 0. ],
       [0. , 0.5, 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ]])

In [36]: np.trace(rho3)
Out[36]: 1.0

In [37]: rho = np.tensordot(rho1, np.tensordot(rho2, rho3, axes=0), axes=0)

In [38]: np.trace(rho.reshape(125, 125))
Out[38]: 0.010000000000000002

In [39]: rho = np.tensordot(rho1, np.tensordot(rho2, rho3, axes=0).reshape(25,25), axes=0)

In [40]: np.trace(rho.reshape(125, 125))
Out[40]: 0.010000000000000002

I can't really find any bug in this code, so I feel I am misunderstanding how tensordot and reshape work. But I didn't really get anything from the docs. Can someone help me out in this?

Answer 1

The tensordot docs does say that axes=0 gives a tensor product:

axes = 0 : tensor product :math:`a\otimes b`, 

an outer product of the arrays. But you have to pay attention to how the resulting dimensions are arranged.

But it says:

The shape of the result consists of the non-contracted axes of the
first tensor, followed by the non-contracted axes of the second.

With axes=0, it in effect creates rho2[:,:,None] and rho3[:,None,:] and does the sum-of-products on the normal dot axes.

In [37]: x=np.tensordot(rho2, rho3, 0)
In [38]: x.shape
Out[38]: (5, 5, 5, 5)

Comparing it to einsum with default 'ijkl' output:

In [40]: y=np.einsum('ij,kl',rho2,rho3)
In [41]: y.shape
Out[41]: (5, 5, 5, 5)
In [42]: np.allclose(x,y)
Out[42]: True

or the dot equivalent:

In [44]: z=np.dot(rho2[:,:,None],rho3[:,None,:])
In [45]: np.allclose(x,z)
Out[45]: True

The first 2 dimensions come from rho2.

If what you really want was the kron, we have to first transpose the dimensions, interleaving ones from the 2 arrays:

In [46]: K=np.kron(rho2,rho3)
In [47]: K.shape
Out[47]: (25, 25)
In [48]: np.trace?
In [49]: np.allclose(x.transpose(0,2,1,3).reshape(25,25),K)
Out[49]: True

Testing the trace:

In [50]: K.trace()
Out[50]: 1.0

The same thing with einsum. Note the order of indices:

In [53]: np.einsum('ij,kl->ikjl',rho2,rho3).reshape(25,25).trace()
Out[53]: 1.0

Another way to look at this is extract the 2d diagonal values from the 4d array:

In [60]: j = np.arange(5); i=j[:,None]
In [61]: x[i,i,j,j]
Out[61]: 
array([[0.1, 0.1, 0. , 0. , 0. ],
       [0.1, 0.1, 0. , 0. , 0. ],
       [0.1, 0.1, 0. , 0. , 0. ],
       [0.1, 0.1, 0. , 0. , 0. ],
       [0.1, 0.1, 0. , 0. , 0. ]])
In [62]: _.sum()
Out[62]: 1.0

Simply reshaping x, puts these values on the diagonal:

In [63]: x[i,j,i,j]
Out[63]: 
array([[0.1, 0. , 0. , 0. , 0. ],
       [0. , 0.1, 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ]])

Equivalently with a 2 step trace:

In [75]: x.trace(0,0,1)
Out[75]: 
array([[0.5, 0. , 0. , 0. , 0. ],
       [0. , 0.5, 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. , 0. ]])
In [76]: x.trace(0,0,1).trace()
Out[76]: 1.0

Answer 2

Short answer:

I guess that you're looking for the kronecker tensor product: np.kron():

rho = np.kron(rho1, np.kron(rho2, rho3))

And this time:

np.trace(rho)
Out[22]: 1.0000000000000002

Details:

Why your solution does not work ?

Because np.reshape() do not reshape your array with the "right" dimension order, you could eventually permute some dimension to get the desired result:

rho = np.moveaxis(np.tensordot(rho1, np.moveaxis(np.tensordot(rho1, rho2, axes=0),[0,1,2,3],[0,2,3,1]).reshape((5,5)), axes=0),[0,1,2,3],[0,2,3,1]).reshape((125,125))

Will output the correct result, but since np.kron exist, it is a bit overkill.

Actually you think that np.tensordot(rho1,rho2,axes=0) is equivalent to:

np.einsum('ik,jl',rho1,rho2)

But in fact np.tensordot(rho1,rho2,axes=0) compute:

np.einsum('ij,kl',rho1,rho2)

So another way to get the right answer will be to use:

np.einsum('ik,jl',rho1,rho2).reshape(5,5)

Answer 3

np.tensordot is not a tensor product, but rather a dot product for tensors. The tensor prouct is analoguous to an outer product.