oxicuda_quantum/statevec/
state.rs1use num_complex::Complex;
2
3use crate::error::{QuantumError, QuantumResult};
4use crate::handle::LcgRng;
5
6type Complex32 = Complex<f32>;
7
8#[derive(Debug, Clone)]
10pub struct StateVector {
11 pub amps: Vec<Complex32>,
12 pub n_qubits: usize,
13}
14
15impl StateVector {
16 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 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 #[must_use]
52 pub fn norm_sq(&self) -> f32 {
53 self.amps.iter().map(|a| a.norm_sqr()).sum()
54 }
55
56 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 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 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 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).expect(
147 "new_zero_state(3) cannot fail: 3 qubits is strictly within the valid range 1..=30",
148 );
149 let n = sv.norm_sq();
150 assert!((n - 1.0).abs() < 1e-6, "norm={n}");
151 }
152
153 #[test]
154 fn invalid_qubit_count() {
155 assert!(StateVector::new_zero_state(0).is_err());
156 }
157
158 #[test]
159 fn from_amps_rejects_unnormalized() {
160 let amps = vec![Complex32::new(1.0, 0.0), Complex32::new(1.0, 0.0)];
161 assert!(StateVector::new_from_amps(amps, 1).is_err());
162 }
163}