1use std::borrow::Cow;
4
5use num_complex::Complex64;
6use rayon::prelude::*;
7
8use super::hash::Gf2Hash;
9use crate::pauli_string::PauliString;
10
11pub const DEFAULT_HASH_SEED: u64 = 0x9E37_79B9_7F4A_7C15;
13
14pub 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
31pub(crate) const MIN_TERMS_PER_TASK: usize = 64;
34
35const PARALLEL_MIN_TERMS: usize = DEFAULT_MIN_BUCKETS * MIN_TERMS_PER_TASK;
37
38pub const DEFAULT_TARGET_BUCKET_LEN: usize = 1024;
40
41pub const DEFAULT_MIN_BUCKETS: usize = 128;
44
45#[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#[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#[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
125pub(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
142fn 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
166pub(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
198pub(super) fn keys_are_disjoint(_: Complex64, _: Complex64) -> Option<Complex64> {
200 debug_assert!(false, "PauliSum: duplicate key across merged runs");
201 None
202}
203
204#[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 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 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 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 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 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 pub fn clear(&mut self) {
311 for columns in self.buckets.iter_mut() {
312 columns.clear();
313 }
314 self.len = 0;
315 }
316
317 pub fn len(&self) -> usize {
319 self.len
320 }
321
322 pub fn is_empty(&self) -> bool {
324 self.len == 0
325 }
326
327 pub fn num_qubits(&self) -> usize {
329 self.num_qubits
330 }
331
332 pub fn num_buckets(&self) -> usize {
334 self.buckets.len()
335 }
336
337 pub fn hash(&self) -> &Gf2Hash<W> {
339 &self.hash
340 }
341
342 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 pub(crate) fn bucket_len(&self, b: usize) -> usize {
350 self.buckets[b].len()
351 }
352
353 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 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 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 pub(crate) fn buckets(&self) -> &[BucketColumns<W>] {
387 &self.buckets
388 }
389
390 pub(crate) fn buckets_mut(&mut self) -> &mut [BucketColumns<W>] {
392 &mut self.buckets
393 }
394
395 pub(crate) fn recount(&mut self) {
397 self.len = self.buckets.iter().map(|c| c.len()).sum();
398 }
399
400 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 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 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 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 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 pub fn retain(&mut self, f: impl Fn(&[u64; W], &[u64; W], Complex64) -> bool + Sync) {
458 self.retain_per_bucket(|_| &f);
459 }
460
461 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 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 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 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 #[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;