Skip to main content

oxicuda_quantum/density/
density.rs

1use num_complex::Complex;
2
3use crate::error::QuantumResult;
4use crate::statevec::state::StateVector;
5
6type Complex32 = Complex<f32>;
7
8/// Density matrix ρ stored as a row-major dim×dim complex matrix.
9#[derive(Debug, Clone)]
10pub struct DensityMatrix {
11    pub rho: Vec<Complex32>,
12    pub dim: usize,
13}
14
15impl DensityMatrix {
16    /// Construct ρ = |ψ⟩⟨ψ| from a pure state vector.
17    #[must_use]
18    pub fn from_pure_state(sv: &StateVector) -> Self {
19        let dim = sv.amps.len();
20        let mut rho = vec![Complex32::new(0.0, 0.0); dim * dim];
21        for i in 0..dim {
22            for j in 0..dim {
23                rho[i * dim + j] = sv.amps[i] * sv.amps[j].conj();
24            }
25        }
26        Self { rho, dim }
27    }
28
29    /// Apply a unitary U: ρ' = U ρ U†.
30    pub fn apply_unitary(&mut self, u: &[Complex32], dim: usize) -> QuantumResult<()> {
31        use crate::error::QuantumError;
32        if dim != self.dim {
33            return Err(QuantumError::DimensionMismatch {
34                expected: self.dim,
35                got: dim,
36            });
37        }
38        // u_rho = U * ρ
39        let mut u_rho = vec![Complex32::new(0.0, 0.0); dim * dim];
40        for i in 0..dim {
41            for j in 0..dim {
42                for k in 0..dim {
43                    u_rho[i * dim + j] += u[i * dim + k] * self.rho[k * dim + j];
44                }
45            }
46        }
47        // result = u_rho * U†
48        let mut result = vec![Complex32::new(0.0, 0.0); dim * dim];
49        for i in 0..dim {
50            for j in 0..dim {
51                for k in 0..dim {
52                    // U†[j,k] = conj(U[k,j])
53                    result[i * dim + j] += u_rho[i * dim + k] * u[j * dim + k].conj();
54                }
55            }
56        }
57        self.rho = result;
58        Ok(())
59    }
60
61    /// Trace Tr(ρ) = Σ_i ρ\[i,i\].
62    #[must_use]
63    pub fn trace(&self) -> Complex32 {
64        (0..self.dim)
65            .map(|i| self.rho[i * self.dim + i])
66            .fold(Complex32::new(0.0, 0.0), |acc, x| acc + x)
67    }
68
69    /// Tr(ρ²) used for purity calculation.
70    #[must_use]
71    pub fn trace_sq(&self) -> Complex32 {
72        let dim = self.dim;
73        let mut tr2 = Complex32::new(0.0, 0.0);
74        for i in 0..dim {
75            for k in 0..dim {
76                // (ρ²)[i,i] = Σ_k ρ[i,k] * ρ[k,i]
77                tr2 += self.rho[i * dim + k] * self.rho[k * dim + i];
78            }
79        }
80        tr2
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn pure_state_trace_is_one() {
90        let sv = StateVector::new_zero_state(2)
91            .expect("n_qubits=2 is a valid qubit count so zero-state construction cannot fail");
92        let dm = DensityMatrix::from_pure_state(&sv);
93        let tr = dm.trace();
94        assert!((tr.re - 1.0).abs() < 1e-6, "trace={}", tr.re);
95    }
96
97    #[test]
98    fn pure_state_trace_sq_is_one() {
99        let sv = StateVector::new_zero_state(1)
100            .expect("n_qubits=1 is a valid qubit count so zero-state construction cannot fail");
101        let dm = DensityMatrix::from_pure_state(&sv);
102        let tr2 = dm.trace_sq();
103        assert!((tr2.re - 1.0).abs() < 1e-5, "trace_sq={}", tr2.re);
104    }
105}