paulistrings/pauli_sum/
storage.rs

1//! [`PauliSum`]'s per-bucket column storage, the bucket-count policy and the merge helpers.
2
3use std::borrow::Cow;
4
5use num_complex::Complex64;
6use rayon::prelude::*;
7
8use super::hash::Gf2Hash;
9use crate::pauli_string::PauliString;
10
11/// Default seed for the partitioning hash.
12pub const DEFAULT_HASH_SEED: u64 = 0x9E37_79B9_7F4A_7C15;
13
14/// Bucket bits for `len` terms: the smallest `b` with `len <= target << b`, clamped below by the parallelism floor.
15pub fn desired_bits(len: usize, target: usize, min_buckets: usize) -> u8 {
16    debug_assert!(target > 0);
17    let worth_splitting = len >= min_buckets.saturating_mul(MIN_TERMS_PER_TASK);
18    let mut floor = 0u8;
19    if worth_splitting {
20        while (1usize << floor) < min_buckets && floor < super::hash::B_MAX_BITS {
21            floor += 1;
22        }
23    }
24    let mut b = floor;
25    while b < super::hash::B_MAX_BITS && len > target.saturating_mul(1usize << b) {
26        b += 1;
27    }
28    b
29}
30
31/// Minimum terms per bucket for the parallelism floor to apply.
32// Not lower without checking a dense and a sparse PTM layer: the bucket count also caps the coset dimension a dense layer's sort needs (research/FINDINGS.md §The dense-PTM bucket cliff is a delta-span rank effect).
33pub(crate) const MIN_TERMS_PER_TASK: usize = 64;
34
35/// Below this many terms the bucket-parallel maintenance passes run serially.
36const PARALLEL_MIN_TERMS: usize = DEFAULT_MIN_BUCKETS * MIN_TERMS_PER_TASK;
37
38/// Default target terms per bucket.
39pub const DEFAULT_TARGET_BUCKET_LEN: usize = 1024;
40
41/// Default floor on the bucket count, fixed rather than thread-derived, so the partition is a function of the sum alone.
42/// Must be `>= 16`, or `desired_bits` stops giving a sum of `<= 1024` terms a single bucket.
43pub const DEFAULT_MIN_BUCKETS: usize = 128;
44
45/// One bucket's columns, whose capacity is retained across layers.
46#[derive(Clone, Debug, Default)]
47pub(crate) struct BucketColumns<const W: usize> {
48    pub(crate) x: Vec<[u64; W]>,
49    pub(crate) z: Vec<[u64; W]>,
50    pub(crate) coeff: Vec<Complex64>,
51}
52
53impl<const W: usize> BucketColumns<W> {
54    pub(super) fn new() -> Self {
55        Self::with_capacity(0)
56    }
57
58    pub(super) fn with_capacity(capacity: usize) -> Self {
59        Self {
60            x: Vec::with_capacity(capacity),
61            z: Vec::with_capacity(capacity),
62            coeff: Vec::with_capacity(capacity),
63        }
64    }
65
66    pub(crate) fn len(&self) -> usize {
67        self.coeff.len()
68    }
69
70    pub(crate) fn clear(&mut self) {
71        self.x.clear();
72        self.z.clear();
73        self.coeff.clear();
74    }
75
76    pub(crate) fn push(&mut self, x: [u64; W], z: [u64; W], c: Complex64) {
77        self.x.push(x);
78        self.z.push(z);
79        self.coeff.push(c);
80    }
81
82    fn extend_from(&mut self, other: &Self, start: usize) {
83        self.x.extend_from_slice(&other.x[start..]);
84        self.z.extend_from_slice(&other.z[start..]);
85        self.coeff.extend_from_slice(&other.coeff[start..]);
86    }
87
88    pub(crate) fn view(&self) -> ColumnsRef<'_, W> {
89        ColumnsRef { x: &self.x, z: &self.z, coeff: &self.coeff }
90    }
91}
92
93/// Borrowed `(x, z, coeff)` columns of equal length, which need not come from one [`BucketColumns`].
94#[derive(Clone, Copy, Debug)]
95pub(crate) struct ColumnsRef<'a, const W: usize> {
96    pub(crate) x: &'a [[u64; W]],
97    pub(crate) z: &'a [[u64; W]],
98    pub(crate) coeff: &'a [Complex64],
99}
100
101/// Mutable `(x, z, coeff)` columns of equal length, which need not come from one [`BucketColumns`].
102#[derive(Debug)]
103pub(crate) struct ColumnsMut<'a, const W: usize> {
104    pub(crate) x: &'a mut [[u64; W]],
105    pub(crate) z: &'a mut [[u64; W]],
106    pub(crate) coeff: &'a mut [Complex64],
107}
108
109impl<'a, const W: usize> ColumnsMut<'a, W> {
110    pub(crate) fn len(&self) -> usize {
111        self.coeff.len()
112    }
113
114    pub(crate) fn split_at(self, at: usize) -> (ColumnsMut<'a, W>, ColumnsMut<'a, W>) {
115        let (x_head, x_tail) = self.x.split_at_mut(at);
116        let (z_head, z_tail) = self.z.split_at_mut(at);
117        let (coeff_head, coeff_tail) = self.coeff.split_at_mut(at);
118        (
119            ColumnsMut { x: x_head, z: z_head, coeff: coeff_head },
120            ColumnsMut { x: x_tail, z: z_tail, coeff: coeff_tail },
121        )
122    }
123}
124
125/// `f` over every item, on the Rayon pool once the sum holds `PARALLEL_MIN_TERMS` terms.
126pub(super) fn map_buckets<C, T>(
127    num_terms: usize,
128    items: C,
129    f: impl Fn(<C as IntoParallelIterator>::Item) -> T + Sync + Send,
130) -> Vec<T>
131where
132    C: IntoParallelIterator + IntoIterator<Item = <C as IntoParallelIterator>::Item>,
133    T: Send,
134{
135    if num_terms < PARALLEL_MIN_TERMS {
136        items.into_iter().map(f).collect()
137    } else {
138        items.into_par_iter().map(f).collect()
139    }
140}
141
142/// Split off the terms whose hash row `new_bit` is set, returning them in order and keeping the rest in place.
143fn refine_bucket<const W: usize>(
144    columns: &mut BucketColumns<W>,
145    hash: &Gf2Hash<W>,
146    new_bit: u8,
147) -> BucketColumns<W> {
148    let mut upper_half = BucketColumns::new();
149    let mut keep = 0usize;
150    for i in 0..columns.len() {
151        if hash.row_parity(&columns.x[i], &columns.z[i], new_bit) == 1 {
152            upper_half.push(columns.x[i], columns.z[i], columns.coeff[i]);
153        } else {
154            columns.x[keep] = columns.x[i];
155            columns.z[keep] = columns.z[i];
156            columns.coeff[keep] = columns.coeff[i];
157            keep += 1;
158        }
159    }
160    columns.x.truncate(keep);
161    columns.z.truncate(keep);
162    columns.coeff.truncate(keep);
163    upper_half
164}
165
166/// Merge two sorted runs; `on_equal` combines a shared key's coefficients, `None` dropping the key.
167pub(super) fn merge_sorted<const W: usize>(
168    a: &BucketColumns<W>,
169    b: &BucketColumns<W>,
170    mut on_equal: impl FnMut(Complex64, Complex64) -> Option<Complex64>,
171) -> BucketColumns<W> {
172    let mut out = BucketColumns::with_capacity(a.len() + b.len());
173    let (mut i, mut j) = (0usize, 0usize);
174    while i < a.len() && j < b.len() {
175        match (&a.x[i], &a.z[i]).cmp(&(&b.x[j], &b.z[j])) {
176            std::cmp::Ordering::Less => {
177                out.push(a.x[i], a.z[i], a.coeff[i]);
178                i += 1;
179            }
180            std::cmp::Ordering::Greater => {
181                out.push(b.x[j], b.z[j], b.coeff[j]);
182                j += 1;
183            }
184            std::cmp::Ordering::Equal => {
185                if let Some(c) = on_equal(a.coeff[i], b.coeff[j]) {
186                    out.push(a.x[i], a.z[i], c);
187                }
188                i += 1;
189                j += 1;
190            }
191        }
192    }
193    out.extend_from(a, i);
194    out.extend_from(b, j);
195    out
196}
197
198/// The `on_equal` of [`merge_sorted`] for runs whose keys are disjoint.
199pub(super) fn keys_are_disjoint(_: Complex64, _: Complex64) -> Option<Complex64> {
200    debug_assert!(false, "PauliSum: duplicate key across merged runs");
201    None
202}
203
204/// Weighted sum of Pauli strings, stored as structure-of-arrays columns partitioned into buckets by a GF(2)-linear hash.
205///
206/// Terms come in canonical order: bucket index ascending, then lexicographic `(x, z)` within a bucket, with no repeated key.
207/// A sum of at most 1024 terms built by [`crate::BuildAccumulator`] has one bucket and so is plain lex-sorted; a larger one interleaves buckets, so compare sums by key ([`Self::get`]), not by position.
208#[derive(Clone, Debug)]
209pub struct PauliSum<const W: usize> {
210    pub(super) buckets: Vec<BucketColumns<W>>,
211    pub(super) hash: Gf2Hash<W>,
212    pub(super) num_qubits: usize,
213    pub(super) len: usize,
214}
215
216impl<const W: usize> PauliSum<W> {
217    /// Partition terms the caller has sorted ascending in `(x, z)` with no repeated key; unsorted input silently breaks the bucket invariant.
218    pub(crate) fn from_key_sorted(
219        x: &[[u64; W]],
220        z: &[[u64; W]],
221        coeff: &[Complex64],
222        hash: Gf2Hash<W>,
223        num_qubits: usize,
224    ) -> Self {
225        let n = coeff.len();
226        let num_buckets = hash.num_buckets();
227
228        let bucket_indices: Vec<u32> =
229            (0..n).into_par_iter().map(|i| hash.bucket_of(&x[i], &z[i])).collect();
230        let mut counts: Vec<usize> = vec![0; num_buckets];
231        for &b in bucket_indices.iter() {
232            counts[b as usize] += 1;
233        }
234
235        let mut buckets: Vec<BucketColumns<W>> =
236            counts.iter().map(|&c| BucketColumns::with_capacity(c)).collect();
237
238        for i in 0..n {
239            buckets[bucket_indices[i] as usize].push(x[i], z[i], coeff[i]);
240        }
241
242        Self { buckets, hash, num_qubits, len: n }
243    }
244
245    /// Empty sum on `num_qubits` qubits, in a single bucket.
246    ///
247    /// # Panics
248    ///
249    /// Panics in debug builds if `num_qubits > 64 · W`.
250    pub fn empty(num_qubits: usize) -> Self {
251        debug_assert!(num_qubits <= 64 * W);
252        Self::empty_with_hash(num_qubits, Gf2Hash::new(num_qubits, 0, DEFAULT_HASH_SEED))
253    }
254
255    /// An empty sum over `num_qubits`, partitioned by `hash`.
256    pub(crate) fn empty_with_hash(num_qubits: usize, hash: Gf2Hash<W>) -> Self {
257        let num_buckets = hash.num_buckets();
258        Self {
259            buckets: (0..num_buckets).map(|_| BucketColumns::new()).collect(),
260            hash,
261            num_qubits,
262            len: 0,
263        }
264    }
265
266    /// Repartition under `hash`, keeping every term; panics if `hash` was built for a different qubit count.
267    pub(crate) fn with_hash(self, hash: Gf2Hash<W>) -> Self {
268        assert_eq!(self.num_qubits, hash.num_qubits(), "PauliSum::with_hash: num_qubits mismatch",);
269        let targets: Vec<Vec<u32>> = map_buckets(self.len, &self.buckets, |columns| {
270            (0..columns.len()).map(|i| hash.bucket_of(&columns.x[i], &columns.z[i])).collect()
271        });
272        let mut scattered: Vec<Vec<(PauliString<W>, Complex64)>> =
273            (0..hash.num_buckets()).map(|_| Vec::new()).collect();
274        for (columns, targets) in self.buckets.iter().zip(&targets) {
275            for (i, &b) in targets.iter().enumerate() {
276                let key = PauliString { x: columns.x[i], z: columns.z[i] };
277                scattered[b as usize].push((key, columns.coeff[i]));
278            }
279        }
280        let buckets = map_buckets(self.len, scattered, |mut terms| {
281            terms.sort_unstable_by(|p, q| p.0.cmp(&q.0));
282            let mut columns = BucketColumns::with_capacity(terms.len());
283            for (key, c) in terms {
284                columns.push(key.x, key.z, c);
285            }
286            columns
287        });
288        Self { buckets, hash, num_qubits: self.num_qubits, len: self.len }
289    }
290
291    /// `self` partitioned exactly as `target` partitions, borrowed when it already is.
292    fn align_to(&self, target: &Gf2Hash<W>) -> Cow<'_, Self> {
293        if !self.hash.same_rows_as(target) {
294            return Cow::Owned(self.clone().with_hash(target.clone()));
295        }
296        if self.hash.bits() == target.bits() {
297            return Cow::Borrowed(self);
298        }
299        let mut out = self.clone();
300        while out.hash.bits() < target.bits() {
301            out.refine();
302        }
303        while out.hash.bits() > target.bits() {
304            out.coarsen();
305        }
306        Cow::Owned(out)
307    }
308
309    /// Drop every term, keeping the hash and the bucket storage.
310    pub fn clear(&mut self) {
311        for columns in self.buckets.iter_mut() {
312            columns.clear();
313        }
314        self.len = 0;
315    }
316
317    /// Total number of terms.
318    pub fn len(&self) -> usize {
319        self.len
320    }
321
322    /// `true` if the sum has no terms.
323    pub fn is_empty(&self) -> bool {
324        self.len == 0
325    }
326
327    /// Number of qubits this sum acts on.
328    pub fn num_qubits(&self) -> usize {
329        self.num_qubits
330    }
331
332    /// Number of buckets, `1 << hash().bits()`.
333    pub fn num_buckets(&self) -> usize {
334        self.buckets.len()
335    }
336
337    /// The partitioning hash.
338    pub fn hash(&self) -> &Gf2Hash<W> {
339        &self.hash
340    }
341
342    /// Borrow bucket `b`'s columns as `(x, z, coeff)`.
343    pub(crate) fn bucket(&self, b: usize) -> (&[[u64; W]], &[[u64; W]], &[Complex64]) {
344        let columns = &self.buckets[b];
345        (&columns.x, &columns.z, &columns.coeff)
346    }
347
348    /// Number of terms in bucket `b`.
349    pub(crate) fn bucket_len(&self, b: usize) -> usize {
350        self.buckets[b].len()
351    }
352
353    /// Double the bucket count, splitting each bucket in two by the hash's new row.
354    pub(crate) fn refine(&mut self) {
355        self.hash.refine();
356        let new_bit = self.hash.bits() - 1;
357        let hash = &self.hash;
358        let mut buckets = std::mem::take(&mut self.buckets);
359        let upper =
360            map_buckets(self.len, &mut buckets, |columns| refine_bucket(columns, hash, new_bit));
361        buckets.extend(upper);
362        self.buckets = buckets;
363        #[cfg(debug_assertions)]
364        self.assert_invariants();
365    }
366
367    /// Halve the bucket count, merging bucket pairs `(i, i + B/2)`.
368    pub(crate) fn coarsen(&mut self) {
369        self.hash.coarsen();
370        let (lower, upper) = self.buckets.split_at(self.buckets.len() / 2);
371        self.buckets = map_buckets(self.len, 0..lower.len(), |b| {
372            merge_sorted(&lower[b], &upper[b], keys_are_disjoint)
373        });
374    }
375
376    /// Refine until the bucket count suits `len()` terms at `target` per bucket, with the `min_buckets` floor once the sum is large enough; never coarsens.
377    pub(crate) fn rebucket(&mut self, target: usize, min_buckets: usize) {
378        debug_assert!(target > 0);
379        let want = desired_bits(self.len, target, min_buckets).max(self.hash.bits());
380        while self.hash.bits() < want {
381            self.refine();
382        }
383    }
384
385    /// The buckets' columns, for read-outs that scan them.
386    pub(crate) fn buckets(&self) -> &[BucketColumns<W>] {
387        &self.buckets
388    }
389
390    /// Mutable access to the buckets, for layers that are applied in place.
391    pub(crate) fn buckets_mut(&mut self) -> &mut [BucketColumns<W>] {
392        &mut self.buckets
393    }
394
395    /// Recompute the cached total after an in-place layer.
396    pub(crate) fn recount(&mut self) {
397        self.len = self.buckets.iter().map(|c| c.len()).sum();
398    }
399
400    /// Iterate every term in canonical order, which is not globally key-sorted.
401    pub fn iter(&self) -> impl Iterator<Item = (&[u64; W], &[u64; W], Complex64)> + '_ {
402        self.buckets.iter().flat_map(|columns| {
403            columns
404                .x
405                .iter()
406                .zip(columns.z.iter())
407                .zip(columns.coeff.iter())
408                .map(|((x, z), c)| (x, z, *c))
409        })
410    }
411
412    /// Copy every term out as three parallel columns, in canonical order.
413    pub fn to_arrays(&self) -> (Vec<[u64; W]>, Vec<[u64; W]>, Vec<Complex64>) {
414        let mut x = Vec::with_capacity(self.len);
415        let mut z = Vec::with_capacity(self.len);
416        let mut coeff = Vec::with_capacity(self.len);
417        for columns in self.buckets.iter() {
418            x.extend_from_slice(&columns.x);
419            z.extend_from_slice(&columns.z);
420            coeff.extend_from_slice(&columns.coeff);
421        }
422        (x, z, coeff)
423    }
424
425    /// Coefficient of the term with key `(x, z)`, or `None` if absent.
426    pub fn get(&self, x: &[u64; W], z: &[u64; W]) -> Option<Complex64> {
427        let columns = &self.buckets[self.hash.bucket_of(x, z) as usize];
428        let mut low = 0usize;
429        let mut high = columns.len();
430        while low < high {
431            let mid = low + (high - low) / 2;
432            match (&columns.x[mid], &columns.z[mid]).cmp(&(x, z)) {
433                std::cmp::Ordering::Less => low = mid + 1,
434                std::cmp::Ordering::Greater => high = mid,
435                std::cmp::Ordering::Equal => return Some(columns.coeff[mid]),
436            }
437        }
438        None
439    }
440
441    /// Coefficient of the identity term, i.e. `tr(O) / 2^n`.
442    pub fn identity_coefficient(&self) -> Complex64 {
443        let zero_key = [0u64; W];
444        self.get(&zero_key, &zero_key).unwrap_or(Complex64::new(0.0, 0.0))
445    }
446
447    /// Multiply every coefficient by `c` in place.
448    pub fn scale(&mut self, c: Complex64) {
449        self.buckets.par_iter_mut().for_each(|columns| {
450            for coeff in columns.coeff.iter_mut() {
451                *coeff *= c;
452            }
453        });
454    }
455
456    /// Keep only the terms for which `f(x, z, coeff)` is `true`; `f` may run on several threads at once.
457    pub fn retain(&mut self, f: impl Fn(&[u64; W], &[u64; W], Complex64) -> bool + Sync) {
458        self.retain_per_bucket(|_| &f);
459    }
460
461    /// [`retain`](Self::retain) with one filter per bucket, `filter_for(b)`, run over bucket `b` in order.
462    pub(crate) fn retain_per_bucket<F>(&mut self, filter_for: impl Fn(usize) -> F + Sync)
463    where
464        F: FnMut(&[u64; W], &[u64; W], Complex64) -> bool,
465    {
466        self.update_per_bucket(|b| {
467            let mut keep = filter_for(b);
468            move |x: &[u64; W], z: &[u64; W], c: Complex64| keep(x, z, c).then_some(c)
469        });
470    }
471
472    /// [`retain_per_bucket`](Self::retain_per_bucket) that also rewrites each kept coefficient: `update_for(b)` maps a term to its new coefficient, or to `None` to drop it.
473    pub(crate) fn update_per_bucket<F>(&mut self, update_for: impl Fn(usize) -> F + Sync)
474    where
475        F: FnMut(&[u64; W], &[u64; W], Complex64) -> Option<Complex64>,
476    {
477        self.buckets.par_iter_mut().enumerate().for_each(|(b, columns)| {
478            let mut update = update_for(b);
479            let n = columns.len();
480            let mut w = 0usize;
481            for r in 0..n {
482                if let Some(coeff) = update(&columns.x[r], &columns.z[r], columns.coeff[r]) {
483                    if w != r {
484                        columns.x[w] = columns.x[r];
485                        columns.z[w] = columns.z[r];
486                    }
487                    columns.coeff[w] = coeff;
488                    w += 1;
489                }
490            }
491            columns.x.truncate(w);
492            columns.z.truncate(w);
493            columns.coeff.truncate(w);
494        });
495        self.recount();
496    }
497
498    /// Hilbert-Schmidt overlap `tr(self† · other) / 2ⁿ`, i.e. `Σ conj(aᵢ)·bᵢ` over the keys the two sums share.
499    ///
500    /// # Panics
501    ///
502    /// Panics if the two sums disagree about `num_qubits`.
503    pub fn overlap(&self, other: &Self) -> Complex64 {
504        assert_eq!(
505            self.num_qubits, other.num_qubits,
506            "PauliSum::overlap: num_qubits mismatch ({} vs {})",
507            self.num_qubits, other.num_qubits,
508        );
509        let rhs = other.align_to(&self.hash);
510        self.buckets
511            .par_iter()
512            .zip(rhs.buckets.par_iter())
513            .map(|(a, b)| {
514                let mut partial = Complex64::new(0.0, 0.0);
515                let (mut i, mut j) = (0usize, 0usize);
516                while i < a.len() && j < b.len() {
517                    match (&a.x[i], &a.z[i]).cmp(&(&b.x[j], &b.z[j])) {
518                        std::cmp::Ordering::Less => i += 1,
519                        std::cmp::Ordering::Greater => j += 1,
520                        std::cmp::Ordering::Equal => {
521                            partial += a.coeff[i].conj() * b.coeff[j];
522                            i += 1;
523                            j += 1;
524                        }
525                    }
526                }
527                partial
528            })
529            .collect::<Vec<_>>()
530            .into_iter()
531            .fold(Complex64::new(0.0, 0.0), |a, b| a + b)
532    }
533
534    /// Sum of two sums, partitioned as `self`; terms summing to exactly zero are dropped.
535    ///
536    /// # Panics
537    ///
538    /// Panics if the two sums disagree about `num_qubits`.
539    pub fn add(&self, other: &Self) -> Self {
540        assert_eq!(
541            self.num_qubits, other.num_qubits,
542            "PauliSum::add: num_qubits mismatch ({} vs {})",
543            self.num_qubits, other.num_qubits,
544        );
545        let rhs = other.align_to(&self.hash);
546        let zero = Complex64::new(0.0, 0.0);
547        let buckets: Vec<BucketColumns<W>> = self
548            .buckets
549            .par_iter()
550            .zip(rhs.buckets.par_iter())
551            .map(|(a, b)| merge_sorted(a, b, |p, q| Some(p + q).filter(|c| *c != zero)))
552            .collect();
553        let len = buckets.iter().map(|c| c.len()).sum();
554        Self { buckets, hash: self.hash.clone(), num_qubits: self.num_qubits, len }
555    }
556
557    /// Assert the structural invariant: every term in its hash bucket, each bucket strictly ascending in `(x, z)`, every key within `num_qubits`.
558    #[cfg(any(test, debug_assertions, feature = "test-utils"))]
559    pub(crate) fn assert_invariants(&self) {
560        assert_eq!(
561            self.buckets.len(),
562            self.hash.num_buckets(),
563            "PauliSum: bucket count disagrees with hash",
564        );
565        let mut total = 0usize;
566        for (b, columns) in self.buckets.iter().enumerate() {
567            assert_eq!(columns.x.len(), columns.z.len());
568            assert_eq!(columns.x.len(), columns.coeff.len());
569            total += columns.len();
570            for i in 0..columns.len() {
571                let got = self.hash.bucket_of(&columns.x[i], &columns.z[i]);
572                assert_eq!(got as usize, b, "PauliSum: term {i} of bucket {b} hashes to {got}",);
573                let term = PauliString::<W> { x: columns.x[i], z: columns.z[i] };
574                assert!(
575                    term.is_within(self.num_qubits),
576                    "PauliSum: term {i} of bucket {b} exceeds num_qubits",
577                );
578            }
579            for i in 1..columns.len() {
580                let prev = (&columns.x[i - 1], &columns.z[i - 1]);
581                let cur = (&columns.x[i], &columns.z[i]);
582                assert!(prev < cur, "PauliSum: bucket {b} out of order at {i}");
583            }
584        }
585        assert_eq!(total, self.len, "PauliSum: cached len disagrees");
586    }
587}
588
589#[cfg(test)]
590mod tests;