14
$\begingroup$

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

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.

Improve this question
$\endgroup$

2 Answers 2

Reset to default
14
$\begingroup$

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]
    ]
   ]
  ]
Improve this answer
$\endgroup$
5
  • 1
    $\begingroup$ 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 $\endgroup$
    – glS
    Commented Mar 13, 2017 at 23:11
  • $\begingroup$ 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} ]? $\endgroup$
    – glS
    Commented Mar 14, 2017 at 12:57
  • $\begingroup$ @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}? $\endgroup$
    – Carl Woll
    Commented Mar 14, 2017 at 14:13
  • $\begingroup$ 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). $\endgroup$
    – glS
    Commented Mar 14, 2017 at 14:43
  • 1
    $\begingroup$ 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). $\endgroup$
    – glS
    Commented Mar 14, 2017 at 17:37
2
$\begingroup$

A less general solution for the case of spin chains and sparse arrays

I faced the problem of finding reduced density matrices for large spin systems, where the Hamiltonian is represented as a $2^N$ by $2^N$ matrix and state vectors are represented as vectors of length $2^N$. You could also have $D$ states per site for a state space of $D^N$, but I'll use $D=2$ as an example.

I found a way to tackle the problem using SparseArrays and Kronecker products. This method is very fast, but I have a feeling working with TensorProduct and TensorContract from the get-go is still faster. This is just a second way to tackle the problem.

The idea is that we can evaluate the density matrix elements as follows: \begin{align*} \langle n|\rho_A|m\rangle &= \langle n|\operatorname{tr}_B\left( \rho_{AB}\right) |m\rangle\\ &=\operatorname{tr}_B\left(\langle n|\rho_{ab}|m\rangle\right)\\ &=\operatorname{tr}_{AB}\left(\rho_{ab}|m\rangle\langle n|\right) \end{align*}

So, if you have $D=2$ and $N=$nsites, the partial trace of a single site after tracing over everything else can be calculated as follows:

id = SparseArray@IdentityMatrix[2]; 
(* The four matrices |0><0|, |1><1|, |1><0| and |0><1|. *)

projector[0, 0] = SparseArray@{{1, 0}, {0, 0}};
projector[1, 1] = SparseArray@{{0, 0}, {0, 1}};
projector[1, 0] = SparseArray@{{0, 0}, {1, 0}};
projector[0, 1] = Transpose[projector[1, 0]];
(* Sparse array kronecker product *)

myProjector[s1_, s2_, site_] := 
   KroneckerProduct @@ 
    Table[If[i == site, projector[s1, s2], id], {i, 1, nsites}];
myPartialTrace[rho_, site_] := 
  Table[Tr[rho.myProjector[s1, s2, site]], {s1, 0, 1}, {s2, 0, 1}];

The sparse kronecker product is calculated in no time at all (even for N=15), so all of the time is being spend inside the trace and matrix multiplication routines.

Here is a full sample implementation showing that I get the same results as partialTrace. I threw in an AbsoluteTiming call, but because I convert the large matrix into a dense matrix with Normal before passing it to partialTrace, the time comparisons aren't fair (I would be comparing a sparse problem to a dense problem).

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]]]];

id = SparseArray@IdentityMatrix[2]; 
(* The four matrices |0><0|, |1><1|, |1><0| and |0><1|. *)

projector[0, 0] = SparseArray@{{1, 0}, {0, 0}};
 projector[1, 1] = SparseArray@{{0, 0}, {0, 1}};
projector[1, 0] = SparseArray@{{0, 0}, {1, 0}};
projector[0, 1] = Transpose[projector[1, 0]];

timings = Table[
   nupspins = Floor[nsites/2];
   Clear[myProjector];
   (* Sparse array kronecker product *)

   myProjector[s1_, s2_, site_] := 
    KroneckerProduct @@ 
     Table[If[i == site, projector[s1, s2], id], {i, 1, nsites}];
   myPartialTrace[rho_, site_] := 
    Table[Tr[rho.myProjector[s1, s2, site]], {s1, 0, 1}, {s2, 0, 1}];

   (* Create an interesting state to work with. 
   This one is a random superposition of all states such that the \
total number of up spins is 1. Sites are indexed in binary, 
   so basis states with nsites=10 and nupspins=
   5 would include 0b0000011111, 0b1010101010, etc. *)

   normalize[vec_] := vec/Sqrt[vec.vec];
   myArray = 
    normalize[
     SparseArray[
      ReleaseHold[
       Table[If[Total[IntegerDigits[i, 2]] == nupspins, 
         i + 1 -> RandomReal[], HoldForm[Sequence[]]], {i, 0, 
         2^nsites - 1}]], {2^nsites}]];

   {AbsoluteTiming[
     partialTrace[Normal@Outer[Times, myArray, myArray], 
      Table[2, {i, 1, nsites}], {1}]],
    AbsoluteTiming[
     myPartialTrace[Outer[Times, myArray, myArray], 1]]}
   , {nsites, 2, 14}];

(* ensure that for all the calculated partial traces, the difference \
is essentially zero. *)
And @@ 
 Thread[Chop[Flatten[timings[[All, 1, 2]] - timings[[All, 2, 2]]]] == 
   0]
Improve this answer
$\endgroup$
1
  • 2
    $\begingroup$ Looks nice. The only tip I would give is that SparseArray @ IdentityMatrix[2] is equivalent to IdentityMatrix[2, SparseArray]. $\endgroup$
    – J. M.'s missing motivation ♦
    Commented May 13, 2020 at 10:34

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.