Question:

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

Ezra: 2 days ago

How can we efficiently compute the partial trace (https://en.wikipedia.org/wiki/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 (http://library.wolfram.com/infocenter/MathSource/8763/#downloads) on the Wolfram Library Archive. Only one question (https://mathematica.stackexchange.com/questions/119148/how-to-compute-the-partial-trace-of-a-4x4-matrix/119152) 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.

Miles: 2 days ago

Looks like you can just use TensorContract (http://reference.wolfram.com/language/ref/TensorContract)/TensorProduct (http://reference.wolfram.com/language/ref/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 (https://mathematica.stackexchange.com/q/17758/45431) and in particular this answer (https://mathematica.stackexchange.com/a/98564/45431) for an efficient version of this approach.

The OP wanted a method to convert a KroneckerProduct (http://reference.wolfram.com/language/ref/KroneckerProduct) representation to a TensorProduct (http://reference.wolfram.com/language/ref/TensorProduct) representation so that TensorContract (http://reference.wolfram.com/language/ref/TensorContract) could be used. For this particular example, you could use Nest (http://reference.wolfram.com/language/ref/Nest) and Partition (http://reference.wolfram.com/language/ref/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 (http://reference.wolfram.com/language/ref/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 (http://reference.wolfram.com/language/ref/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]
]
]
]