pythonnumpyreshapetensordot

Why is trace not preserved in tensordot

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?

Solution

  • 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)