oxicuda_quantum/density/
partial_trace.rs

1use num_complex::Complex;
2
3use crate::density::density::DensityMatrix;
4use crate::error::{QuantumError, QuantumResult};
5
6type Complex32 = Complex<f32>;
7
8/// Compute the reduced density matrix by tracing out all qubits not in `keep_qubits`.
9///
10/// `n_total` is the total number of qubits; `keep_qubits` is sorted in ascending order.
11/// The output density matrix has dimension `2^|keep_qubits|`.
12pub fn partial_trace(
13    dm: &DensityMatrix,
14    keep_qubits: &[usize],
15    n_total: usize,
16) -> QuantumResult<DensityMatrix> {
17    let full_dim = 1usize << n_total;
18    if dm.dim != full_dim {
19        return Err(QuantumError::DimensionMismatch {
20            expected: full_dim,
21            got: dm.dim,
22        });
23    }
24
25    for &q in keep_qubits {
26        if q >= n_total {
27            return Err(QuantumError::QubitIndexOutOfRange {
28                index: q,
29                n_qubits: n_total,
30            });
31        }
32    }
33
34    let n_keep = keep_qubits.len();
35    let reduced_dim = 1usize << n_keep;
36    let mut rho_red = vec![Complex32::new(0.0, 0.0); reduced_dim * reduced_dim];
37
38    // Build the set of traced-out (discarded) qubit indices
39    let mut trace_qubits: Vec<usize> = (0..n_total).filter(|q| !keep_qubits.contains(q)).collect();
40    trace_qubits.sort_unstable();
41    let n_trace = trace_qubits.len();
42    let trace_dim = 1usize << n_trace;
43
44    // For each pair of keep-basis indices (i_keep, j_keep), sum over trace index
45    for ri in 0..reduced_dim {
46        for rj in 0..reduced_dim {
47            let mut sum = Complex32::new(0.0, 0.0);
48
49            for tk in 0..trace_dim {
50                // Reconstruct full indices by interleaving keep and trace bits
51                let full_i = interleave_bits(ri, tk, keep_qubits, &trace_qubits, n_total);
52                let full_j = interleave_bits(rj, tk, keep_qubits, &trace_qubits, n_total);
53                sum += dm.rho[full_i * full_dim + full_j];
54            }
55
56            rho_red[ri * reduced_dim + rj] = sum;
57        }
58    }
59
60    Ok(DensityMatrix {
61        rho: rho_red,
62        dim: reduced_dim,
63    })
64}
65
66/// Reconstruct a full-system basis index from keep-bits and trace-bits.
67fn interleave_bits(
68    keep_idx: usize,
69    trace_idx: usize,
70    keep_qubits: &[usize],
71    trace_qubits: &[usize],
72    n_total: usize,
73) -> usize {
74    let _ = n_total;
75    let mut result = 0usize;
76    for (k, &q) in keep_qubits.iter().enumerate() {
77        if (keep_idx >> k) & 1 == 1 {
78            result |= 1 << q;
79        }
80    }
81    for (k, &q) in trace_qubits.iter().enumerate() {
82        if (trace_idx >> k) & 1 == 1 {
83            result |= 1 << q;
84        }
85    }
86    result
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use crate::statevec::state::StateVector;
93
94    #[test]
95    fn partial_trace_2q_product_state() {
96        // |0⟩ ⊗ |0⟩ → trace out qubit 1 → still |0⟩⟨0|
97        let sv = StateVector::new_zero_state(2).unwrap();
98        let dm = DensityMatrix::from_pure_state(&sv);
99        let red = partial_trace(&dm, &[0], 2).unwrap();
100        assert_eq!(red.dim, 2);
101        assert!((red.rho[0].re - 1.0).abs() < 1e-5);
102    }
103}
oxicuda_quantum/density/partial_trace.rs

oxicuda_quantum/density/
partial_trace.rs
