gammalooprs/integrands/
mod.rs

1pub mod evaluation;
2pub mod process;
3
4use crate::integrands::evaluation::{EvaluationResult, RawBatchEvaluationResult};
5use crate::integrands::process::{EvaluationTarget, ProcessIntegrand};
6use crate::integrands::process::{amplitude, cross_section};
7use crate::model::Model;
8use crate::observables::{
9    ObservableAccumulatorBundle, ObservableFileFormat, ObservableSnapshotBundle,
10};
11use crate::settings::runtime::IntegratorSettings;
12use crate::utils::F;
13#[cfg(test)]
14use crate::{is_interrupted, settings::RuntimeSettings};
15
16use color_eyre::Result;
17use enum_dispatch::enum_dispatch;
18use spenso::algebra::complex::Complex;
19#[cfg(test)]
20use symbolica::numerical_integration::ContinuousGrid;
21use symbolica::numerical_integration::{Grid, Sample};
22
23#[enum_dispatch]
24pub trait HasIntegrand {
25    /// Constructs the continuous or discrete integration grid required by this integrand.
26    fn create_grid(&self) -> Grid<F<f64>>;
27
28    /// Returns the integrand name used in state inspection, logs, and status output.
29    fn name(&self) -> String;
30
31    /// Evaluates one weighted grid sample against the supplied model and numerical controls.
32    fn evaluate_sample(
33        &mut self,
34        sample: &Sample<F<f64>>,
35        model: &Model,
36        wgt: F<f64>,
37        iter: usize,
38        use_arb_prec: bool,
39        max_eval: Complex<F<f64>>,
40    ) -> Result<EvaluationResult>;
41
42    /// Returns the number of continuous sampling dimensions used by the integrand.
43    fn get_n_dim(&self) -> usize;
44
45    /// Returns integrator settings specialized for this integrand, or the defaults.
46    fn get_integrator_settings(&self) -> IntegratorSettings {
47        IntegratorSettings::default()
48    }
49
50    /// Merges iteration-local result state from another worker integrand.
51    // In case your integrand supports observable, then overload this function to combine the observables
52    fn merge_results<I: HasIntegrand>(&mut self, _other: &mut I, _iter: usize) {}
53
54    /// Finalizes or writes iteration-local result state after an integration iteration.
55    // In case your integrand supports observable, then overload this function to write the observables to file
56    fn update_results(&mut self, _iter: usize) {}
57}
58
59/// The only production integrand container. Sampling acceptance tests use
60/// [`ProcessIntegrand`] with a reference overlay, so no standalone unit-volume
61/// or profile integrand is kept as a second owner of parameterization logic.
62#[derive(Clone)]
63#[allow(clippy::large_enum_variant)]
64pub enum Integrand {
65    /// Generated amplitude or cross-section process integrand.
66    ProcessIntegrand(Box<ProcessIntegrand>),
67    #[cfg(test)]
68    TestProbe(TestProbeIntegrand),
69}
70
71impl Integrand {
72    pub fn evaluate_samples_raw(
73        &mut self,
74        samples: &[Sample<F<f64>>],
75        target: EvaluationTarget<'_>,
76        iter: usize,
77        use_arb_prec: bool,
78        stop_on_interrupt: bool,
79        max_eval: Complex<F<f64>>,
80    ) -> Result<RawBatchEvaluationResult> {
81        match self {
82            Integrand::ProcessIntegrand(integrand) => integrand.evaluate_samples_raw(
83                target,
84                samples,
85                iter,
86                use_arb_prec,
87                stop_on_interrupt,
88                max_eval,
89            ),
90            #[cfg(test)]
91            Integrand::TestProbe(_) => {
92                let EvaluationTarget::Physical(model) = target else {
93                    return Err(color_eyre::eyre::eyre!(
94                        "reference overlays require a generated process"
95                    ));
96                };
97                let mut results = Vec::with_capacity(samples.len());
98                for sample in samples {
99                    if stop_on_interrupt && is_interrupted() {
100                        break;
101                    }
102                    results.push(self.evaluate_sample(
103                        sample,
104                        model,
105                        sample.get_weight(),
106                        iter,
107                        use_arb_prec,
108                        max_eval,
109                    )?);
110                }
111                Ok(RawBatchEvaluationResult {
112                    statistics: evaluation::StatisticsCounter::from_evaluation_results(&results),
113                    samples: results,
114                })
115            }
116        }
117    }
118
119    #[allow(clippy::too_many_arguments)]
120    #[allow(clippy::too_many_arguments)]
121    pub fn evaluate_samples_raw_with_estimate(
122        &mut self,
123        samples: &[Sample<F<f64>>],
124        target: EvaluationTarget<'_>,
125        iter: usize,
126        use_arb_prec: bool,
127        stop_on_interrupt: bool,
128        max_eval: Complex<F<f64>>,
129        integral_estimate: Option<(f64, f64)>,
130    ) -> Result<RawBatchEvaluationResult> {
131        match self {
132            Integrand::ProcessIntegrand(integrand) => integrand.evaluate_samples_raw_with_estimate(
133                target,
134                samples,
135                iter,
136                use_arb_prec,
137                stop_on_interrupt,
138                max_eval,
139                integral_estimate,
140            ),
141            #[cfg(test)]
142            Integrand::TestProbe(_) => self.evaluate_samples_raw(
143                samples,
144                target,
145                iter,
146                use_arb_prec,
147                stop_on_interrupt,
148                max_eval,
149            ),
150        }
151    }
152
153    pub fn process_evaluation_result(&mut self, result: &EvaluationResult) {
154        match self {
155            Integrand::ProcessIntegrand(integrand) => integrand.process_evaluation_result(result),
156            #[cfg(test)]
157            Integrand::TestProbe(_) => {}
158        }
159    }
160
161    pub fn merge_runtime_results(&mut self, other: &mut Integrand) -> Result<()> {
162        match (self, other) {
163            (Integrand::ProcessIntegrand(lhs), Integrand::ProcessIntegrand(rhs)) => {
164                lhs.merge_event_processing_runtime(rhs)
165            }
166            #[cfg(test)]
167            _ => Ok(()),
168        }
169    }
170
171    pub fn update_runtime_results(&mut self, iter: usize) {
172        match self {
173            Integrand::ProcessIntegrand(integrand) => {
174                integrand.update_event_processing_runtime(iter)
175            }
176            #[cfg(test)]
177            Integrand::TestProbe(_) => {}
178        }
179    }
180
181    pub fn observable_accumulator_bundle(&self) -> Option<ObservableAccumulatorBundle> {
182        match self {
183            Integrand::ProcessIntegrand(integrand) => integrand.observable_accumulator_bundle(),
184            #[cfg(test)]
185            Integrand::TestProbe(_) => None,
186        }
187    }
188
189    pub fn has_observables(&self) -> bool {
190        match self {
191            Integrand::ProcessIntegrand(integrand) => integrand.has_observables(),
192            #[cfg(test)]
193            Integrand::TestProbe(_) => false,
194        }
195    }
196
197    pub fn observable_snapshot_bundle(&self) -> Option<ObservableSnapshotBundle> {
198        match self {
199            Integrand::ProcessIntegrand(integrand) => integrand.observable_snapshot_bundle(),
200            #[cfg(test)]
201            Integrand::TestProbe(_) => None,
202        }
203    }
204
205    pub fn build_observable_snapshots_for_result(
206        &self,
207        result: &EvaluationResult,
208    ) -> Option<ObservableSnapshotBundle> {
209        match self {
210            Integrand::ProcessIntegrand(integrand) => {
211                integrand.build_observable_snapshots_for_result(result)
212            }
213            #[cfg(test)]
214            Integrand::TestProbe(_) => None,
215        }
216    }
217
218    pub fn write_observable_snapshots(
219        &self,
220        path: impl AsRef<std::path::Path>,
221        format: ObservableFileFormat,
222    ) -> Result<()> {
223        match self {
224            Integrand::ProcessIntegrand(integrand) => {
225                integrand.write_observable_snapshots(path, format)
226            }
227            #[cfg(test)]
228            Integrand::TestProbe(_) => Ok(()),
229        }
230    }
231}
232
233impl HasIntegrand for Integrand {
234    fn name(&self) -> String {
235        match self {
236            Integrand::ProcessIntegrand(i) => i.name(),
237            #[cfg(test)]
238            Integrand::TestProbe(integrand) => integrand.name(),
239        }
240    }
241
242    fn create_grid(&self) -> Grid<F<f64>> {
243        match self {
244            Integrand::ProcessIntegrand(integrand) => integrand.create_grid(),
245            #[cfg(test)]
246            Integrand::TestProbe(integrand) => integrand.create_grid(),
247        }
248    }
249
250    fn evaluate_sample(
251        &mut self,
252        sample: &Sample<F<f64>>,
253        model: &Model,
254        wgt: F<f64>,
255        iter: usize,
256        use_arb_prec: bool,
257        max_eval: Complex<F<f64>>,
258    ) -> Result<EvaluationResult> {
259        match self {
260            Integrand::ProcessIntegrand(integrand) => {
261                integrand.evaluate_sample(sample, model, wgt, iter, use_arb_prec, max_eval)
262            }
263            #[cfg(test)]
264            Integrand::TestProbe(integrand) => {
265                integrand.evaluate_sample(sample, model, wgt, iter, use_arb_prec, max_eval)
266            }
267        }
268    }
269
270    fn get_n_dim(&self) -> usize {
271        match self {
272            Integrand::ProcessIntegrand(integrand) => integrand.get_n_dim(),
273            #[cfg(test)]
274            Integrand::TestProbe(integrand) => integrand.get_n_dim(),
275        }
276    }
277
278    fn get_integrator_settings(&self) -> IntegratorSettings {
279        match self {
280            Integrand::ProcessIntegrand(integrand) => integrand.get_integrator_settings(),
281            #[cfg(test)]
282            Integrand::TestProbe(integrand) => integrand.get_integrator_settings(),
283        }
284    }
285}
286
287/// Minimal test-only integrand used for exercising the integration state and
288/// monitoring machinery. It intentionally has no standalone production owner;
289/// process-level acceptance tests cover actual sampling and Jacobians.
290#[cfg(test)]
291#[derive(Clone)]
292pub struct TestProbeIntegrand {
293    settings: RuntimeSettings,
294    n_dim: usize,
295}
296
297#[cfg(test)]
298impl TestProbeIntegrand {
299    pub(crate) fn new(settings: RuntimeSettings, n_dim: usize) -> Self {
300        Self { settings, n_dim }
301    }
302}
303
304#[cfg(test)]
305impl HasIntegrand for TestProbeIntegrand {
306    fn create_grid(&self) -> Grid<F<f64>> {
307        Grid::Continuous(
308            ContinuousGrid::new(
309                self.n_dim,
310                self.settings.integrator.n_bins,
311                self.settings.integrator.min_samples_for_update,
312                self.settings.integrator.bin_number_evolution.clone(),
313                self.settings.integrator.train_on_avg,
314            )
315            .expect("test-probe integration requires valid continuous-grid settings"),
316        )
317    }
318
319    fn name(&self) -> String {
320        "TestProbeIntegrand".to_string()
321    }
322
323    fn evaluate_sample(
324        &mut self,
325        _sample: &Sample<F<f64>>,
326        _model: &Model,
327        wgt: F<f64>,
328        _iter: usize,
329        _use_arb_prec: bool,
330        _max_eval: Complex<F<f64>>,
331    ) -> Result<EvaluationResult> {
332        let mut result = EvaluationResult::zero();
333        result.integrand_result = Complex::new_re(F(1.0));
334        result.parameterization_jacobian = Some(F(1.0));
335        result.integrator_weight = wgt;
336        Ok(result)
337    }
338
339    fn get_n_dim(&self) -> usize {
340        self.n_dim
341    }
342}
gammalooprs/integrands/mod.rs

gammalooprs/integrands/
mod.rs
