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    /// An inclusive range of closed fermion loops, excluding ghost loops.
472    FermionLoopCountRange((usize, usize)),
473    FactorizedLoopTopologiesCountRange((usize, usize)),
474}
475
476impl fmt::Display for FeynGenFilter {
477    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
478        write!(
479            f,
480            "{}",
481            match self {
482                Self::SelfEnergyFilter(opts) => format!("NoExternalSelfEnergy({})", opts),
483                Self::ParticleVeto(pdgs) => format!(
484                    "ParticleVeto({})",
485                    pdgs.iter()
486                        .map(|x| x.to_string())
487                        .collect::<Vec<String>>()
488                        .join("|")
489                ),
490                Self::VertexVeto(vetos) => format!(
491                    "VertexVeto({})",
492                    vetos
493                        .iter()
494                        .map(|x| x.to_string())
495                        .collect::<Vec<String>>()
496                        .join("|")
497                ),
498                Self::VertexAllow(allowed) => format!(
499                    "VertexAllow({})",
500                    allowed
501                        .iter()
502                        .map(|x| x.to_string())
503                        .collect::<Vec<String>>()
504                        .join("|")
505                ),
506                Self::SpectatorRange(r) => format!("SpectatorRange({:?})", r),
507                Self::BlobRange(r) => format!("BlobRange({:?})", r),
508                Self::MaxNumberOfBridges(n) => format!("MaxNumberOfBridges({})", n),
509                Self::TadpolesFilter(opts) => format!("NoTadpoles({})", opts),
510                Self::ZeroSnailsFilter(opts) => format!("NoZeroSnails({})", opts),
511                Self::CouplingOrders(orders) => format!(
512                    "CouplingOrders({})",
513                    orders
514                        .iter()
515                        .map(|(k, (v_min, v_max_opt))| {
516                            if let Some(v_max) = v_max_opt {
517                                if v_min == v_max {
518                                    format!("{}=={}", k, v_min)
519                                } else {
520                                    format!("{}=[{}..{}]", k, v_min, v_max)
521                                }
522                            } else {
523                                format!("{}>={}", k, v_min)
524                            }
525                        })
526                        .collect::<Vec<String>>()
527                        .join("|")
528                ),
529                Self::PerturbativeOrders(orders) => format!(
530                    "PerturbativeOrders({})",
531                    orders
532                        .iter()
533                        .map(|(k, v)| format!("{}={}", k, v))
534                        .collect::<Vec<String>>()
535                        .join("|")
536                ),
537                Self::LoopCountRange((loop_count_min, loop_count_max)) =>
538                    format!("LoopCountRange({{{},{}}})", loop_count_min, loop_count_max),
539                Self::FermionLoopCountRange((loop_count_min, loop_count_max)) => format!(
540                    "FermionLoopCountRange({{{},{}}})",
541                    loop_count_min, loop_count_max
542                ),
543                Self::FactorizedLoopTopologiesCountRange((loop_count_min, loop_count_max)) =>
544                    format!(
545                        "NFactorizableLoopRange({{{},{}}})",
546                        loop_count_min, loop_count_max
547                    ),
548                Self::SewedFilter(SewedFilterOptions { filter_tadpoles }) => format!(
549                    "SewedCrossSectionFilter(filter_tadpoles={{{}}})",
550                    filter_tadpoles
551                ),
552            }
553        )
554    }
555}
556
557#[cfg(test)]
558pub mod test;