gammalooprs/integrands/
mod.rs

1pub mod builtin;
2pub mod evaluation;
3pub mod process;
4
5use crate::integrands::evaluation::{
6    EvaluationMetaData, EvaluationResult, RawBatchEvaluationResult, StabilityResult,
7    StabilityStatus,
8};
9// use crate::integrands::process::ProcessIntegrandImpl;
10use crate::integrands::builtin::h_function::{HFunctionTestIntegrand, HFunctionTestSettings};
11use crate::integrands::process::ProcessIntegrand;
12use crate::integrands::process::{amplitude, cross_section};
13use crate::model::Model;
14use crate::momentum::FourMomentum;
15use crate::observables::{
16    ObservableAccumulatorBundle, ObservableFileFormat, ObservableSnapshotBundle,
17};
18use crate::utils::{F, FloatLike};
19use crate::{
20    is_interrupted,
21    settings::{
22        RuntimeSettings,
23        runtime::{IntegratorSettings, Precision},
24    },
25    utils,
26};
27
28use bincode_trait_derive::{Decode, Encode};
29use color_eyre::Result;
30use enum_dispatch::enum_dispatch;
31use schemars::JsonSchema;
32use serde::{Deserialize, Serialize};
33use spenso::algebra::complex::Complex;
34use std::fmt::{Display, Formatter};
35use std::time::Duration;
36use symbolica::numerical_integration::{ContinuousGrid, Grid, Sample};
37#[allow(unused_imports)]
38use tracing::{debug, error, info, instrument, trace, warn};
39
40#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
41#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode, PartialEq, JsonSchema)]
42// #[trait_decode(trait= GammaLoopContext)]
43#[allow(non_snake_case)]
44#[serde(tag = "type")]
45pub enum IntegrandSettings {
46    #[serde(rename = "unit_surface")]
47    UnitSurface(UnitSurfaceSettings),
48    #[serde(rename = "unit_volume")]
49    UnitVolume(UnitVolumeSettings),
50    #[serde(rename = "h_function_test")]
51    HFunctionTest(HFunctionTestSettings),
52}
53
54impl Display for IntegrandSettings {
55    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
56        match self {
57            IntegrandSettings::UnitSurface(_) => write!(f, "unit_surface"),
58            IntegrandSettings::UnitVolume(_) => write!(f, "unit_volume"),
59            IntegrandSettings::HFunctionTest(_) => {
60                write!(f, "h_function_test")
61            }
62        }
63    }
64}
65
66impl Default for IntegrandSettings {
67    fn default() -> IntegrandSettings {
68        IntegrandSettings::UnitSurface(UnitSurfaceSettings { n_3d_momenta: 11 })
69    }
70}
71
72#[enum_dispatch]
73pub trait HasIntegrand {
74    /// Constructs the continuous or discrete integration grid required by this integrand.
75    fn create_grid(&self) -> Grid<F<f64>>;
76
77    /// Returns the integrand name used in state inspection, logs, and status output.
78    fn name(&self) -> String;
79
80    /// Evaluates one weighted grid sample against the supplied model and numerical controls.
81    fn evaluate_sample(
82        &mut self,
83        sample: &Sample<F<f64>>,
84        model: &Model,
85        wgt: F<f64>,
86        iter: usize,
87        use_arb_prec: bool,
88        max_eval: Complex<F<f64>>,
89    ) -> Result<EvaluationResult>;
90
91    /// Returns the number of continuous sampling dimensions used by the integrand.
92    fn get_n_dim(&self) -> usize;
93
94    /// Returns integrator settings specialized for this integrand, or the defaults.
95    fn get_integrator_settings(&self) -> IntegratorSettings {
96        IntegratorSettings::default()
97    }
98
99    /// Merges iteration-local result state from another worker integrand.
100    // In case your integrand supports observable, then overload this function to combine the observables
101    fn merge_results<I: HasIntegrand>(&mut self, _other: &mut I, _iter: usize) {}
102
103    /// Finalizes or writes iteration-local result state after an integration iteration.
104    // In case your integrand supports observable, then overload this function to write the observables to file
105    fn update_results(&mut self, _iter: usize) {}
106}
107
108#[derive(Clone)]
109pub enum Integrand {
110    /// Built-in diagnostic integrand over a unit hypersurface.
111    UnitSurface(UnitSurfaceIntegrand),
112    /// Built-in diagnostic integrand over a unit integration volume.
113    UnitVolume(UnitVolumeIntegrand),
114    /// One-dimensional diagnostic integrand for the configured H-function profile.
115    HFunctionTest(HFunctionTestIntegrand),
116    // ProcessIntegrandImpl(ProcessIntegrandImpl),
117    /// Generated amplitude or cross-section process integrand.
118    ProcessIntegrand(Box<ProcessIntegrand>),
119}
120
121impl Integrand {
122    pub fn evaluate_samples_raw(
123        &mut self,
124        samples: &[Sample<F<f64>>],
125        model: &Model,
126        iter: usize,
127        use_arb_prec: bool,
128        stop_on_interrupt: bool,
129        max_eval: Complex<F<f64>>,
130    ) -> Result<RawBatchEvaluationResult> {
131        match self {
132            Integrand::ProcessIntegrand(integrand) => integrand.evaluate_samples_raw(
133                model,
134                samples,
135                iter,
136                use_arb_prec,
137                stop_on_interrupt,
138                max_eval,
139            ),
140            _ => {
141                let mut results = Vec::with_capacity(samples.len());
142                for sample in samples {
143                    if stop_on_interrupt && is_interrupted() {
144                        break;
145                    }
146                    results.push(self.evaluate_sample(
147                        sample,
148                        model,
149                        sample.get_weight(),
150                        iter,
151                        use_arb_prec,
152                        max_eval,
153                    )?);
154                    if stop_on_interrupt && is_interrupted() {
155                        break;
156                    }
157                }
158                Ok(RawBatchEvaluationResult {
159                    statistics: evaluation::StatisticsCounter::from_evaluation_results(&results),
160                    samples: results,
161                })
162            }
163        }
164    }
165
166    pub fn process_evaluation_result(&mut self, result: &EvaluationResult) {
167        if let Integrand::ProcessIntegrand(integrand) = self {
168            integrand.process_evaluation_result(result);
169        }
170    }
171
172    pub fn merge_runtime_results(&mut self, other: &mut Integrand) -> Result<()> {
173        match (self, other) {
174            (Integrand::ProcessIntegrand(lhs), Integrand::ProcessIntegrand(rhs)) => {
175                lhs.merge_event_processing_runtime(rhs)
176            }
177            _ => Ok(()),
178        }
179    }
180
181    pub fn update_runtime_results(&mut self, iter: usize) {
182        if let Integrand::ProcessIntegrand(integrand) = self {
183            integrand.update_event_processing_runtime(iter);
184        }
185    }
186
187    pub fn observable_accumulator_bundle(&self) -> Option<ObservableAccumulatorBundle> {
188        match self {
189            Integrand::ProcessIntegrand(integrand) => integrand.observable_accumulator_bundle(),
190            _ => None,
191        }
192    }
193
194    pub fn has_observables(&self) -> bool {
195        match self {
196            Integrand::ProcessIntegrand(integrand) => integrand.has_observables(),
197            _ => false,
198        }
199    }
200
201    pub fn observable_snapshot_bundle(&self) -> Option<ObservableSnapshotBundle> {
202        match self {
203            Integrand::ProcessIntegrand(integrand) => integrand.observable_snapshot_bundle(),
204            _ => None,
205        }
206    }
207
208    pub fn build_observable_snapshots_for_result(
209        &self,
210        result: &EvaluationResult,
211    ) -> Option<ObservableSnapshotBundle> {
212        match self {
213            Integrand::ProcessIntegrand(integrand) => {
214                integrand.build_observable_snapshots_for_result(result)
215            }
216            _ => None,
217        }
218    }
219
220    pub fn write_observable_snapshots(
221        &self,
222        path: impl AsRef<std::path::Path>,
223        format: ObservableFileFormat,
224    ) -> Result<()> {
225        match self {
226            Integrand::ProcessIntegrand(integrand) => {
227                integrand.write_observable_snapshots(path, format)
228            }
229            _ => Ok(()),
230        }
231    }
232}
233
234impl HasIntegrand for Integrand {
235    fn name(&self) -> String {
236        match self {
237            Integrand::UnitSurface(_) => "UnitSurface".to_string(),
238            Integrand::UnitVolume(_) => "UnitVolume".to_string(),
239            Integrand::HFunctionTest(_) => "HFunctionTest".to_string(),
240            // Integrand::ProcessIntegrandImpl(_) => "ProcessIntegrandImpl".to_string(),
241            Integrand::ProcessIntegrand(i) => i.name(),
242        }
243    }
244
245    fn create_grid(&self) -> Grid<F<f64>> {
246        match self {
247            Integrand::UnitSurface(integrand) => integrand.create_grid(),
248            Integrand::UnitVolume(integrand) => integrand.create_grid(),
249            Integrand::HFunctionTest(integrand) => integrand.create_grid(),
250            // Integrand::ProcessIntegrandImpl(integrand) => integrand.create_grid(),
251            Integrand::ProcessIntegrand(integrand) => integrand.create_grid(),
252        }
253    }
254
255    fn evaluate_sample(
256        &mut self,
257        sample: &Sample<F<f64>>,
258        model: &Model,
259        wgt: F<f64>,
260        iter: usize,
261        use_arb_prec: bool,
262        max_eval: Complex<F<f64>>,
263    ) -> Result<EvaluationResult> {
264        match self {
265            Integrand::UnitSurface(integrand) => {
266                integrand.evaluate_sample(sample, model, wgt, iter, use_arb_prec, max_eval)
267            }
268            Integrand::UnitVolume(integrand) => {
269                integrand.evaluate_sample(sample, model, wgt, iter, use_arb_prec, max_eval)
270            }
271            Integrand::HFunctionTest(integrand) => {
272                integrand.evaluate_sample(sample, model, wgt, iter, use_arb_prec, max_eval)
273            }
274            // Integrand::ProcessIntegrandImpl(integrand) => {
275            //     integrand.evaluate_sample(sample,model, wgt, iter, use_f128, max_eval)
276            // }
277            Integrand::ProcessIntegrand(integrand) => {
278                integrand.evaluate_sample(sample, model, wgt, iter, use_arb_prec, max_eval)
279            }
280        }
281    }
282
283    fn get_n_dim(&self) -> usize {
284        match self {
285            Integrand::UnitSurface(integrand) => integrand.get_n_dim(),
286            Integrand::UnitVolume(integrand) => integrand.get_n_dim(),
287            Integrand::HFunctionTest(integrand) => integrand.get_n_dim(),
288            // Integrand::ProcessIntegrandImpl(integrand) => integrand.get_n_dim(),
289            Integrand::ProcessIntegrand(integrand) => integrand.get_n_dim(),
290        }
291    }
292
293    fn get_integrator_settings(&self) -> IntegratorSettings {
294        match self {
295            Integrand::UnitSurface(integrand) => integrand.get_integrator_settings(),
296            Integrand::UnitVolume(integrand) => integrand.get_integrator_settings(),
297            Integrand::HFunctionTest(integrand) => integrand.get_integrator_settings(),
298            // Integrand::ProcessIntegrandImpl(integrand) => integrand.get_integrator_settings(),
299            Integrand::ProcessIntegrand(integrand) => integrand.get_integrator_settings(),
300        }
301    }
302}
303
304pub(crate) fn integrand_factory(settings: &RuntimeSettings) -> Integrand {
305    match settings.hard_coded_integrand.as_ref().unwrap().clone() {
306        IntegrandSettings::UnitSurface(integrand_settings) => Integrand::UnitSurface(
307            UnitSurfaceIntegrand::new(settings.clone(), integrand_settings),
308        ),
309        IntegrandSettings::UnitVolume(integrand_settings) => Integrand::UnitVolume(
310            UnitVolumeIntegrand::new(settings.clone(), integrand_settings),
311        ),
312        IntegrandSettings::HFunctionTest(integrand_settings) => Integrand::HFunctionTest(
313            HFunctionTestIntegrand::new(settings.clone(), integrand_settings),
314        ),
315    }
316}
317
318#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
319#[derive(Debug, Clone, Default, Serialize, Deserialize, Encode, Decode, PartialEq, JsonSchema)]
320// #[trait_decode(trait= GammaLoopContext)]
321pub struct UnitSurfaceSettings {
322    /// Number of independent three-momenta used to construct the unit-surface test dimension.
323    pub n_3d_momenta: usize,
324}
325
326#[derive(Clone)]
327pub struct UnitSurfaceIntegrand {
328    pub settings: RuntimeSettings,
329    pub n_dim: usize,
330    pub n_3d_momenta: usize,
331    pub surface: F<f64>,
332}
333
334#[allow(unused)]
335impl UnitSurfaceIntegrand {
336    pub(crate) fn new(
337        settings: RuntimeSettings,
338        integrand_settings: UnitSurfaceSettings,
339    ) -> UnitSurfaceIntegrand {
340        let n_dim = integrand_settings.n_3d_momenta * 3 - 1;
341        let surface = utils::compute_surface_and_volume(
342            integrand_settings.n_3d_momenta * 3 - 1,
343            F(settings.kinematics.e_cm),
344        )
345        .0;
346        UnitSurfaceIntegrand {
347            settings,
348            n_3d_momenta: integrand_settings.n_3d_momenta,
349            n_dim,
350            surface,
351        }
352    }
353
354    fn evaluate_numerator<T: FloatLike>(&self, loop_momenta: &[FourMomentum<F<T>>]) -> F<T> {
355        loop_momenta[0].temporal.value.one()
356    }
357
358    fn parameterize<T: FloatLike>(&self, xs: &[F<T>]) -> (Vec<[F<T>; 3]>, F<T>) {
359        let zero = xs[0].zero();
360        utils::global_parameterize(
361            xs,
362            F::<T>::from_f64(self.settings.kinematics.e_cm * self.settings.kinematics.e_cm),
363            &self
364                .settings
365                .sampling
366                .get_parameterization_settings()
367                .unwrap(),
368        )
369    }
370}
371
372#[allow(unused)]
373impl HasIntegrand for UnitSurfaceIntegrand {
374    fn name(&self) -> String {
375        "UnitSurfaceIntegrand".to_string()
376    }
377
378    fn create_grid(&self) -> Grid<F<f64>> {
379        Grid::Continuous(ContinuousGrid::new(
380            self.n_dim,
381            self.settings.integrator.n_bins,
382            self.settings.integrator.min_samples_for_update,
383            self.settings.integrator.bin_number_evolution.clone(),
384            self.settings.integrator.train_on_avg,
385        ))
386    }
387
388    fn get_n_dim(&self) -> usize {
389        self.n_dim
390    }
391
392    fn evaluate_sample(
393        &mut self,
394        sample: &Sample<F<f64>>,
395        model: &Model,
396        wgt: F<f64>,
397        iter: usize,
398        use_arb_prec: bool,
399        max_eval: Complex<F<f64>>,
400    ) -> Result<EvaluationResult> {
401        let start_evaluate_sample = std::time::Instant::now();
402
403        let xs = match sample {
404            Sample::Continuous(_w, v) => v,
405            _ => panic!("Wrong sample type"),
406        };
407        let mut sample_xs = vec![F(self.settings.kinematics.e_cm)];
408        sample_xs.extend(xs);
409
410        let before_parameterization = std::time::Instant::now();
411        let (moms, jac) = self.parameterize(sample_xs.as_slice());
412        let mut loop_momenta = vec![];
413        for m in &moms {
414            loop_momenta.push(FourMomentum::from_args(
415                ((m[0] + m[1] + m[2]) * (m[0] + m[1] + m[2])).sqrt(),
416                m[0],
417                m[1],
418                m[2],
419            ));
420        }
421
422        let parameterization_time = before_parameterization.elapsed();
423
424        let before_evaluation = std::time::Instant::now();
425        let mut itg_wgt = self.evaluate_numerator(loop_momenta.as_slice());
426        // Normalize the integral
427        itg_wgt /= self.surface;
428
429        info!("Sampled loop momenta:");
430        for (i, l) in loop_momenta.iter().enumerate() {
431            info!("k{} = ( {:-23})", i, format!("{:+.16e}", l),);
432        }
433        info!("Integrator weight : {:+.16e}", wgt);
434        info!("Integrand weight  : {:+.16e}", itg_wgt);
435        info!("Sampling jacobian : {:+.16e}", jac);
436        info!("Final contribution: {:+.16e}", itg_wgt * jac);
437
438        let is_nan = itg_wgt.is_nan();
439
440        let evaluation_time = before_evaluation.elapsed();
441
442        let evaluation_metadata = EvaluationMetaData {
443            total_timing: start_evaluate_sample.elapsed(),
444            integrand_evaluation_time: evaluation_time,
445            evaluator_evaluation_time: Duration::ZERO,
446            parameterization_time,
447            event_processing_time: Duration::ZERO,
448            generated_event_count: 0,
449            accepted_event_count: 0,
450            relative_instability_error: Complex::new_zero(),
451            is_nan,
452            loop_momenta_escalation: None,
453            stability_results: vec![StabilityResult {
454                precision: Precision::Double,
455                estimated_relative_accuracy: None,
456                status: StabilityStatus::Unknown,
457                total_time: start_evaluate_sample.elapsed(),
458            }],
459            threshold_counterterm_error: None,
460            radial_root_diagnostics: Default::default(),
461        };
462
463        Ok(EvaluationResult {
464            integrand_result: Complex::new(itg_wgt, F(0.)),
465            parameterization_jacobian: Some(jac),
466            integrator_weight: wgt,
467            event_groups: Default::default(),
468            evaluation_metadata,
469        })
470    }
471}
472
473#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
474#[derive(Debug, Clone, Default, Serialize, Deserialize, Encode, Decode, PartialEq, JsonSchema)]
475// #[trait_decode(trait= GammaLoopContext)]
476pub struct UnitVolumeSettings {
477    /// Number of independent three-momenta used to construct the unit-volume test dimension.
478    pub n_3d_momenta: usize,
479}
480
481#[derive(Clone)]
482pub struct UnitVolumeIntegrand {
483    pub settings: RuntimeSettings,
484    pub n_dim: usize,
485    pub n_3d_momenta: usize,
486    pub volume: F<f64>,
487}
488
489#[allow(unused)]
490impl UnitVolumeIntegrand {
491    pub(crate) fn new(
492        settings: RuntimeSettings,
493        integrand_settings: UnitVolumeSettings,
494    ) -> UnitVolumeIntegrand {
495        let n_dim = utils::get_n_dim_for_n_loop_momenta(
496            &settings.sampling,
497            integrand_settings.n_3d_momenta,
498            None,
499        );
500        let volume = utils::compute_surface_and_volume(
501            integrand_settings.n_3d_momenta * 3,
502            F(settings.kinematics.e_cm),
503        )
504        .1;
505        UnitVolumeIntegrand {
506            settings,
507            n_3d_momenta: integrand_settings.n_3d_momenta,
508            n_dim,
509            volume,
510        }
511    }
512
513    fn evaluate_numerator<T: FloatLike>(&self, loop_momenta: &[FourMomentum<F<T>>]) -> F<T> {
514        let zero = loop_momenta[0].temporal.value.zero();
515        if loop_momenta
516            .iter()
517            .map(|l| l.spatial.norm_squared())
518            .reduce(|acc, e| acc + &e)
519            .unwrap_or(zero.clone())
520            .sqrt()
521            > F::<T>::from_f64(self.settings.kinematics.e_cm)
522        {
523            zero
524        } else {
525            zero.one()
526        }
527    }
528
529    fn parameterize<T: FloatLike>(&self, xs: &[F<T>]) -> (Vec<[F<T>; 3]>, F<T>) {
530        let zero = xs[0].zero();
531        utils::global_parameterize(
532            xs,
533            F::<T>::from_f64(self.settings.kinematics.e_cm * self.settings.kinematics.e_cm),
534            &self
535                .settings
536                .sampling
537                .get_parameterization_settings()
538                .unwrap(),
539        )
540    }
541}
542
543#[allow(unused)]
544impl HasIntegrand for UnitVolumeIntegrand {
545    fn name(&self) -> String {
546        "UnitVolumeIntegrand".to_string()
547    }
548    fn create_grid(&self) -> Grid<F<f64>> {
549        Grid::Continuous(ContinuousGrid::new(
550            self.n_dim,
551            self.settings.integrator.n_bins,
552            self.settings.integrator.min_samples_for_update,
553            self.settings.integrator.bin_number_evolution.clone(),
554            self.settings.integrator.train_on_avg,
555        ))
556    }
557
558    fn get_n_dim(&self) -> usize {
559        self.n_dim
560    }
561
562    fn evaluate_sample(
563        &mut self,
564        sample: &Sample<F<f64>>,
565        model: &Model,
566        wgt: F<f64>,
567        iter: usize,
568        use_arb_prec: bool,
569        max_eval: Complex<F<f64>>,
570    ) -> Result<EvaluationResult> {
571        let start_evaluate_sample = std::time::Instant::now();
572
573        let xs = match sample {
574            Sample::Continuous(_w, v) => v,
575            _ => panic!("Wrong sample type"),
576        };
577
578        let before_parameterization = std::time::Instant::now();
579
580        let (moms, jac) = self.parameterize(xs);
581        let mut loop_momenta = vec![];
582        for m in &moms {
583            loop_momenta.push(FourMomentum::new(F(0.).into(), (*m).into()));
584        }
585
586        let parameterization_time = before_parameterization.elapsed();
587
588        let before_evaluation = std::time::Instant::now();
589        let mut itg_wgt = self.evaluate_numerator(loop_momenta.as_slice());
590        // Normalize the integral
591        itg_wgt /= self.volume;
592        info!("Sampled loop momenta:");
593        for (i, l) in loop_momenta.iter().enumerate() {
594            info!("k{} = ( {:-23})", i, format!("{:+.16e}", l),);
595        }
596        info!("Integrator weight : {:+.16e}", wgt);
597        info!("Integrand weight  : {:+.16e}", itg_wgt);
598        info!("Sampling jacobian : {:+.16e}", jac);
599        info!("Final contribution: {:+.16e}", itg_wgt * jac);
600
601        let is_nan = itg_wgt.is_nan();
602
603        let evaluation_time = before_evaluation.elapsed();
604
605        let evaluation_metadata = EvaluationMetaData {
606            total_timing: start_evaluate_sample.elapsed(),
607            integrand_evaluation_time: evaluation_time,
608            evaluator_evaluation_time: Duration::ZERO,
609            parameterization_time,
610            event_processing_time: Duration::ZERO,
611            generated_event_count: 0,
612            accepted_event_count: 0,
613            relative_instability_error: Complex::new_zero(),
614            is_nan,
615            loop_momenta_escalation: None,
616            stability_results: vec![StabilityResult {
617                precision: Precision::Double,
618                estimated_relative_accuracy: None,
619                status: StabilityStatus::Unknown,
620                total_time: start_evaluate_sample.elapsed(),
621            }],
622            threshold_counterterm_error: None,
623            radial_root_diagnostics: Default::default(),
624        };
625
626        Ok(EvaluationResult {
627            integrand_result: Complex::new(itg_wgt, F(0.)),
628            parameterization_jacobian: Some(jac),
629            integrator_weight: wgt,
630            event_groups: Default::default(),
631            evaluation_metadata,
632        })
633    }
634}
gammalooprs/integrands/mod.rs

gammalooprs/integrands/
mod.rs
