Partial Trace of a Multi-qubit System
by
Mark S. Tame
School of Mathematics and Physics, Queen's University, Belfast.  UK.
m.tame@qub.ac.uk

Introduction

For many situations in Quantum Information Theory it is necessary to trace out a subgroup of qubits from a larger group. For example, suppose Alice, Bob and Charlie each have a qubit with the total system described by the density matrix (Subscript[ρ, ABC]). If we would like  to examine more closely the subsystem of qubits shared by just Alice and Bob we need to take the trace:
 (Subscript[ρ, AB])=(Subscript[Tr, C]) {(Subscript[ρ, ABC]) }
However for large multi-qubit systems, calculating this by hand becomes quite unmanageable. What the TraceSystem  function does, is allow you to trace out an arbitrary number of qubits from a large system of N-qubits by getting Mathematica do the work for you:
TraceSystem[(Subscript[ρ, sys]), {1,3,6,...}] = (Subscript[Tr, [1, 3, 6,  ... ]]){(Subscript[ρ, sys])}
Where (Subscript[ρ, sys]) is in the computational basis eg. for 2 qubits we have:
({{Basis, <00|, <01|, \(\(\(<\)\(10\)\)\(|\)\), <11|}, {|00>, a, b, c, d}, {|01>, e, f, g, h}, {|10>, i, j, k, l}, {|11>, m, n, o, p}})                            (Subscript[ρ, sys])= (({{a, b, c, d}, {e, f, g, h}, {i, j, k, l}, {m, n, o, p}}))
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 TraceSystem  function to trace out an arbitrary number of qubits from an N-qubit 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]
TraceSystem[D_,  
        s_] :=  (

Qubits = 
          Reverse[Sort[s]]; 
TrkM = 
          D; 

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

For[q = 1, 
          q < z, q++, 
n = 
            Log[2, (Dimensions[
                    TrkM][[1]])]; 
M = 
            TrkM; 
k = 
            Qubits[[q]]; 
If[
            k == 
              n, 
TrkM = {}; 
For[p = 1, 
              p < (2^n) + 1, 
              p = p + 2, 
TrkM = 
                  Append[TrkM, 
                    Take[M[[p, All]], {1, (2^n), 2}] + 
                      Take[M[[p + 1, All]], {2, (2^n), 
                          2}]];
 ], 
For[j = 0, 
                j < (n - 
                      k), j++, 
b = {0}; 
For[i = 1, 
                  i < (2^n) + 
                      1, i++, 
If[(Mod[(IntegerDigits[i - 1, 2, 
                                    n][[n]] + IntegerDigits[i - 1, 2, 
                                    n][[n - j - 1]]), 2]) == 
                        1  &&  Count[b,  i]   == 0,  
                    Permut = {i, (FromDigits[
                              SwapParts[(IntegerDigits[i - 1, 2,  
                                    n]), {n}, {n - j - 1}], 2] + 
                            1)}; 
b = 
                      Append[b, (FromDigits[
                              SwapParts[(IntegerDigits[i - 1, 2,  
                                    n]), {n}, {n - j - 1}], 2] + 
                            1)]; 
c = 
                      Range[(2^n)]; 
perm = 
                      SwapParts[
                        c, {i}, {(FromDigits[
                                SwapParts[(IntegerDigits[i - 1, 2,  
                                      n]), {n}, {n - j - 1}], 2] + 
                              1)}]; 

M = M[[perm, 
                          perm]];

 ]    
]         ; 
TrkM = {}; 
For[p = 1, p < (2^n) + 1, 
                  p = p + 2, 
TrkM = 
                      Append[TrkM, 
                        Take[M[[p, All]], {1, (2^n), 2}] + 
                          Take[M[[p + 1, All]], {2, (2^n), 
                              
                              2}]];
]
   ];
]

]

; 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 eg. |0000...0><0000...0| corresponds to element (Subscript[ρ, 00]), with qubits numbered from 1 to N. TraceSystem has 2 arguments, the density matrix ρ and the list of qubits to trace out eg. {2, 3}. This 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 system and we are tracing out the last qubit (qubit 3):

TraceSystem[(({{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}] //MatrixForm

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

- Example 2.
In this example we have a three qubit system and we are tracing out qubits 2 and 3:

TraceSystem[(({{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}})),{2, 3}] //MatrixForm

Leaving us with the density matrix for qubit 1.

- Example 3.
In this example we have a four qubit system and we are tracing out qubits 1 and 4:

Psi = {{a1}, {b1}, {c1}, {d1}, {e1}, {f1}, {g1}, {h1}, {i1}, {j1}, {k1}, {l1}, {m1}, {n1}, {o1}, {p1}};
rho = Psi . Transpose[Psi];
rho  // MatrixForm

TraceSystem[rho, {1, 4}]  // MatrixForm

Leaving us with the density matrix for qubits 2 and 3.

How does it Work?

At the bottom of this section, the code for TraceSystem has parts highlighted in different colours. These parts are explained below:
Blue - This part applies a "pick out" method if the last qubit (k==n)  is chosen for tracing out. Basically what Take[M[[p, All]], {1, (2^n), 2}] does is take all the elements in row p of the density matrix M , starting from 1 and going to (2^n) in steps of 2 and producing a row vector with (2^n)/2 entries as a result. This row vector is then added to a similar "pick out" row vector from row p+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 qubit k to be traced out (in the case where it is not the last qubit) has a basis corresponding to the last qubit and the same "pick out" method as before is applied (see the  White  part). Basically all the qubits from k  to n have their basis changed in a way such that the basis for qubit  k+i-> (k + i - 1)  for i = 1-> (n - k - 1) and k-> n . An example of this basis change is given below for n=5 and k=2 .

(({{0, 0, 0, 0, 0}, {0, 0, 0, 1, 0}, {0, 0, 0, 0, 1}, {0, 0, 0, 1, 1}, {0, 0, 1, 0, 0}, {0, 0, 1, 1, 0}, {0, 0, 1, 0, 1}, {0, 0, 1, 1, 1}, {0, 1, 0, 0, 0}, {0, 1, 0, 1, 0}, {0, 1, 0, 0, 1}, {0, 1, 0, 1, 1}, {0, 1, 1, 0, 0}, {0, 1, 1, 1, 0}, {0, 1, 1, 0, 1}, {0, 1, 1, 1, 1}, {1, 0, 0, 0, 0}, {1, 0, 0, 1, 0}, {1, 0, 0, 0, 1}, {1, 0, 0, 1, 1}, {1, 0, 1, 0, 0}, {1, 0, 1, 1, 0}, {1, 0, 1, 0, 1}, {1, 0, 1, 1, 1}, {1, 1, 0, 0, 0}, {1, 1, 0, 1, 0}, {1, 1, 0, 0, 1}, {1, 1, 0, 1, 1}, {1, 1, 1, 0, 0}, {1, 1, 1, 1, 0}, {1, 1, 1, 0, 1}, {1, 1, 1, 1, 1}}))   ->   (({{0, 0, 0, 0, 0}, {0, 0, 1, 0, 0}, {0, 0, 0, 0, 1}, {0, 0, 1, 0, 1}, {0, 0, 0, 1, 0}, {0, 0, 1, 1, 0}, {0, 0, 0, 1, 1}, {0, 0, 1, 1, 1}, {0, 1, 0, 0, 0}, {0, 1, 1, 0, 0}, {0, 1, 0, 0, 1}, {0, 1, 1, 0, 1}, {0, 1, 0, 1, 0}, {0, 1, 1, 1, 0}, {0, 1, 0, 1, 1}, {0, 1, 1, 1, 1}, {1, 0, 0, 0, 0}, {1, 0, 1, 0, 0}, {1, 0, 0, 0, 1}, {1, 0, 1, 0, 1}, {1, 0, 0, 1, 0}, {1, 0, 1, 1, 0}, {1, 0, 0, 1, 1}, {1, 0, 1, 1, 1}, {1, 1, 0, 0, 0}, {1, 1, 1, 0, 0}, {1, 1, 0, 0, 1}, {1, 1, 1, 0, 1}, {1, 1, 0, 1, 0}, {1, 1, 1, 1, 0}, {1, 1, 0, 1, 1}, {1, 1, 1, 1, 1}}))   ->   (({{0, 0, 0, 0, 0}, {0, 1, 0, 0, 0}, {0, 0, 0, 0, 1}, {0, 1, 0, 0, 1}, {0, 0, 0, 1, 0}, {0, 1, 0, 1, 0}, {0, 0, 0, 1, 1}, {0, 1, 0, 1, 1}, {0, 0, 1, 0, 0}, {0, 1, 1, 0, 0}, {0, 0, 1, 0, 1}, {0, 1, 1, 0, 1}, {0, 0, 1, 1, 0}, {0, 1, 1, 1, 0}, {0, 0, 1, 1, 1}, {0, 1, 1, 1, 1}, {1, 0, 0, 0, 0}, {1, 1, 0, 0, 0}, {1, 0, 0, 0, 1}, {1, 1, 0, 0, 1}, {1, 0, 0, 1, 0}, {1, 1, 0, 1, 0}, {1, 0, 0, 1, 1}, {1, 1, 0, 1, 1}, {1, 0, 1, 0, 0}, {1, 1, 1, 0, 0}, {1, 0, 1, 0, 1}, {1, 1, 1, 0, 1}, {1, 0, 1, 1, 0}, {1, 1, 1, 1, 0}, {1, 0, 1, 1, 1}, {1, 1, 1, 1, 1}}))
(Overscript[stays, ↑])        (Overscript[swapped, ⋃])                     (Overscript[swapped, ⋃])                     (Overscript[swapped, ⋃])          

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 ]] on line 28 does. The use of Qubits=Reverse[Sort[s]]  on line 4 is to ensure that the trace procedure for a group of qubits to be traced out, is carried out sequentially, starting with the closest one to qubit n .

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

Qubits = Reverse[Sort[s]];
TrkM = D;

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

For[q = 1,q < z,q++,
n = Log[2, (Dimensions[TrkM][[1]])];
M = TrkM;
k = Qubits[[q]];
If[k == n,
TrkM = {};

                      
                      For[p = 1, p < (2^n) + 1, 
                        p = p + 2, 
TrkM = 
                            Append[TrkM, 
                              Take[M[[p, All]], {1, (2^n), 2}] + 
                                Take[M[[p + 1, All]], {2, (2^n), 
                                    2}]];
 ],
For[j = 0,j < (n - k),j++,
b = {0};
For[i = 1, 
                                i < (2^n) + 
                                    1, i++, 
If[(Mod[(IntegerDigits[i - 1, 2, n][[n]] + IntegerDigits[i - 1, 2, 
                                        n][[n - j - 1]]), 2]) == 
                                      1  &&  Count[b,  i]   == 
                                      0,  Permut = {i, (FromDigits[
                                        SwapParts[(IntegerDigits[i - 1, 
                                        2,  n]), {n}, {n - j - 1}], 2] + 
                                        1)}; 
b = 
                                    Append[
                                      b, (FromDigits[
                                        SwapParts[(IntegerDigits[i - 1, 
                                        2,  n]), {n}, {n - j - 1}], 2] + 
                                        1)]; 
c = 
                                    Range[(2^n)]; 
perm = 
                                    SwapParts[
                                      c, {i}, {(FromDigits[
                                        SwapParts[(IntegerDigits[i - 1, 
                                        2,  n]), {n}, {n - j - 1}], 2] + 
                                        1)}]; 

M = M[[perm, 
                                        perm]];

 ]    
]         ;
TrkM = {};
For[p = 1, p < (2^n) + 1, 
                                p = p + 2, 
TrkM = 
                                    Append[TrkM, 
                                      Take[M[[p, All]], {1, (2^n), 2}] + 
                                        Take[
                                        M[[p + 1, All]], {2, (2^n), 
                                        2}]];
]
   ];
]

]

;Return[TrkM])