Partial Trace of a Multi-QuDit System
by
Mark Tame
School of Chemistry and Physics, University of KwaZulu-Natal, South Africa.
markstame@gmail.com

Introduction

This code is a generalization of the Multi-qubit partial trace to the case of qudits (d-level systems). It works for both pure and mixed states.
Please see the qubit notebook (available from Mathsource) for more details.
Background: Suppose Alice, Bob and Charlie each have a qudit with the total system described by the density matrix (Subscript[ρ, ABC]). If we would like  to examine more closely the subsystem of qudits shared by just Alice and Bob we need to take the trace:
 (Subscript[ρ, AB])=(Subscript[Tr, C]){(Subscript[ρ, ABC]) }
However for large multi-qudit systems, calculating this by hand becomes quite unmanageable. What the dTraceSystem  function does, is allow you to trace out an arbitrary number of qudits from a large system of n-qudits by getting Mathematica to do the work for you:
dTraceSystem[(Subscript[ρ, sys]), {1,3,6,...}, d] = (Subscript[Tr, [1,3,6,...]]){(Subscript[ρ, sys])}
Where d is the dimension of the qudits and (Subscript[ρ, sys]) is in the computational basis, e.g. for 2 qutrits (d=3) we have:
       
 ({{Basis, <00|, <01|, <02|, <10|, <11|, <12|, <20|, <21|, <22|}, {|00>, a, b, c, d, e, f, g, h, i}, {|01>, j, k, l, m, n, o, p, q, r}, {|02>, s, t, u, v, w, x, y, z, aa}, {|10>, bb, cc, dd, ee, ff, gg, hh, ii, jj}, {|11>, kk, ll, mm, nn, oo, pp, qq, rr, ss}, {|12>, tt, uu, vv, ww, xx, yy, zz, aaa, bbb}, {|20>, ccc, ddd, eee, fff, ggg, hhh, iii, jjj, kkk}, {|21>, lll, mmm, nnn, ooo, ppp, qqq, rrr, sss, ttt}, {|22>, uuu, vvv, www, xxx, yyy, zzz, aaaa, bbbb, cccc}})
 
  (Subscript[ρ, sys])= (({{a, b, c, d, e, f, g, h, i}, {j, k, l, m, n, o, p, q, r}, {s, t, u, v, w, x, y, z, aa}, {bb, cc, dd, ee, ff, gg, hh, ii, jj}, {kk, ll, mm, nn, oo, pp, qq, rr, ss}, {tt, uu, vv, ww, xx, yy, zz, aaa, bbb}, {ccc, ddd, eee, fff, ggg, hhh, iii, jjj, kkk}, {lll, mmm, nnn, ooo, ppp, qqq, rrr, sss, ttt}, {uuu, vvv, www, xxx, yyy, zzz, aaaa, bbbb, cccc}}))
To use the function, you must evaluate it first in The Code section. Then it is straightforward to use, see the Examples section for its possible uses.

The Code

Here is the code for the dTraceSystem  function to trace out an arbitrary number of qudits from an n-qudit quantum system. Evaluate the cell below and you can use it. For more help, please see the Examples section:

SwapParts[expr_, pos1_, pos2_] := ReplacePart[#,#, {pos1,pos2}, {pos2,pos1}]&[expr]
dTraceSystem[D_, s_, dimen_]:= (

Qudits=Reverse[Sort[s]];
TrkM=D;
dim=dimen;

z=(Dimensions[Qudits][[1]]+1);

For[q=1,q<z,q++,
n=Log[dim,(Dimensions[TrkM][[1]])];
M=TrkM;
k=Qudits[[q]];
If[k==n,
TrkM={};
For[p=1,p<(dim^n)+1,p=p+dim,
TrkM=Append[TrkM,(∑_(h=0)^(dim-1))Take[M[[p+h,All]],{1+h,(dim^n),dim}]];
 ],

For[j=0,j<(n-k),j++,
b={0};
For[i=1,i<(dim^n)+1,i++,
If[IntegerDigits[i-1,dim,n][[n]]≠IntegerDigits[i-1,dim,n][[n-j-1]] && Count[b, i]  ==0, 
b=Append[b,(FromDigits[SwapParts[(IntegerDigits[i-1,dim, n]),{n},{n-j-1}],dim]+1)];
c=Range[(dim^n)];
perm=SwapParts[c,{i},{(FromDigits[SwapParts[(IntegerDigits[i-1,dim, n]),{n},{n-j-1}],dim]+1)}];
M=M[[perm,perm]];
 ]    
];
   ];

TrkM={};
For[p=1,p<(dim^n)+1,p=p+dim,
TrkM=Append[TrkM,(∑_(h=0)^(dim-1))Take[M[[p+h,All]],{1+h,(dim^n),dim}]];
]
]
]

;Return[TrkM])

Examples

Here are some examples of how it works. Just to remind you again, the density matrix must be in the computational basis, e.g. |0000...0><0000...0| corresponds to element (Subscript[ρ, 0000...0]), with qudits numbered from 1 to n. dTraceSystem has 3 arguments, the density matrix ρ, the list of qudits to trace out, e.g. {2, 3} and the dimension d of the qudits. The list can be as long as you want and entered in any order, as it is reordered within the function anyway.
- Example 1.
In this example we have a three qubit (d=2) system and we are tracing out the last qubit (qubit 3):

dTraceSystem[(({{a1, a2, a3, a4, a5, a6, a7, a8}, {b1, b2, b3, b4, b5, b6, b7, b8}, {c1, c2, c3, c4, c5, c6, c7, c8}, {d1, d2, d3, d4, d5, d6, d7, d8}, {e1, e2, e3, e4, e5, e6, e7, e8}, {f1, f2, f3, f4, f5, f6, f7, f8}, {g1, g2, g3, g4, g5, g6, g7, g8}, {h1, h2, h3, h4, h5, h6, h7, h8}})),{3},2] //MatrixForm

This leaves us with the density matrix for qubits 1 and 2.

- Example 2.
In this example we have a four qutrit (d=3) system |+>|ψ>|+>|+>, where |ψ>=α|0>+β|1>+γ|2> and |+>=(1/(Sqrt[3]))(|0>+|1>+|2>). We are tracing out qutrits 1, 3 and 4:

Off[General::obspkg]
<<LinearAlgebra`MatrixManipulation`;
<<LinearAlgebra`Orthogonalization`;
VDP[a_List?VectorQ, b_List?VectorQ] :=Flatten[Outer[Times, a, b]];
KETBRA[a_List?VectorQ, b_List?VectorQ] := (Transpose[{a}] ).(Conjugate[Transpose[(Transpose[{b}])]]);
Psiplus=(1/(Sqrt[3])){1,1,1};
Psi={α,β,γ};
PsiTot=VDP[Psiplus,VDP[Psi,VDP[Psiplus,Psiplus]]];
rho=KETBRA[PsiTot,PsiTot];
rho //MatrixForm

Simplify[dTraceSystem[rho,{1,3,4},3]]//MatrixForm

Leaving us with the density matrix for qutrit 2.

How does it Work?

At the bottom of this section, the code for dTraceSystem has parts highlighted in different colours. 
Blue - This part applies a "pick out" method if the last qudit (k==n)  is chosen for tracing out.  (∑_(h=0)^(d-1))Take[M[[p+h,All]],{1+h,(d^n),d}] for h=0, takes all the elements in row p of the density matrix M , starting from 1 and going to (d^n) in steps of d and producing a row vector with (d^n)/d entries as a result. This row vector is then added to a similar "pick out" row vector for each h, i.e. from row p+1, ...,p+(d-1) and we produce a row of the reduced density matrix TrkM. For example if p=1 we are producing the first row of TrkM .

Grey - This part changes the basis of the density matrix so that the qudit k to be traced out (in the case where it is not the last qudit) has a basis corresponding to the last qudit and the same "pick out" method as before is applied (see the  White  part). Basically all the qudits from k  to n have their basis changed in a way such that the basis for qudit  k+i-> (k + i - 1)  for i = 1-> (n - k - 1) and k-> n . 

The coefficents in the density matrix are therefore exchanged according to the basis change (row and column swaps), this is what M=M [[ perm, perm ]] does. The use of Qudits=Reverse[Sort[s]]  is to ensure that the trace procedure for a group of qudits to be traced out, is carried out sequentially, starting with the closest one to qudit n.
Please see the Qubit version notebook (available from Mathsource) for more details on the basis change.

SwapParts[expr_, pos1_, pos2_] := ReplacePart[#,#, {pos1,pos2}, {pos2,pos1}]&[expr]
dTraceSystem[D_, s_, dimen_]:= (

Qudits=Reverse[Sort[s]];
TrkM=D;
dim=dimen;

z=(Dimensions[Qudits][[1]]+1);

For[q=1,q<z,q++,
n=Log[dim,(Dimensions[TrkM][[1]])];
M=TrkM;
k=Qudits[[q]];
If[k==n,
TrkM={};
For[p=1,p<(dim^n)+1,p=p+dim,
TrkM=Append[TrkM,(∑_(h=0)^(dim-1))Take[M[[p+h,All]],{1+h,(dim^n),dim}]];
 ],

For[j=0,j<(n-k),j++,
b={0};
For[i=1,i<(dim^n)+1,i++,
If[IntegerDigits[i-1,dim,n][[n]]≠IntegerDigits[i-1,dim,n][[n-j-1]] && Count[b, i]  ==0, 
b=Append[b,(FromDigits[SwapParts[(IntegerDigits[i-1,dim, n]),{n},{n-j-1}],dim]+1)];
c=Range[(dim^n)];
perm=SwapParts[c,{i},{(FromDigits[SwapParts[(IntegerDigits[i-1,dim, n]),{n},{n-j-1}],dim]+1)}];
M=M[[perm,perm]];
 ]    
];
   ];

TrkM={};
For[p=1,p<(dim^n)+1,p=p+dim,
TrkM=Append[TrkM,(∑_(h=0)^(dim-1))Take[M[[p+h,All]],{1+h,(dim^n),dim}]];
]
]
]

;Return[TrkM])