# How to efficiently compute the partial trace of a matrix?

How can we efficiently compute the partial trace of a matrix with Mathematica?

There is some Mathematica code around to compute this, but most of it seems outdates and not very well written. See for example this code on the Wolfram Library Archive. Only one question seems to have been asked here about this problem, but it was about a very special case.

Here is my solution to the problem:

makeIterators[iterators_, lengths_, indices_] := Join @@ Table[
Table[
{iter[k], lengths[[k]]},
{k, indices}
],
{iter, iterators}
];
indicesToIndex[indices_List, lengths_List] := 1 + Total@MapIndexed[
#1 Times @@ lengths[[First@#2 + 1 ;;]] &,
indices - 1
];

ClearAll[partialTrace];
partialTrace[matrix_, lengths_, indicesToKeep_] := Module[{i, j},
With[{indicesToTrace =
Complement[Range@Length@lengths, indicesToKeep]},
With[{
iteratorsInFinalMatrix =
Sequence @@ makeIterators[{i, j}, lengths, indicesToKeep],
iteratorsToTrace =
Sequence @@ makeIterators[{i}, lengths, indicesToTrace]
},
Do[
Plus @@ Flatten @ Table[
matrix[[
indicesToIndex[i /@ Range@Length@lengths, lengths],

indicesToIndex[
j /@ Range@Length@lengths /. j[n_] :> i[n], lengths]
]],
iteratorsToTrace
] // Sow,
iteratorsInFinalMatrix
] // Reap // Last // First //
Partition[#, Times @@ lengths[[indicesToKeep]]] &
]
]
]


This solution basically replicates the steps one would naturally do when computing the partial trace by hand, but I don't like it very much (in particular having to programmatically create the iterators for the Table and Do). To name a few problems, it cannot be compiled nor parallelized.

Here is an example of its operation:

testMatrix = KroneckerProduct[Array[a, {2, 2}], Array[b, {2, 2}], Array[c, {2, 2}]];
partialTrace[testMatrix, {2, 2, 2}, {1, 3}] // TraditionalForm


which gives

While an algorithm working for symbolic inputs is nice, I'm mostly interested in a function working efficiently for (potentially big) numerical matrices.

To be clear: the problem is that of finding an algorithm to compute the partial trace of a matrix. That is, the inputs will be one matrix, the set of dimensions of the bases, and the dimensions to keep (or those to trace away, equivalently). A solution working on the nested structure given by TensorProduct is a valid answer only as long as one also provides a mean to convert back and forth between the regular matrix representation and the TensorProduct representation.

Looks like you can just use TensorContract/TensorProduct:

TensorContract[
TensorProduct[Array[a,{2,2}],Array[b,{2,2}],Array[c,{2,2}]],
{3,4}
] //ArrayFlatten //TeXForm


$\tiny \begin{pmatrix} a(1,1) b(1,1) c(1,1)+a(1,1) b(2,2) c(1,1) & a(1,1) b(1,1) c(1,2)+a(1,1) b(2,2) c(1,2) & a(1,2) b(1,1) c(1,1)+a(1,2) b(2,2) c(1,1) & a(1,2) b(1,1) c(1,2)+a(1,2) b(2,2) c(1,2) \\ a(1,1) b(1,1) c(2,1)+a(1,1) b(2,2) c(2,1) & a(1,1) b(1,1) c(2,2)+a(1,1) b(2,2) c(2,2) & a(1,2) b(1,1) c(2,1)+a(1,2) b(2,2) c(2,1) & a(1,2) b(1,1) c(2,2)+a(1,2) b(2,2) c(2,2) \\ a(2,1) b(1,1) c(1,1)+a(2,1) b(2,2) c(1,1) & a(2,1) b(1,1) c(1,2)+a(2,1) b(2,2) c(1,2) & a(2,2) b(1,1) c(1,1)+a(2,2) b(2,2) c(1,1) & a(2,2) b(1,1) c(1,2)+a(2,2) b(2,2) c(1,2) \\ a(2,1) b(1,1) c(2,1)+a(2,1) b(2,2) c(2,1) & a(2,1) b(1,1) c(2,2)+a(2,1) b(2,2) c(2,2) & a(2,2) b(1,1) c(2,1)+a(2,2) b(2,2) c(2,1) & a(2,2) b(1,1) c(2,2)+a(2,2) b(2,2) c(2,2) \\ \end{pmatrix}$

You can look at Ways to compute inner products of tensors and in particular this answer for an efficient version of this approach.

The OP wanted a method to convert a KroneckerProduct representation to a TensorProduct representation so that TensorContract could be used. For this particular example, you could use Nest and Partition to do this. Here I use this method on a random 1000 x 1000 matrix:

TensorContract[
Nest[Partition[#, {10, 10}]&, RandomReal[1, {1000, 1000}], 2],
{3,4}
]


Another possibility is to use ArrayReshape, although a little massaging is necessary for this approach:

r1 = Transpose[
ArrayReshape[data, {10,10,10,10,10,10}],
{1,3,5,2,4,6}
]; //AbsoluteTiming
r2 = Nest[Partition[#, {10,10}]&, data, 2]; //AbsoluteTiming
r1===r2


{0.005594, Null}

{0.0286, Null}

True

To convert back you would use ArrayFlatten or Flatten.

To wrap it all up into a function:

partialTrace[matrix_, lengths_, indicesToKeep_] := With[{
indicesToTrace = Complement[Range@Length@lengths, indicesToKeep]
},
With[{
matrixInTPForm = Transpose[
ArrayReshape[matrix, Join[lengths, lengths]],
Join @@ Transpose@Partition[Range[2 Length@lengths], 2]
]
},
Flatten[
TensorContract[matrixInTPForm,
{2 # - 1, 2 #} & /@ indicesToTrace
],
Transpose@
Partition[Range[2 Length@lengths - 2 Length@indicesToTrace], 2]
]
]
]

• Not really. The problem is to compute the partial trace given a matrix in the form obtained from a KroneckerProduct operation (in the actual application it will be obtained by other means, but that is the structure it will have). This method can work as long as you provide also the code to convert the matrix to and from the TensorProduct structure. This was actually my initial solution (not the one I posted in the question), but I think this method is not particularly efficient, mostly because of the cost of converting big matrices to and from the nested form given by TensorProduct – glS Mar 13 '17 at 23:11
• for example, how do you compute with your function the partial trace of a randomly generated matrix? Like what I would do with partialTrace[ RandomReal[{0, 1}, {1000, 1000}], {10, 10, 10}, {1, 3} ]? – glS Mar 14 '17 at 12:57
• @gIS I have some ideas on working directly with matrics, but they are not ready yet. One question: Are the indices you're contracting over always of the same matrix? I.e., in my notation, always {3,4} or {1,2}, but not say {1, 4}? – Carl Woll Mar 14 '17 at 14:13
• if I get what you mean, yes. By definition the partial trace involves contracting pairs of indices associated to the same Hilbert space (you can think of it as associated to the same "matrix" only if the matrix you are partial tracing is a simple tensor product of other matrices. In general it will not be, like in the case of my comment above). – glS Mar 14 '17 at 14:43
• I took the liberty of editing your post to wrap up your method into a single function with the same interface I used for mine. Please feel free to revert it back if you don't agree/like my edit. This function is already 2 orders of magnitudes faster than my version for a 1000x1000 matrix (.02s vs 5s on my laptop). – glS Mar 14 '17 at 17:37