oxicuda_quantum/density/
density.rs1use num_complex::Complex;
2
3use crate::error::QuantumResult;
4use crate::statevec::state::StateVector;
5
6type Complex32 = Complex<f32>;
7
8#[derive(Debug, Clone)]
10pub struct DensityMatrix {
11 pub rho: Vec<Complex32>,
12 pub dim: usize,
13}
14
15impl DensityMatrix {
16 #[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 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 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 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 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 #[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 #[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 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).unwrap();
91 let dm = DensityMatrix::from_pure_state(&sv);
92 let tr = dm.trace();
93 assert!((tr.re - 1.0).abs() < 1e-6, "trace={}", tr.re);
94 }
95
96 #[test]
97 fn pure_state_trace_sq_is_one() {
98 let sv = StateVector::new_zero_state(1).unwrap();
99 let dm = DensityMatrix::from_pure_state(&sv);
100 let tr2 = dm.trace_sq();
101 assert!((tr2.re - 1.0).abs() < 1e-5, "trace_sq={}", tr2.re);
102 }
103}