Skip to main content

oxicuda_quantum/statevec/
state.rs

1use num_complex::Complex;
2
3use crate::error::{QuantumError, QuantumResult};
4use crate::handle::LcgRng;
5
6type Complex32 = Complex<f32>;
7
8/// Pure quantum state represented as a complex amplitude vector.
9#[derive(Debug, Clone)]
10pub struct StateVector {
11    pub amps: Vec<Complex32>,
12    pub n_qubits: usize,
13}
14
15impl StateVector {
16    /// Construct |0⟩^n = [1, 0, 0, …].
17    pub fn new_zero_state(n_qubits: usize) -> QuantumResult<Self> {
18        if n_qubits == 0 {
19            return Err(QuantumError::InvalidQubitCount { n: 0 });
20        }
21        if n_qubits > 30 {
22            return Err(QuantumError::InvalidQubitCount { n: n_qubits });
23        }
24        let dim = 1usize << n_qubits;
25        let mut amps = vec![Complex32::new(0.0, 0.0); dim];
26        amps[0] = Complex32::new(1.0, 0.0);
27        Ok(Self { amps, n_qubits })
28    }
29
30    /// Construct from a given amplitude vector; validates length = 2^n and unit norm.
31    pub fn new_from_amps(amps: Vec<Complex32>, n_qubits: usize) -> QuantumResult<Self> {
32        if n_qubits == 0 || n_qubits > 30 {
33            return Err(QuantumError::InvalidQubitCount { n: n_qubits });
34        }
35        let expected = 1usize << n_qubits;
36        if amps.len() != expected {
37            return Err(QuantumError::DimensionMismatch {
38                expected,
39                got: amps.len(),
40            });
41        }
42        let sv = Self { amps, n_qubits };
43        let norm = sv.norm_sq().sqrt();
44        if (norm - 1.0).abs() > 1e-4 {
45            return Err(QuantumError::NonNormalizedState { norm });
46        }
47        Ok(sv)
48    }
49
50    /// Sum of squared magnitudes: ‖ψ‖².
51    #[must_use]
52    pub fn norm_sq(&self) -> f32 {
53        self.amps.iter().map(|a| a.norm_sqr()).sum()
54    }
55
56    /// Normalize in place; no-op if already unit.
57    pub fn normalize_inplace(&mut self) {
58        let norm = self.norm_sq().sqrt();
59        if norm > 1e-12 {
60            let inv = 1.0 / norm;
61            for a in &mut self.amps {
62                *a *= inv;
63            }
64        }
65    }
66
67    /// Inner product ⟨self|other⟩ = Σ conj(self\[i\]) * other\[i\].
68    pub fn inner_product(&self, other: &Self) -> QuantumResult<Complex32> {
69        if self.n_qubits != other.n_qubits {
70            return Err(QuantumError::DimensionMismatch {
71                expected: self.amps.len(),
72                got: other.amps.len(),
73            });
74        }
75        Ok(self
76            .amps
77            .iter()
78            .zip(other.amps.iter())
79            .map(|(a, b)| a.conj() * b)
80            .fold(Complex32::new(0.0, 0.0), |acc, x| acc + x))
81    }
82
83    /// Probability of measuring `qubit` in state `outcome`.
84    pub fn measure_prob(&self, qubit: usize, outcome: bool) -> QuantumResult<f32> {
85        if qubit >= self.n_qubits {
86            return Err(QuantumError::QubitIndexOutOfRange {
87                index: qubit,
88                n_qubits: self.n_qubits,
89            });
90        }
91        let mask = 1usize << qubit;
92        let prob = self
93            .amps
94            .iter()
95            .enumerate()
96            .filter(|(i, _)| ((i & mask) != 0) == outcome)
97            .map(|(_, a)| a.norm_sqr())
98            .sum();
99        Ok(prob)
100    }
101
102    /// Sample a measurement of `qubit`, returning (outcome, post-measurement StateVector).
103    pub fn sample_measure(&self, qubit: usize, rng: &mut LcgRng) -> QuantumResult<(bool, Self)> {
104        let p1 = self.measure_prob(qubit, true)?;
105        let r = rng.next_f32();
106        let outcome = r < p1;
107
108        let mask = 1usize << qubit;
109        let mut new_amps = self.amps.clone();
110        let mut norm_sq = 0.0_f32;
111
112        for (i, a) in new_amps.iter_mut().enumerate() {
113            let bit_set = (i & mask) != 0;
114            if bit_set != outcome {
115                *a = Complex32::new(0.0, 0.0);
116            } else {
117                norm_sq += a.norm_sqr();
118            }
119        }
120
121        let norm = norm_sq.sqrt();
122        if norm < 1e-12 {
123            return Err(QuantumError::MeasurementFailed);
124        }
125        let inv = 1.0 / norm;
126        for a in &mut new_amps {
127            *a *= inv;
128        }
129
130        Ok((
131            outcome,
132            Self {
133                amps: new_amps,
134                n_qubits: self.n_qubits,
135            },
136        ))
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn zero_state_norm_is_one() {
146        let sv = StateVector::new_zero_state(3).unwrap();
147        let n = sv.norm_sq();
148        assert!((n - 1.0).abs() < 1e-6, "norm={n}");
149    }
150
151    #[test]
152    fn invalid_qubit_count() {
153        assert!(StateVector::new_zero_state(0).is_err());
154    }
155
156    #[test]
157    fn from_amps_rejects_unnormalized() {
158        let amps = vec![Complex32::new(1.0, 0.0), Complex32::new(1.0, 0.0)];
159        assert!(StateVector::new_from_amps(amps, 1).is_err());
160    }
161}