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