oxicuda-quantum 0.1.6

Quantum simulation primitives for OxiCUDA — state-vector simulation, Pauli expectation values, VQE/QAOA, density matrices, Trotter-Suzuki, quantum kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use num_complex::Complex;

use crate::density::density::DensityMatrix;
use crate::error::{QuantumError, QuantumResult};

type Complex32 = Complex<f32>;

/// Compute the reduced density matrix by tracing out all qubits not in `keep_qubits`.
///
/// `n_total` is the total number of qubits; `keep_qubits` is sorted in ascending order.
/// The output density matrix has dimension `2^|keep_qubits|`.
pub fn partial_trace(
    dm: &DensityMatrix,
    keep_qubits: &[usize],
    n_total: usize,
) -> QuantumResult<DensityMatrix> {
    let full_dim = 1usize << n_total;
    if dm.dim != full_dim {
        return Err(QuantumError::DimensionMismatch {
            expected: full_dim,
            got: dm.dim,
        });
    }

    for &q in keep_qubits {
        if q >= n_total {
            return Err(QuantumError::QubitIndexOutOfRange {
                index: q,
                n_qubits: n_total,
            });
        }
    }

    let n_keep = keep_qubits.len();
    let reduced_dim = 1usize << n_keep;
    let mut rho_red = vec![Complex32::new(0.0, 0.0); reduced_dim * reduced_dim];

    // Build the set of traced-out (discarded) qubit indices
    let mut trace_qubits: Vec<usize> = (0..n_total).filter(|q| !keep_qubits.contains(q)).collect();
    trace_qubits.sort_unstable();
    let n_trace = trace_qubits.len();
    let trace_dim = 1usize << n_trace;

    // For each pair of keep-basis indices (i_keep, j_keep), sum over trace index
    for ri in 0..reduced_dim {
        for rj in 0..reduced_dim {
            let mut sum = Complex32::new(0.0, 0.0);

            for tk in 0..trace_dim {
                // Reconstruct full indices by interleaving keep and trace bits
                let full_i = interleave_bits(ri, tk, keep_qubits, &trace_qubits, n_total);
                let full_j = interleave_bits(rj, tk, keep_qubits, &trace_qubits, n_total);
                sum += dm.rho[full_i * full_dim + full_j];
            }

            rho_red[ri * reduced_dim + rj] = sum;
        }
    }

    Ok(DensityMatrix {
        rho: rho_red,
        dim: reduced_dim,
    })
}

/// Reconstruct a full-system basis index from keep-bits and trace-bits.
fn interleave_bits(
    keep_idx: usize,
    trace_idx: usize,
    keep_qubits: &[usize],
    trace_qubits: &[usize],
    n_total: usize,
) -> usize {
    let _ = n_total;
    let mut result = 0usize;
    for (k, &q) in keep_qubits.iter().enumerate() {
        if (keep_idx >> k) & 1 == 1 {
            result |= 1 << q;
        }
    }
    for (k, &q) in trace_qubits.iter().enumerate() {
        if (trace_idx >> k) & 1 == 1 {
            result |= 1 << q;
        }
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::statevec::state::StateVector;

    #[test]
    fn partial_trace_2q_product_state() {
        // |0⟩ ⊗ |0⟩ → trace out qubit 1 → still |0⟩⟨0|
        let sv = StateVector::new_zero_state(2).unwrap();
        let dm = DensityMatrix::from_pure_state(&sv);
        let red = partial_trace(&dm, &[0], 2).unwrap();
        assert_eq!(red.dim, 2);
        assert!((red.rho[0].re - 1.0).abs() < 1e-5);
    }
}