Skip to main content

ppvm_tableau/
sparsevec.rs

1// SPDX-FileCopyrightText: 2026 The PPVM Authors
2// SPDX-License-Identifier: Apache-2.0
3
4use num::complex::{Complex, ComplexFloat};
5use num::traits::{One, Zero};
6
7/// A sparse vector keyed by index type `I` with values of type `T`.
8///
9/// Used by [`GeneralizedTableau`](crate::data::GeneralizedTableau) to
10/// store the coefficient over each branching bitstring. The default
11/// implementation is `Vec<(T, I)>`; alternative backings (BTreeMap,
12/// HashMap, etc.) can be provided by downstream code.
13pub trait SparseVector<T, I>: Clone + IntoIterator<Item = (T, I)> {
14    /// Construct an empty sparse vector.
15    fn new() -> Self;
16    /// Inserts an element without checking whether the index already exists.
17    fn unsafe_insert(&mut self, index: I, value: T);
18    /// Add `value` into the entry at `index`, creating it if absent.
19    fn add_or_insert(&mut self, index: I, value: T);
20    /// Retrieve the value at `index`, or zero if absent.
21    fn get(&self, index: &I) -> T;
22    /// Number of stored entries.
23    fn len(&self) -> usize;
24    /// `true` if no entries are stored.
25    fn is_empty(&self) -> bool;
26    /// Multiply every entry's value by `factor`.
27    fn mul_by(&mut self, factor: T);
28    /// Multiply the value at `index` by `factor`. No-op if absent.
29    fn mul_element_by(&mut self, index: I, factor: T);
30    /// Drop entries whose magnitude is at most `|cutoff|`.
31    fn trim(&mut self, cutoff: T);
32    /// Drop entries failing the predicate `f`.
33    fn retain(&mut self, f: impl FnMut(&(T, I)) -> bool);
34    /// L2-normalize the vector in place. Panics on zero norm.
35    fn normalize(&mut self);
36    /// Reserve capacity for at least `additional` more entries. Backings
37    /// that don't support pre-allocation can leave this as a no-op.
38    fn reserve(&mut self, _additional: usize) {}
39    /// Borrow the stored entries as an iterator without consuming the vector.
40    fn iter<'a>(&'a self) -> impl Iterator<Item = &'a (T, I)>
41    where
42        T: 'a,
43        I: 'a;
44}
45
46impl<T, I> SparseVector<T, I> for Vec<(T, I)>
47where
48    T: std::ops::AddAssign + std::ops::MulAssign + One + ComplexFloat + Zero,
49    f64: std::ops::Div<<T as ComplexFloat>::Real>,
50    I: std::cmp::PartialEq + Clone,
51{
52    fn new() -> Self {
53        Vec::new()
54    }
55
56    fn unsafe_insert(&mut self, index: I, value: T) {
57        self.push((value, index));
58    }
59
60    fn add_or_insert(&mut self, index: I, value: T) {
61        for (v, i) in self.iter_mut() {
62            if *i == index {
63                *v += value;
64                return;
65            }
66        }
67        self.push((value, index));
68    }
69
70    fn get(&self, index: &I) -> T {
71        for (v, i) in self.iter() {
72            if i == index {
73                return *v;
74            }
75        }
76        T::zero()
77    }
78
79    fn len(&self) -> usize {
80        self.len()
81    }
82
83    fn is_empty(&self) -> bool {
84        self.is_empty()
85    }
86
87    fn mul_by(&mut self, factor: T) {
88        for (v, _) in self.iter_mut() {
89            *v *= factor;
90        }
91    }
92
93    fn mul_element_by(&mut self, index: I, factor: T) {
94        for (v, i) in self.iter_mut() {
95            if *i == index {
96                *v *= factor;
97                return;
98            }
99        }
100    }
101
102    fn trim(&mut self, cutoff: T) {
103        // TODO: make cutoff real
104        let c_re = cutoff.re();
105        let c_im = cutoff.im();
106        let cutoff_sq = c_re * c_re + c_im * c_im;
107        self.retain(|(element, _)| {
108            let e_re = element.re();
109            let e_im = element.im();
110            e_re * e_re + e_im * e_im > cutoff_sq
111        });
112    }
113
114    fn retain(&mut self, f: impl FnMut(&(T, I)) -> bool) {
115        Vec::retain(self, f);
116    }
117
118    fn reserve(&mut self, additional: usize) {
119        Vec::reserve(self, additional);
120    }
121
122    fn iter<'a>(&'a self) -> impl Iterator<Item = &'a (T, I)>
123    where
124        T: 'a,
125        I: 'a,
126    {
127        <[(T, I)]>::iter(self)
128    }
129
130    fn normalize(&mut self) {
131        // `re*re + im*im` directly; `abs() * abs()` would compute
132        // `hypot(re, im)` (a sqrt) only to immediately square it again.
133        let norm: T::Real = self.iter().fold(T::Real::zero(), |acc, (v, _)| {
134            let re = v.re();
135            let im = v.im();
136            acc + re * re + im * im
137        });
138
139        if norm == T::Real::zero() {
140            panic!("Zero norm encountered during normalization");
141        }
142        if norm == T::Real::one() {
143            return;
144        }
145        let norm_sqrt = norm.sqrt();
146        let inv_norm_sqrt = T::Real::one() / norm_sqrt;
147        let scale: T = T::from(Complex::new(inv_norm_sqrt, T::Real::zero()))
148            .expect("Failed to convert scale factor");
149        for (v, _) in self.iter_mut() {
150            *v *= scale;
151        }
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use num::complex::Complex64;
159
160    #[test]
161    fn test_sparse_vector_new() {
162        let vec: Vec<(Complex64, usize)> = SparseVector::new();
163        assert!(vec.is_empty());
164        assert_eq!(vec.len(), 0);
165    }
166
167    #[test]
168    fn test_unsafe_insert() {
169        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
170        vec.unsafe_insert(0, Complex64::new(1.0, 0.0));
171        vec.unsafe_insert(2, Complex64::new(2.0, 1.0));
172
173        assert_eq!(vec.len(), 2);
174        assert_eq!(vec.get(&0), Complex64::new(1.0, 0.0));
175        assert_eq!(vec.get(&2), Complex64::new(2.0, 1.0));
176        assert_eq!(vec.get(&1), Complex64::new(0.0, 0.0)); // Non-existent index returns zero
177    }
178
179    #[test]
180    fn test_add_or_insert() {
181        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
182
183        // First insert
184        vec.add_or_insert(0, Complex64::new(1.0, 0.0));
185        assert_eq!(vec.get(&0), Complex64::new(1.0, 0.0));
186
187        // Add to existing
188        vec.add_or_insert(0, Complex64::new(2.0, 1.0));
189        assert_eq!(vec.get(&0), Complex64::new(3.0, 1.0));
190
191        // Insert new
192        vec.add_or_insert(5, Complex64::new(0.5, -0.5));
193        assert_eq!(vec.get(&5), Complex64::new(0.5, -0.5));
194    }
195
196    #[test]
197    fn test_mul_element_by() {
198        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
199        vec.unsafe_insert(0, Complex64::new(2.0, 1.0));
200        vec.unsafe_insert(1, Complex64::new(1.0, 0.0));
201
202        vec.mul_element_by(0, Complex64::new(2.0, 0.0));
203        assert_eq!(vec.get(&0), Complex64::new(4.0, 2.0));
204        assert_eq!(vec.get(&1), Complex64::new(1.0, 0.0)); // Unchanged
205
206        // Multiply non-existent element (should do nothing)
207        vec.mul_element_by(99, Complex64::new(5.0, 0.0));
208        assert_eq!(vec.get(&99), Complex64::new(0.0, 0.0));
209    }
210
211    #[test]
212    fn test_trim() {
213        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
214        vec.unsafe_insert(0, Complex64::new(1.0, 0.0));
215        vec.unsafe_insert(1, Complex64::new(0.001, 0.0));
216        vec.unsafe_insert(2, Complex64::new(0.0001, 0.0));
217        vec.unsafe_insert(3, Complex64::new(2.0, 1.0));
218
219        assert_eq!(vec.len(), 4);
220
221        vec.trim(Complex64::new(0.01, 0.0));
222        assert_eq!(vec.len(), 2); // Only elements with abs > 0.01 remain
223        assert_eq!(vec.get(&0), Complex64::new(1.0, 0.0));
224        assert_eq!(vec.get(&3), Complex64::new(2.0, 1.0));
225        assert_eq!(vec.get(&1), Complex64::new(0.0, 0.0)); // Trimmed
226    }
227
228    #[test]
229    fn test_retain() {
230        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
231        vec.unsafe_insert(0, Complex64::new(1.0, 0.0));
232        vec.unsafe_insert(1, Complex64::new(2.0, 0.0));
233        vec.unsafe_insert(2, Complex64::new(3.0, 0.0));
234        vec.unsafe_insert(3, Complex64::new(4.0, 0.0));
235
236        // Keep only even indices
237        vec.retain(|(_, idx)| idx % 2 == 0);
238
239        assert_eq!(vec.len(), 2);
240        assert_eq!(vec.get(&0), Complex64::new(1.0, 0.0));
241        assert_eq!(vec.get(&2), Complex64::new(3.0, 0.0));
242        assert_eq!(vec.get(&1), Complex64::new(0.0, 0.0)); // Removed
243    }
244
245    #[test]
246    fn test_normalize() {
247        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
248        vec.unsafe_insert(0, Complex64::new(3.0, 0.0));
249        vec.unsafe_insert(1, Complex64::new(4.0, 0.0));
250
251        // Norm should be sqrt(3^2 + 4^2) = 5
252        vec.normalize();
253
254        // After normalization, norm should be 1
255        let norm: f64 = vec
256            .iter()
257            .map(|(v, _)| v.abs() * v.abs())
258            .sum::<f64>()
259            .sqrt();
260
261        assert!(
262            (norm - 1.0).abs() < 1e-10,
263            "Norm should be 1 after normalization, got {}",
264            norm
265        );
266
267        // Check values
268        assert!((vec.get(&0).re - 0.6).abs() < 1e-10); // 3/5 = 0.6
269        assert!((vec.get(&1).re - 0.8).abs() < 1e-10); // 4/5 = 0.8
270    }
271
272    #[test]
273    fn test_normalize_complex() {
274        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
275        vec.unsafe_insert(0, Complex64::new(1.0, 1.0));
276        vec.unsafe_insert(1, Complex64::new(1.0, -1.0));
277
278        // |1+i| + |1-i| = sqrt(2) + sqrt(2) = 2*sqrt(2)
279        // After normalization: 1/2 * sqrt(2)
280        vec.normalize();
281
282        let norm: f64 = vec
283            .iter()
284            .map(|(v, _)| v.abs() * v.abs())
285            .sum::<f64>()
286            .sqrt();
287
288        assert!(
289            (norm - 1.0).abs() < 1e-10,
290            "Norm should be 1 after normalization"
291        );
292    }
293
294    #[test]
295    fn test_normalize_already_normalized() {
296        let mut vec: Vec<(Complex64, usize)> = SparseVector::new();
297        vec.unsafe_insert(0, Complex64::new(1.0, 0.0));
298
299        vec.normalize(); // Should return early since norm is already 1
300
301        assert_eq!(vec.get(&0), Complex64::new(1.0, 0.0));
302    }
303
304    use bnum::types::U256;
305    #[test]
306    fn test_bigint() {
307        let mut vec: Vec<(Complex64, U256)> = SparseVector::new();
308
309        vec.unsafe_insert(U256::from(0u8), Complex64::new(1.0, 0.0));
310
311        vec.normalize(); // Should return early since norm is already 1
312
313        assert_eq!(vec.get(&U256::from(0u8)), Complex64::new(1.0, 0.0));
314    }
315}