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 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(_) => {} }
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 ParticleVeto(Vec<i64>),
460 VertexAllow(Vec<String>),
461 VertexVeto(Vec<String>),
462 MaxNumberOfBridges(usize),
463 CouplingOrders(HashMap<String, (usize, Option<usize>)>),
465 LoopCountRange((usize, usize)),
467 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;