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

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.

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