paulistrings/truncation/
tree.rs

1//! [`BuiltinTruncation`], every builtin policy and combinator as one runtime value.
2
3use super::builtin::{
4    or_operand_error, And, ApproxTopN, CoefficientThreshold, CollapseSample, Or, TopN, WeightCutoff,
5};
6use super::TruncationPolicy;
7use crate::collectives::Collectives;
8use crate::pauli_sum::PauliSum;
9use num_complex::Complex64;
10use std::sync::Arc;
11
12/// Any builtin policy or combinator as one runtime value, the form the device drivers take.
13///
14/// Each variant truncates exactly as the builtin it names, and every builtin converts into one with [`From`]; a custom [`TruncationPolicy`] does not.
15/// A tree holding an exact `TopN` cannot run above one partition ([`supports_partitioned`](Self::supports_partitioned) is `false`).
16/// `Or` combines per-term filters only: build it with [`or`](Self::or) or [`try_or`](Self::try_or), which reject an operand with a layer pass; a hand-built `Or` variant that holds one panics in [`finalizes_layer`](TruncationPolicy::finalizes_layer) and [`supports_partitioned`](Self::supports_partitioned).
17/// `Clone` shares a [`CollapseSample`] rather than copying it, so every clone continues the same trajectory.
18#[derive(Clone, Debug, PartialEq)]
19pub enum BuiltinTruncation {
20    /// No truncation.
21    Keep,
22    /// [`CoefficientThreshold`]`(eps)`.
23    Coefficient(f64),
24    /// [`WeightCutoff`]`(k)`.
25    Weight(u32),
26    /// [`TopN`]`(n)`.
27    TopN(usize),
28    /// [`ApproxTopN`]`(n)`.
29    ApproxTopN(usize),
30    /// [`CollapseSample`].
31    CollapseSample(Arc<CollapseSample>),
32    /// [`And`] of two policies.
33    And(Box<BuiltinTruncation>, Box<BuiltinTruncation>),
34    /// [`Or`] of two per-term filters.
35    Or(Box<BuiltinTruncation>, Box<BuiltinTruncation>),
36}
37
38impl BuiltinTruncation {
39    /// `first` or `second`, both per-term filters.
40    ///
41    /// # Panics
42    ///
43    /// If either operand has a layer pass (see [`try_or`](Self::try_or)).
44    pub fn or(first: BuiltinTruncation, second: BuiltinTruncation) -> Self {
45        Self::try_or(first, second).unwrap_or_else(|message| panic!("{message}"))
46    }
47
48    /// `first` or `second`, or the reason it is rejected: an operand with a layer pass (`TopN`, `ApproxTopN`, `CollapseSample`, or an `And` containing one).
49    pub fn try_or(first: BuiltinTruncation, second: BuiltinTruncation) -> Result<Self, String> {
50        match or_operand_error(first.has_layer_pass(), second.has_layer_pass()) {
51            Some(message) => Err(message),
52            None => Ok(Self::Or(Box::new(first), Box::new(second))),
53        }
54    }
55
56    /// Whether the tree runs a layer pass; panics on an `Or` holding one.
57    fn has_layer_pass(&self) -> bool {
58        match self {
59            Self::TopN(_) | Self::ApproxTopN(_) | Self::CollapseSample(_) => true,
60            Self::And(a, b) => a.has_layer_pass() || b.has_layer_pass(),
61            Self::Or(a, b) => {
62                if let Some(message) = or_operand_error(a.has_layer_pass(), b.has_layer_pass()) {
63                    panic!("{message}");
64                }
65                false
66            }
67            Self::Keep | Self::Coefficient(_) | Self::Weight(_) => false,
68        }
69    }
70
71    /// Whether the tree runs above one partition: `false` exactly when it holds an exact `TopN`.
72    ///
73    /// # Panics
74    ///
75    /// On an `Or` variant holding a layer pass.
76    pub fn supports_partitioned(&self) -> bool {
77        match self {
78            Self::TopN(_) => false,
79            Self::And(a, b) => a.supports_partitioned() && b.supports_partitioned(),
80            Self::Or(..) => !self.has_layer_pass(),
81            Self::Keep
82            | Self::Coefficient(_)
83            | Self::Weight(_)
84            | Self::ApproxTopN(_)
85            | Self::CollapseSample(_) => true,
86        }
87    }
88
89    /// Whether a [`CollapseSample`] appears anywhere in the tree, `Or` branches included.
90    pub fn contains_collapse_sample(&self) -> bool {
91        match self {
92            Self::CollapseSample(_) => true,
93            Self::And(a, b) | Self::Or(a, b) => {
94                a.contains_collapse_sample() || b.contains_collapse_sample()
95            }
96            Self::Keep
97            | Self::Coefficient(_)
98            | Self::Weight(_)
99            | Self::TopN(_)
100            | Self::ApproxTopN(_) => false,
101        }
102    }
103}
104
105impl<const W: usize> TruncationPolicy<W> for BuiltinTruncation {
106    fn keep_term(&self, x: &[u64; W], z: &[u64; W], c: Complex64) -> bool {
107        match self {
108            Self::Keep => true,
109            Self::Coefficient(eps) => CoefficientThreshold(*eps).keep_term(x, z, c),
110            Self::Weight(k) => WeightCutoff(*k).keep_term(x, z, c),
111            Self::TopN(n) => <TopN as TruncationPolicy<W>>::keep_term(&TopN(*n), x, z, c),
112            Self::ApproxTopN(n) => {
113                <ApproxTopN as TruncationPolicy<W>>::keep_term(&ApproxTopN(*n), x, z, c)
114            }
115            Self::CollapseSample(s) => {
116                <CollapseSample as TruncationPolicy<W>>::keep_term(s, x, z, c)
117            }
118            Self::And(a, b) => a.keep_term(x, z, c) && b.keep_term(x, z, c),
119            Self::Or(a, b) => a.keep_term(x, z, c) || b.keep_term(x, z, c),
120        }
121    }
122
123    fn finalize_layer(&self, sum: &mut PauliSum<W>) {
124        match self {
125            Self::TopN(n) => TopN(*n).finalize_layer(sum),
126            Self::ApproxTopN(n) => ApproxTopN(*n).finalize_layer(sum),
127            Self::CollapseSample(s) => {
128                <CollapseSample as TruncationPolicy<W>>::finalize_layer(s, sum)
129            }
130            Self::And(a, b) => {
131                a.finalize_layer(sum);
132                b.finalize_layer(sum);
133            }
134            Self::Keep | Self::Coefficient(_) | Self::Weight(_) | Self::Or(_, _) => {}
135        }
136    }
137
138    fn finalizes_layer(&self) -> bool {
139        self.has_layer_pass()
140    }
141
142    /// Arm for arm [`finalize_layer`](TruncationPolicy::finalize_layer), each leaf by its own collective form.
143    fn finalize_layer_partitioned(&self, local: &mut PauliSum<W>, collectives: &dyn Collectives) {
144        match self {
145            Self::TopN(n) => <TopN as TruncationPolicy<W>>::finalize_layer_partitioned(
146                &TopN(*n),
147                local,
148                collectives,
149            ),
150            Self::ApproxTopN(n) => <ApproxTopN as TruncationPolicy<W>>::finalize_layer_partitioned(
151                &ApproxTopN(*n),
152                local,
153                collectives,
154            ),
155            Self::CollapseSample(s) => {
156                <CollapseSample as TruncationPolicy<W>>::finalize_layer_partitioned(
157                    s,
158                    local,
159                    collectives,
160                )
161            }
162            Self::And(a, b) => {
163                <Self as TruncationPolicy<W>>::finalize_layer_partitioned(a, local, collectives);
164                <Self as TruncationPolicy<W>>::finalize_layer_partitioned(b, local, collectives);
165            }
166            Self::Keep | Self::Coefficient(_) | Self::Weight(_) | Self::Or(_, _) => {}
167        }
168    }
169
170    fn supports_partitioned(&self) -> bool {
171        BuiltinTruncation::supports_partitioned(self)
172    }
173}
174
175impl From<CoefficientThreshold> for BuiltinTruncation {
176    fn from(p: CoefficientThreshold) -> Self {
177        Self::Coefficient(p.0)
178    }
179}
180
181impl From<WeightCutoff> for BuiltinTruncation {
182    fn from(p: WeightCutoff) -> Self {
183        Self::Weight(p.0)
184    }
185}
186
187impl From<TopN> for BuiltinTruncation {
188    fn from(p: TopN) -> Self {
189        Self::TopN(p.0)
190    }
191}
192
193impl From<ApproxTopN> for BuiltinTruncation {
194    fn from(p: ApproxTopN) -> Self {
195        Self::ApproxTopN(p.0)
196    }
197}
198
199impl From<CollapseSample> for BuiltinTruncation {
200    fn from(p: CollapseSample) -> Self {
201        Self::CollapseSample(Arc::new(p))
202    }
203}
204
205impl From<Arc<CollapseSample>> for BuiltinTruncation {
206    fn from(p: Arc<CollapseSample>) -> Self {
207        Self::CollapseSample(p)
208    }
209}
210
211impl<A: Into<BuiltinTruncation>, B: Into<BuiltinTruncation>> From<And<A, B>> for BuiltinTruncation {
212    fn from(p: And<A, B>) -> Self {
213        Self::And(Box::new(p.0.into()), Box::new(p.1.into()))
214    }
215}
216
217/// # Panics
218///
219/// If either operand has a layer pass, as [`BuiltinTruncation::or`].
220impl<A: Into<BuiltinTruncation>, B: Into<BuiltinTruncation>> From<Or<A, B>> for BuiltinTruncation {
221    fn from(p: Or<A, B>) -> Self {
222        Self::or(p.0.into(), p.1.into())
223    }
224}
225
226/// A borrowed policy converts as its clone does.
227impl<P: Clone + Into<BuiltinTruncation>> From<&P> for BuiltinTruncation {
228    fn from(p: &P) -> Self {
229        p.clone().into()
230    }
231}
232
233#[cfg(test)]
234mod tests;
paulistrings/truncation/tree.rs

paulistrings/truncation/
tree.rs
