paulistrings/truncation/
mod.rs

1//! The [`TruncationPolicy`] trait and the built-in policies.
2
3pub(crate) mod builtin;
4mod tree;
5
6pub use builtin::{And, ApproxTopN, CoefficientThreshold, CollapseSample, Or, TopN, WeightCutoff};
7pub use tree::BuiltinTruncation;
8
9use crate::collectives::Collectives;
10use crate::pauli_sum::PauliSum;
11use num_complex::Complex64;
12
13/// A truncation strategy: a per-term filter, a per-layer pass, or both.
14///
15/// Built-ins: [`CoefficientThreshold`] and [`WeightCutoff`] filter terms, [`TopN`], [`ApproxTopN`] and [`CollapseSample`] run a layer pass, [`And`] and [`Or`] compose, and [`BuiltinTruncation`] holds any of them as one value.
16///
17/// Above one partition every partition sees a disjoint slice of the layer, so a layer pass must be collective: a policy whose `finalizes_layer()` is `true` runs partitioned only if it overrides [`finalize_layer_partitioned`](Self::finalize_layer_partitioned) and [`supports_partitioned`](Self::supports_partitioned).
18/// Every built-in does except exact [`TopN`], which is not yet supported partitioned; [`ApproxTopN`] is the partitioned alternative.
19pub trait TruncationPolicy<const W: usize>: Send + Sync {
20    /// Per-term filter, run inside the merge on every output, so it must be cheap.
21    ///
22    /// `c` is the coefficient summed over every contribution to the key `(x, z)`.
23    fn keep_term(&self, _x: &[u64; W], _z: &[u64; W], _c: Complex64) -> bool {
24        true
25    }
26
27    /// Pass over the whole sum after each layer; [`PauliSum::retain`] is the easy way to write a filter here.
28    fn finalize_layer(&self, _sum: &mut PauliSum<W>) {}
29
30    /// Whether [`finalize_layer`](Self::finalize_layer) does anything, a hint that lets the engine skip it.
31    ///
32    /// Return `false` only if `finalize_layer` is a no-op: the engine may then skip it, which is a wrong answer otherwise.
33    fn finalizes_layer(&self) -> bool {
34        true
35    }
36
37    /// The layer pass over one partition's slice of the layer, run on every layer on every partition in lock-step.
38    ///
39    /// The default runs [`finalize_layer`](Self::finalize_layer) at one partition and nothing above it.
40    /// An override must issue the same collective calls on every partition, never skipping one on a partition with nothing to do, and decide only from all-reduced values.
41    ///
42    /// # Panics
43    ///
44    /// Above one partition, if the policy reports [`finalizes_layer()`](Self::finalizes_layer) and has not overridden this method.
45    fn finalize_layer_partitioned(&self, local: &mut PauliSum<W>, collectives: &dyn Collectives) {
46        if collectives.size() == 1 {
47            self.finalize_layer(local);
48            return;
49        }
50        assert!(
51            !self.finalizes_layer(),
52            "{} has a layer pass but no partitioned one: every partition sees only a disjoint \
53             slice of the layer, so the pass has to be collective. Override \
54             `finalize_layer_partitioned` and `supports_partitioned` (see `ApproxTopN`, which \
55             all-reduces its octave histogram), or use `ApproxTopN` itself.",
56            std::any::type_name_of_val(self),
57        );
58    }
59
60    /// Whether [`finalize_layer_partitioned`](Self::finalize_layer_partitioned) runs above one partition; every partitioned driver checks it before the first layer.
61    ///
62    /// The default is `true` exactly when there is no layer pass; a policy that overrides `finalize_layer_partitioned` with a collective pass returns `true`.
63    fn supports_partitioned(&self) -> bool {
64        !self.finalizes_layer()
65    }
66}
67
68/// Panics unless `policy` can run on `partitions` partitions.
69pub(crate) fn assert_supports_partitions<const W: usize, T>(policy: &T, partitions: usize)
70where
71    T: TruncationPolicy<W> + ?Sized,
72{
73    assert!(
74        partitions == 1 || policy.supports_partitioned(),
75        "{} cannot run on {partitions} partitions: its layer pass has no collective form. \
76         Partitioned exact `TopN` is not yet supported; `ApproxTopN` is the partitioned \
77         alternative.",
78        std::any::type_name_of_val(policy),
79    );
80}
81
82#[cfg(test)]
83mod tests;
paulistrings/truncation/mod.rs

paulistrings/truncation/
mod.rs
