Skip to main content

gammalooprs/feyngen/
mod.rs

1pub mod diagram_generator;
2
3use ahash::{AHashMap, HashMap};
4use bincode_trait_derive::Decode;
5use bincode_trait_derive::Encode;
6use diagram_generator::EdgeColor;
7use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
8
9use rayon::iter::IntoParallelRefMutIterator;
10use rayon::iter::ParallelIterator;
11use schemars::JsonSchema;
12use serde::{Deserialize, Serialize};
13use smartstring::{LazyCompact, SmartString};
14use std::ops::RangeInclusive;
15use std::{fmt, str::FromStr};
16use symbolica::atom::Atom;
17use symbolica::graph::Graph as SymbolicaGraph;
18use thiserror::Error;
19
20use crate::{graph::LmbError, model::Model};
21
22#[derive(Error, Debug)]
23pub enum FeynGenError {
24    #[error("generation interrupted by user")]
25    Interrupted,
26    #[error("{0}")]
27    GenericError(String),
28    #[error("failed to build loop momentum basis for graph '{graph_name}'")]
29    LoopMomentumBasisError {
30        graph_name: String,
31        #[source]
32        source: LmbError,
33    },
34    #[error("Could not convert symbolica graph symmetry factor to an integer: {0}")]
35    SymmetryFactorError(String),
36    #[error("Could not numerically evaluate numerator: {0}")]
37    NumeratorEvaluationError(String),
38    #[error(transparent)]
39    Eyre(#[from] color_eyre::Report),
40}
41
42#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema, Encode, Decode)]
43
44pub struct GraphGroupingOptions {
45    pub numerical_sample_seed: u16,
46    pub number_of_numerical_samples: usize,
47    pub differentiate_particle_masses_only: bool,
48    pub fully_numerical_substitution_when_comparing_numerators: bool,
49    pub test_canonized_numerator: bool,
50    pub symmetric_polarizations: bool,
51}
52
53impl Default for GraphGroupingOptions {
54    fn default() -> Self {
55        Self {
56            numerical_sample_seed: 3,
57            number_of_numerical_samples: 5,
58            differentiate_particle_masses_only: true,
59            fully_numerical_substitution_when_comparing_numerators: false,
60            test_canonized_numerator: false,
61            symmetric_polarizations: false,
62        }
63    }
64}
65
66impl fmt::Display for GraphGroupingOptions {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        write!(
69            f,
70            "differentiate_masses_only={}, test_canonized_numerator={}, #samples={}, seed={}, fully_numerical_substitution={}, symmetric_polarizations={}",
71            self.numerical_sample_seed,
72            self.number_of_numerical_samples,
73            self.differentiate_particle_masses_only,
74            self.test_canonized_numerator,
75            self.fully_numerical_substitution_when_comparing_numerators,
76            self.symmetric_polarizations
77        )
78    }
79}
80
81#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema, Encode, Decode)]
82pub enum NumeratorAwareGraphGroupingOption {
83    NoGrouping,
84    OnlyDetectZeroes,
85    GroupIdenticalGraphUpToSign(GraphGroupingOptions),
86    GroupIdenticalGraphUpToScalarRescaling(GraphGroupingOptions),
87}
88
89impl fmt::Display for NumeratorAwareGraphGroupingOption {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        write!(
92            f,
93            "{}",
94            match self {
95                Self::NoGrouping => "no grouping",
96                Self::OnlyDetectZeroes => "only detect zero numerators",
97                Self::GroupIdenticalGraphUpToSign(_opts) => "up to a sign",
98                Self::GroupIdenticalGraphUpToScalarRescaling(_opts) => {
99                    "up to a scalar rescaling"
100                }
101            }
102        )
103    }
104}
105
106impl NumeratorAwareGraphGroupingOption {
107    pub(crate) fn get_options(&self) -> Option<&GraphGroupingOptions> {
108        match self {
109            Self::NoGrouping => None,
110            Self::OnlyDetectZeroes => None,
111            Self::GroupIdenticalGraphUpToSign(opts) => Some(opts),
112            Self::GroupIdenticalGraphUpToScalarRescaling(opts) => Some(opts),
113        }
114    }
115
116    #[allow(dead_code)]
117    pub(crate) fn description(&self) -> String {
118        format!(
119            "{}{}",
120            self,
121            self.get_options().map_or("".into(), |o| format!("({})", o))
122        )
123    }
124}
125
126#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema, Encode, Decode, Copy)]
127pub enum GenerationType {
128    Amplitude,
129    CrossSection,
130}
131
132impl fmt::Display for GenerationType {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        write!(
135            f,
136            "{}",
137            match self {
138                Self::Amplitude => "Amplitude",
139                Self::CrossSection => "Cross-section",
140            }
141        )
142    }
143}
144
145impl FromStr for GenerationType {
146    type Err = FeynGenError;
147
148    fn from_str(s: &str) -> Result<Self, FeynGenError> {
149        match s {
150            "amplitude" => Ok(Self::Amplitude),
151            "cross_section" => Ok(Self::CrossSection),
152            _ => Err(FeynGenError::GenericError(format!(
153                "Invalid generation type: {}",
154                s
155            ))),
156        }
157    }
158}
159
160pub(crate) fn get_coupling_orders<NodeColor: diagram_generator::NodeColorFunctions>(
161    graph: &SymbolicaGraph<NodeColor, EdgeColor>,
162    model: &Model,
163) -> AHashMap<SmartString<LazyCompact>, usize> {
164    let mut coupling_orders = AHashMap::default();
165    for node in graph.nodes() {
166        for (k, v) in node.data.coupling_orders(model) {
167            *coupling_orders.entry(k).or_insert(0) += v;
168        }
169    }
170    coupling_orders
171}
172
173#[derive(Debug, Clone, Encode, Decode, Serialize, Deserialize, JsonSchema, PartialEq)]
174pub struct FeynGenFilters(pub Vec<FeynGenFilter>);
175
176impl FeynGenFilters {
177    pub(crate) fn get_max_bridge(&self) -> Option<usize> {
178        self.0.iter().find_map(|f| match f {
179            FeynGenFilter::MaxNumberOfBridges(n) => Some(*n),
180            _ => None,
181        })
182    }
183
184    pub(crate) fn get_blob_range(&self) -> Option<&RangeInclusive<usize>> {
185        self.0.iter().find_map(|f| {
186            if let FeynGenFilter::BlobRange(v) = f {
187                Some(v)
188            } else {
189                None
190            }
191        })
192    }
193
194    pub(crate) fn get_spectator_range(&self) -> Option<&RangeInclusive<usize>> {
195        self.0.iter().find_map(|f| {
196            if let FeynGenFilter::SpectatorRange(v) = f {
197                Some(v)
198            } else {
199                None
200            }
201        })
202    }
203
204    #[allow(dead_code)]
205    pub fn allow_tadpoles(&self) -> bool {
206        !self
207            .0
208            .iter()
209            .any(|f| matches!(f, FeynGenFilter::TadpolesFilter(_)))
210    }
211
212    pub(crate) fn filter_cross_section_tadpoles(&self) -> bool {
213        self.0.iter().any(|f| {
214            matches!(
215                f,
216                FeynGenFilter::SewedFilter(SewedFilterOptions {
217                    filter_tadpoles: true,
218                    ..
219                })
220            )
221        })
222    }
223
224    pub(crate) fn get_particle_vetos(&self) -> Option<&[i64]> {
225        self.0.iter().find_map(|f| {
226            if let FeynGenFilter::ParticleVeto(v) = f {
227                Some(v.as_slice())
228            } else {
229                None
230            }
231        })
232    }
233
234    pub(crate) fn get_coupling_orders(&self) -> Option<&HashMap<String, (usize, Option<usize>)>> {
235        self.0.iter().find_map(|f| {
236            if let FeynGenFilter::CouplingOrders(o) = f {
237                Some(o)
238            } else {
239                None
240            }
241        })
242    }
243
244    pub(crate) fn get_perturbative_orders(&self) -> Option<&HashMap<String, usize>> {
245        self.0.iter().find_map(|f| {
246            if let FeynGenFilter::PerturbativeOrders(o) = f {
247                Some(o)
248            } else {
249                None
250            }
251        })
252    }
253
254    pub(crate) fn get_loop_count_range(&self) -> Option<(usize, usize)> {
255        self.0.iter().find_map(|f| {
256            if let FeynGenFilter::LoopCountRange(o) = f {
257                Some(*o)
258            } else {
259                None
260            }
261        })
262    }
263
264    pub(crate) fn get_fermion_loop_count_range(&self) -> Option<(usize, usize)> {
265        self.0.iter().find_map(|f: &FeynGenFilter| {
266            if let FeynGenFilter::FermionLoopCountRange(o) = f {
267                Some(*o)
268            } else {
269                None
270            }
271        })
272    }
273
274    #[allow(clippy::type_complexity)]
275    pub(crate) fn apply_filters<
276        NodeColor: diagram_generator::NodeColorFunctions + Send + Sync + Clone,
277    >(
278        &self,
279        graphs: &mut Vec<(SymbolicaGraph<NodeColor, EdgeColor>, Atom)>,
280        model: &Model,
281        pool: &rayon::ThreadPool,
282        progress_bar_style: &ProgressStyle,
283    ) -> Result<(), FeynGenError> {
284        for filter in self.0.iter() {
285            match filter {
286                FeynGenFilter::CouplingOrders(orders) => {
287                    let bar = ProgressBar::new(graphs.len() as u64);
288                    bar.set_style(progress_bar_style.clone());
289                    bar.set_message("Applying coupling orders constraints...");
290                    pool.install(|| {
291                        *graphs = graphs
292                            .par_iter_mut()
293                            .progress_with(bar.clone())
294                            .filter(|(g, _)| {
295                                let graph_coupling_orders = get_coupling_orders(g, model);
296
297                                // if a {
298                                //     info!(
299                                //         "Coupling orders constraints satisfied for graph {}",
300                                //         g.to_dot()
301                                //     );
302                                //     info!("{:?}", graph_coupling_orders);
303                                // }
304                                orders.iter().all(|(k, (v_min, v_max))| {
305                                    graph_coupling_orders
306                                        .get(&SmartString::from(k))
307                                        .map_or(0 == *v_min, |o| {
308                                            *o >= *v_min && (*v_max).is_none_or(|max| *o <= max)
309                                        })
310                                })
311                            })
312                            .map(|(g, sf)| (g.clone(), sf.clone()))
313                            .collect::<Vec<_>>()
314                    });
315                    bar.finish_and_clear();
316                }
317                FeynGenFilter::LoopCountRange((loop_count_min, loop_count_max)) => {
318                    graphs.retain(|(g, _)| {
319                        g.num_loops() >= *loop_count_min && g.num_loops() <= *loop_count_max
320                    });
321                }
322                FeynGenFilter::PerturbativeOrders(_)
323                | FeynGenFilter::MaxNumberOfBridges(_)
324                | FeynGenFilter::SelfEnergyFilter(_)
325                | FeynGenFilter::TadpolesFilter(_)
326                | FeynGenFilter::ZeroSnailsFilter(_)
327                | FeynGenFilter::FermionLoopCountRange(_)
328                | FeynGenFilter::SewedFilter(_)
329                | FeynGenFilter::FactorizedLoopTopologiesCountRange(_)
330                | FeynGenFilter::BlobRange(_)
331                | FeynGenFilter::SpectatorRange(_)
332                | FeynGenFilter::VertexAllow(_)
333                | FeynGenFilter::VertexVeto(_)
334                | FeynGenFilter::ParticleVeto(_) => {} // These other filters are implemented directly during diagram generation
335            }
336        }
337
338        Ok(())
339    }
340}
341
342#[derive(Debug, Clone, Copy, Encode, Decode, Serialize, Deserialize, JsonSchema, PartialEq)]
343pub struct SelfEnergyFilterOptions {
344    pub veto_self_energy_of_massive_lines: bool,
345    pub veto_self_energy_of_massless_lines: bool,
346    pub veto_only_scaleless_self_energy: bool,
347}
348
349impl Default for SelfEnergyFilterOptions {
350    fn default() -> Self {
351        Self {
352            veto_self_energy_of_massive_lines: true,
353            veto_self_energy_of_massless_lines: true,
354            veto_only_scaleless_self_energy: false,
355        }
356    }
357}
358
359impl fmt::Display for SelfEnergyFilterOptions {
360    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
361        let mut descr = vec![];
362        if self.veto_self_energy_of_massive_lines && !self.veto_self_energy_of_massless_lines {
363            descr.push("only massive legs")
364        } else if !self.veto_self_energy_of_massive_lines && self.veto_self_energy_of_massless_lines
365        {
366            descr.push("only massless legs")
367        };
368        if self.veto_only_scaleless_self_energy {
369            descr.push("only scaleless self-energies")
370        };
371        write!(f, "{}", descr.join(" | "))
372    }
373}
374
375#[derive(Debug, Clone, Copy, Encode, Decode, Serialize, Deserialize, JsonSchema, PartialEq)]
376pub struct SnailFilterOptions {
377    pub veto_snails_attached_to_massive_lines: bool,
378    pub veto_snails_attached_to_massless_lines: bool,
379    pub veto_only_scaleless_snails: bool,
380}
381
382impl Default for SnailFilterOptions {
383    fn default() -> Self {
384        Self {
385            veto_snails_attached_to_massive_lines: false,
386            veto_snails_attached_to_massless_lines: true,
387            veto_only_scaleless_snails: false,
388        }
389    }
390}
391
392impl fmt::Display for SnailFilterOptions {
393    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
394        let mut descr = vec![];
395        if self.veto_snails_attached_to_massive_lines
396            && !self.veto_snails_attached_to_massless_lines
397        {
398            descr.push("only attached to massive legs")
399        } else if !self.veto_snails_attached_to_massive_lines
400            && self.veto_snails_attached_to_massless_lines
401        {
402            descr.push("only attached to massless legs")
403        };
404        if self.veto_only_scaleless_snails {
405            descr.push("only scaleless snails")
406        };
407        write!(f, "{}", descr.join(" | "))
408    }
409}
410
411#[derive(Debug, Clone, Copy, Encode, Decode, Serialize, Deserialize, JsonSchema, PartialEq)]
412pub struct TadpolesFilterOptions {
413    pub veto_tadpoles_attached_to_massive_lines: bool,
414    pub veto_tadpoles_attached_to_massless_lines: bool,
415    pub veto_only_scaleless_tadpoles: bool,
416}
417
418impl Default for TadpolesFilterOptions {
419    fn default() -> Self {
420        Self {
421            veto_tadpoles_attached_to_massive_lines: true,
422            veto_tadpoles_attached_to_massless_lines: true,
423            veto_only_scaleless_tadpoles: false,
424        }
425    }
426}
427
428impl fmt::Display for TadpolesFilterOptions {
429    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
430        let mut descr = vec![];
431        if self.veto_tadpoles_attached_to_massive_lines
432            && !self.veto_tadpoles_attached_to_massless_lines
433        {
434            descr.push("only attached to massive legs")
435        } else if !self.veto_tadpoles_attached_to_massive_lines
436            && self.veto_tadpoles_attached_to_massless_lines
437        {
438            descr.push("only attached to massless legs")
439        };
440        if self.veto_only_scaleless_tadpoles {
441            descr.push("only scaleless tadpoles")
442        };
443        write!(f, "{}", descr.join(" | "))
444    }
445}
446
447#[derive(Debug, Clone, Copy, Encode, Decode, Serialize, Deserialize, JsonSchema, PartialEq)]
448pub struct SewedFilterOptions {
449    pub filter_tadpoles: bool,
450}
451
452#[derive(Debug, Clone, Encode, Decode, Serialize, Deserialize, JsonSchema, PartialEq)]
453pub enum FeynGenFilter {
454    SelfEnergyFilter(SelfEnergyFilterOptions),
455    TadpolesFilter(TadpolesFilterOptions),
456    ZeroSnailsFilter(SnailFilterOptions),
457    SewedFilter(SewedFilterOptions),
458    /// A list of vetoed pdgs
459    ParticleVeto(Vec<i64>),
460    VertexAllow(Vec<String>),
461    VertexVeto(Vec<String>),
462    MaxNumberOfBridges(usize),
463    /// A map between the coupling order name and a range of orders, inclusive, with an optional upper bound
464    CouplingOrders(HashMap<String, (usize, Option<usize>)>),
465    /// A range of loop counts, inclusive
466    LoopCountRange((usize, usize)),
467    /// A range of blob counts, inclusive
468    BlobRange(RangeInclusive<usize>),
469    SpectatorRange(RangeInclusive<usize>),
470    PerturbativeOrders(HashMap<String, usize>),
471    FermionLoopCountRange((usize, usize)),
472    FactorizedLoopTopologiesCountRange((usize, usize)),
473}
474
475impl fmt::Display for FeynGenFilter {
476    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
477        write!(
478            f,
479            "{}",
480            match self {
481                Self::SelfEnergyFilter(opts) => format!("NoExternalSelfEnergy({})", opts),
482                Self::ParticleVeto(pdgs) => format!(
483                    "ParticleVeto({})",
484                    pdgs.iter()
485                        .map(|x| x.to_string())
486                        .collect::<Vec<String>>()
487                        .join("|")
488                ),
489                Self::VertexVeto(vetos) => format!(
490                    "VertexVeto({})",
491                    vetos
492                        .iter()
493                        .map(|x| x.to_string())
494                        .collect::<Vec<String>>()
495                        .join("|")
496                ),
497                Self::VertexAllow(allowed) => format!(
498                    "VertexAllow({})",
499                    allowed
500                        .iter()
501                        .map(|x| x.to_string())
502                        .collect::<Vec<String>>()
503                        .join("|")
504                ),
505                Self::SpectatorRange(r) => format!("SpectatorRange({:?})", r),
506                Self::BlobRange(r) => format!("BlobRange({:?})", r),
507                Self::MaxNumberOfBridges(n) => format!("MaxNumberOfBridges({})", n),
508                Self::TadpolesFilter(opts) => format!("NoTadpoles({})", opts),
509                Self::ZeroSnailsFilter(opts) => format!("NoZeroSnails({})", opts),
510                Self::CouplingOrders(orders) => format!(
511                    "CouplingOrders({})",
512                    orders
513                        .iter()
514                        .map(|(k, (v_min, v_max_opt))| {
515                            if let Some(v_max) = v_max_opt {
516                                if v_min == v_max {
517                                    format!("{}=={}", k, v_min)
518                                } else {
519                                    format!("{}=[{}..{}]", k, v_min, v_max)
520                                }
521                            } else {
522                                format!("{}>={}", k, v_min)
523                            }
524                        })
525                        .collect::<Vec<String>>()
526                        .join("|")
527                ),
528                Self::PerturbativeOrders(orders) => format!(
529                    "PerturbativeOrders({})",
530                    orders
531                        .iter()
532                        .map(|(k, v)| format!("{}={}", k, v))
533                        .collect::<Vec<String>>()
534                        .join("|")
535                ),
536                Self::LoopCountRange((loop_count_min, loop_count_max)) =>
537                    format!("LoopCountRange({{{},{}}})", loop_count_min, loop_count_max),
538                Self::FermionLoopCountRange((loop_count_min, loop_count_max)) => format!(
539                    "FermionLoopCountRange({{{},{}}})",
540                    loop_count_min, loop_count_max
541                ),
542                Self::FactorizedLoopTopologiesCountRange((loop_count_min, loop_count_max)) =>
543                    format!(
544                        "NFactorizableLoopRange({{{},{}}})",
545                        loop_count_min, loop_count_max
546                    ),
547                Self::SewedFilter(SewedFilterOptions { filter_tadpoles }) => format!(
548                    "SewedCrossSectionFilter(filter_tadpoles={{{}}})",
549                    filter_tadpoles
550                ),
551            }
552        )
553    }
554}
555
556#[cfg(test)]
557pub mod test;