How to efficiently compute the partial trace of a matrix?

Ezra: 2 days ago

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[
     {iter[k], lengths[[k]]},
     {k, indices}
    {iter, iterators}
indicesToIndex[indices_List, lengths_List] := 1 + Total@MapIndexed[
     #1 Times @@ lengths[[First@#2 + 1 ;;]] &,
     indices - 1

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

              j /@ Range@Length@lengths /. j[n_] :> i[n], lengths]
            ] // Sow,
         ] // 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

enter image description here

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 ( (

] //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:

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

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}],
]; //AbsoluteTiming
r2 = Nest[Partition[#, {10,10}]&, data, 2]; //AbsoluteTiming

{0.005594, Null}

{0.0286, Null}


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]
    matrixInTPForm = Transpose[
      ArrayReshape[matrix, Join[lengths, lengths]],
      Join @@ Transpose@Partition[Range[2 Length@lengths], 2]
     {2 # - 1, 2 #} & /@ indicesToTrace
     Partition[Range[2 Length@lengths - 2 Length@indicesToTrace], 2]