gammalooprs/graph/
edge.rs

1use std::{
2    fmt::Display,
3    sync::{Arc, Mutex},
4};
5
6use eyre::Context;
7use linnet::{
8    half_edge::{
9        HedgeGraph,
10        involution::{EdgeData, EdgeIndex, Flow, Hedge, HedgePair, Orientation},
11    },
12    parser::{DotEdgeData, DotHedgeData, DotVertexData},
13};
14use spenso::{
15    algebra::complex::Complex,
16    structure::{IndexLess, ScalarStructure, representation::LibraryRep},
17};
18use symbolica::{domains::float::Complex as SymComplex, prelude::*};
19
20use crate::{
21    feyngen::diagram_generator::EdgeColor,
22    integrands::process::ParamBuilder,
23    model::{ArcParticle, Model, UFOSymbol},
24    momentum::{Helicity, sample::LoopIndex},
25    numerator::aind::{Aind, NewAind},
26    utils::{F, FloatLike, GS},
27    uv::uv_graph::UVE,
28};
29
30use super::parse::{StripParse, ToQuoted};
31use crate::graph::Autogen;
32use color_eyre::Result;
33use eyre::eyre;
34
35#[derive(Debug, Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
36#[trait_decode(trait = crate::GammaLoopContext)]
37pub enum PossibleParticle {
38    Particle(ArcParticle),
39    MassOverriddenParticle { particle: ArcParticle, mass: Atom },
40    JustMass { expr: Atom },
41}
42
43impl From<ArcParticle> for PossibleParticle {
44    fn from(particle: ArcParticle) -> Self {
45        PossibleParticle::Particle(particle)
46    }
47}
48
49impl From<Atom> for PossibleParticle {
50    fn from(atom: Atom) -> Self {
51        PossibleParticle::JustMass { expr: atom }
52    }
53}
54
55impl From<()> for PossibleParticle {
56    fn from(_: ()) -> Self {
57        PossibleParticle::JustMass { expr: Atom::Zero }
58    }
59}
60
61impl PossibleParticle {
62    pub fn reverse(&self, model: &Model) -> Self {
63        match self {
64            PossibleParticle::Particle(p) => PossibleParticle::Particle(p.get_anti_particle(model)),
65            PossibleParticle::MassOverriddenParticle { particle, mass } => {
66                PossibleParticle::MassOverriddenParticle {
67                    particle: particle.get_anti_particle(model),
68                    mass: mass.clone(),
69                }
70            }
71            PossibleParticle::JustMass { expr } => {
72                PossibleParticle::JustMass { expr: expr.clone() }
73            }
74        }
75    }
76
77    pub fn orientation(&self) -> Orientation {
78        self.particle()
79            .map(|a| {
80                if a.is_fermion() {
81                    if a.pdg_code < 0 {
82                        Orientation::Reversed
83                    } else {
84                        Orientation::Default
85                    }
86                } else {
87                    Orientation::Undirected
88                }
89            })
90            .unwrap_or(Orientation::Undirected)
91    }
92
93    pub fn is_fermion(&self) -> bool {
94        match self {
95            PossibleParticle::Particle(p) => p.is_fermion(),
96            PossibleParticle::MassOverriddenParticle { particle, .. } => particle.is_fermion(),
97            PossibleParticle::JustMass { .. } => false,
98        }
99    }
100
101    pub fn is_anticommutating(&self) -> bool {
102        self.particle().is_some_and(|p| p.is_anticommutating())
103    }
104
105    pub fn is_self_antiparticle(&self) -> bool {
106        match self {
107            PossibleParticle::Particle(p) => p.is_self_antiparticle(),
108            PossibleParticle::MassOverriddenParticle { particle, .. } => {
109                particle.is_self_antiparticle()
110            }
111            PossibleParticle::JustMass { .. } => false,
112        }
113    }
114    pub fn mass_atom(&self) -> Atom {
115        match &self {
116            PossibleParticle::JustMass { expr, .. } => expr.clone(),
117            PossibleParticle::Particle(p) => p.mass.0.into(),
118            PossibleParticle::MassOverriddenParticle { mass, .. } => mass.clone(),
119        }
120    }
121
122    /// Literal ZERO, including constants canonicalized at model construction,
123    /// is absent. Adjustable named parameters remain present even at value zero.
124    pub(crate) fn chemical_potential_atom(&self) -> Option<Atom> {
125        self.particle()
126            .and_then(|particle| particle.chemical_potential)
127            .map(|mu| {
128                Atom::var(mu.0.0)
129                    .replace(UFOSymbol::zero().0)
130                    .with(Atom::Zero)
131            })
132            .filter(|mu| !mu.is_zero())
133    }
134    // pub fn just_mass(mass:Atom,)
135
136    pub fn zero() -> Self {
137        ().into()
138    }
139
140    pub(crate) fn override_mass(self, mass: Option<Atom>) -> Self {
141        let Some(mass) = mass else {
142            return self;
143        };
144
145        match self {
146            PossibleParticle::JustMass { .. } => PossibleParticle::JustMass { expr: mass },
147            PossibleParticle::MassOverriddenParticle { particle, .. }
148            | PossibleParticle::Particle(particle) => {
149                PossibleParticle::MassOverriddenParticle { particle, mass }
150            }
151        }
152    }
153
154    pub(crate) fn color_reps(&self, flow: Flow) -> IndexLess {
155        self.particle()
156            .map(|p| p.color_reps(flow))
157            .unwrap_or_else(IndexLess::scalar_structure)
158    }
159
160    pub(crate) fn spin_reps(&self) -> IndexLess<LibraryRep, Aind> {
161        self.particle()
162            .map(|p| p.spin_reps())
163            .unwrap_or_else(IndexLess::scalar_structure)
164    }
165
166    pub(crate) fn particle(&self) -> Option<ArcParticle> {
167        match self {
168            PossibleParticle::Particle(particle)
169            | PossibleParticle::MassOverriddenParticle { particle, .. } => Some(particle.clone()),
170            _ => None,
171        }
172    }
173
174    pub(crate) fn is_massless(&self) -> bool {
175        match self {
176            PossibleParticle::JustMass { expr } => expr.is_zero(),
177            PossibleParticle::Particle(p) => p.is_massless(),
178            PossibleParticle::MassOverriddenParticle { mass, .. } => mass.is_zero(),
179        }
180    }
181
182    pub(crate) fn is_massive(&self) -> bool {
183        !self.is_massless()
184    }
185}
186
187#[derive(Debug, Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
188#[trait_decode(trait = crate::GammaLoopContext)]
189#[derive(Default)]
190pub struct EdgeExtraData {
191    /// This controls the power of the momentum in the momtrop sampler. This does *not* affect other aspects of the treatment of the graph by gammaloop.
192    pub momtrop_edge_power: Option<Atom>,
193    /// This controls the power of the momentum in the vakint evaluation of the graph. This does *not* affect other aspects of the treatment of the graph by gammaloop.
194    pub vakint_edge_power: Option<isize>,
195}
196
197#[derive(Debug, Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
198#[trait_decode(trait = crate::GammaLoopContext)]
199pub struct Edge {
200    // #[bincode(with_serde)]
201    pub name: Autogen<String>,
202    pub extra_data: EdgeExtraData,
203    // pub edge_type: EdgeType,
204    // pub propagator: ArcPropagator,
205    pub particle: PossibleParticle,
206    pub mass: EdgeMass,
207    pub num: Autogen<Atom>,
208    // pub spin_num: Atom,
209    pub dod: Autogen<i32>,
210    pub is_dummy: bool, // #[bincode(with_serde)]
211}
212
213#[derive(Debug, Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
214#[trait_decode(trait = crate::GammaLoopContext)]
215pub enum EdgeMass {
216    Zero,
217    Value(Complex<F<f64>>),
218    ModelVar(Symbol),
219    Evaluator(Arc<Mutex<ExpressionEvaluator<Complex<F<f64>>>>>),
220}
221
222impl PartialEq for EdgeMass {
223    fn eq(&self, other: &Self) -> bool {
224        match (self, other) {
225            (EdgeMass::Zero, EdgeMass::Zero) => true,
226            (EdgeMass::Value(a), EdgeMass::Value(b)) => a == b,
227            (EdgeMass::ModelVar(a), EdgeMass::ModelVar(b)) => a == b,
228            _ => false,
229        }
230    }
231}
232
233impl EdgeMass {
234    pub fn from_atom(atom: Atom, model: &Model, paramb: &ParamBuilder) -> Result<Self> {
235        if atom.is_zero() {
236            return Ok(EdgeMass::Zero);
237        } else if let AtomView::Var(v) = atom.as_view() {
238            if model.contains_symbol(&UFOSymbol(v.get_symbol())) {
239                return Ok(EdgeMass::ModelVar(v.get_symbol()));
240            }
241        } else if let Ok(a) = SymComplex::<Float>::try_from(&atom) {
242            return Ok(EdgeMass::Value(Complex {
243                re: F(a.re.to_f64()),
244                im: F(a.im.to_f64()),
245            }));
246        }
247
248        let params: Vec<Atom> = (&paramb.pairs)
249            .into_iter()
250            .flat_map(|p| p.params.clone())
251            .collect();
252
253        let a = atom
254            .evaluator(&params)
255            .function_map(paramb.fn_map.clone())
256            .optimization_settings(OptimizationSettings::default())
257            .build()
258            .map_err(|a| eyre!(a))?;
259
260        Ok(EdgeMass::Evaluator(Arc::new(Mutex::new(a.map_coeff(
261            &|r| Complex::new(F::from(&r.re), F::from(&r.im)),
262        )))))
263    }
264
265    pub fn value<T: FloatLike>(
266        &self,
267        model: &Model,
268        paramb: &ParamBuilder,
269    ) -> Option<Complex<F<T>>> {
270        match self {
271            EdgeMass::Zero => None,
272            EdgeMass::Value(v) => Some(*v),
273            EdgeMass::ModelVar(s) => model.get_symbol_value(UFOSymbol(*s)),
274            EdgeMass::Evaluator(a) => Some(a.lock().unwrap().evaluate_single(&paramb.values[0])),
275        }
276        .map(|a| a.map_ref(|a| F::from_ff64(*a)))
277    }
278}
279
280impl UVE for Edge {
281    fn mass_atom(&self) -> Atom {
282        match &self.particle {
283            PossibleParticle::JustMass { expr, .. } => {
284                expr.replace(UFOSymbol::zero().0).with(Atom::Zero)
285            }
286            PossibleParticle::Particle(p) => Atom::var(p.mass.0.0)
287                .replace(UFOSymbol::zero().0)
288                .with(Atom::Zero),
289            PossibleParticle::MassOverriddenParticle { mass, .. } => {
290                mass.replace(UFOSymbol::zero().0).with(Atom::Zero)
291            }
292        }
293    }
294
295    fn particle_pdg_code(&self) -> Option<isize> {
296        self.particle().map(|particle| particle.pdg_code)
297    }
298
299    fn is_massive(&self) -> bool {
300        self.particle.is_massive()
301    }
302
303    fn is_fermion(&self) -> bool {
304        self.particle.is_fermion()
305    }
306
307    fn chemical_potential_atom(&self) -> Option<Atom> {
308        self.particle.chemical_potential_atom()
309    }
310}
311
312impl Edge {
313    pub fn random_helicity(&self, seed: u64) -> Helicity {
314        if let PossibleParticle::Particle(p) = &self.particle {
315            p.random_helicity(seed)
316        } else {
317            Helicity::ZERO
318        }
319    }
320
321    pub(crate) fn particle(&self) -> Option<ArcParticle> {
322        self.particle.particle()
323    }
324
325    pub(crate) fn mass_value<T: FloatLike>(
326        &self,
327        model: &Model,
328        paramb: &ParamBuilder,
329    ) -> Option<Complex<F<T>>> {
330        self.mass.value(model, paramb)
331    }
332}
333
334impl From<&ParseEdge> for DotEdgeData {
335    fn from(value: &ParseEdge) -> Self {
336        let mut e = DotEdgeData::empty();
337        if let Some(name) = &value.name {
338            e.add_statement("name", name.clone());
339        }
340        match &value.particle {
341            PossibleParticle::Particle(p) => {
342                e.add_statement("particle", format!("{}", p.name));
343            }
344            PossibleParticle::JustMass { expr, .. } => {
345                e.add_statement("mass", expr.to_quoted());
346            }
347            PossibleParticle::MassOverriddenParticle { mass, particle, .. } => {
348                e.add_statement("mass", mass.to_quoted());
349                e.add_statement("particle", format!("{}", particle.name));
350            }
351        }
352        if let Some(lmb_id) = &value.lmb_id {
353            e.add_statement("lmb_id", usize::from(*lmb_id));
354        }
355        if let Some(cut) = &value.is_cut {
356            e.add_statement("is_cut", usize::from(*cut));
357        }
358        if let Some(dod) = &value.dod {
359            e.add_statement("dod", *dod);
360        }
361        if let Some(num) = &value.num {
362            e.add_statement("num", num.to_quoted());
363        }
364
365        if let Some(mep) = value.momtrop_edge_power.as_ref() {
366            e.add_statement("momtrop_edge_power", mep.to_canonical_string());
367        }
368
369        if let Some(vak) = value.vakint_edge_power {
370            e.add_statement("vakint_edge_power", vak);
371        }
372
373        if value.is_dummy {
374            e.add_statement("is_dummy", value.is_dummy);
375        }
376        // e.add_statement("color_num", value.color_num.to_quoted());
377        e
378    }
379}
380
381impl From<&Edge> for DotEdgeData {
382    fn from(value: &Edge) -> Self {
383        value.to_dot_data(false)
384    }
385}
386
387impl From<&Edge> for ParseEdge {
388    fn from(value: &Edge) -> Self {
389        value.to_parse_edge(false)
390    }
391}
392
393impl Edge {
394    pub(crate) fn to_parse_edge(&self, include_autogenerated_fields: bool) -> ParseEdge {
395        ParseEdge {
396            name: self
397                .name
398                .clone()
399                .option_with_generated(include_autogenerated_fields),
400            particle: self.particle.clone(),
401            dod: self.dod.option_with_generated(include_autogenerated_fields),
402            is_dummy: self.is_dummy,
403            lmb_id: None,
404            num: self
405                .num
406                .clone()
407                .option_with_generated(include_autogenerated_fields),
408            is_cut: None,
409            momtrop_edge_power: self.extra_data.momtrop_edge_power.clone(),
410            vakint_edge_power: self.extra_data.vakint_edge_power,
411        }
412    }
413
414    pub(crate) fn to_dot_data(&self, include_autogenerated_fields: bool) -> DotEdgeData {
415        let parse_edge = self.to_parse_edge(include_autogenerated_fields);
416        let mut e: DotEdgeData = (&parse_edge).into();
417        if include_autogenerated_fields {
418            if self.name.autogenerated {
419                e.add_statement("name_autogen", true);
420            }
421            if self.dod.autogenerated {
422                e.add_statement("dod_autogen", true);
423            }
424            if self.num.autogenerated {
425                e.add_statement("num_autogen", true);
426            }
427        }
428        e
429    }
430}
431
432#[derive(Debug, Clone)]
433pub struct ParseEdge {
434    pub name: Option<String>,
435    pub particle: PossibleParticle,
436
437    pub dod: Option<i32>,
438    pub is_dummy: bool,
439    pub lmb_id: Option<LoopIndex>,
440
441    /// User provided numerator
442    pub num: Option<Atom>,
443    /// Specifies the incoming initial state hedge
444    pub is_cut: Option<Hedge>,
445    pub momtrop_edge_power: Option<Atom>,
446    pub vakint_edge_power: Option<isize>,
447}
448
449impl ParseEdge {
450    pub fn from_symbolica_edge(
451        model: &Model,
452        edge_color: &EdgeColor,
453        is_cut: Option<Hedge>,
454    ) -> Self {
455        let particle = model.get_particle_from_pdg(edge_color.pdg);
456        let mut e = ParseEdge::new(particle);
457        e.is_cut = is_cut;
458        e
459    }
460}
461
462impl Display for ParseEdge {
463    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
464        DotEdgeData::from(self).fmt(f)
465    }
466}
467
468impl UVE for ParseEdge {
469    fn mass_atom(&self) -> Atom {
470        match &self.particle {
471            PossibleParticle::JustMass { expr, .. } => expr.clone(),
472            PossibleParticle::Particle(p) => p.mass.0.into(),
473            PossibleParticle::MassOverriddenParticle { mass, .. } => mass.clone(),
474        }
475    }
476
477    fn particle_pdg_code(&self) -> Option<isize> {
478        self.particle.particle().map(|particle| particle.pdg_code)
479    }
480
481    fn is_massive(&self) -> bool {
482        self.particle.is_massive()
483    }
484
485    fn is_fermion(&self) -> bool {
486        self.particle.is_fermion()
487    }
488
489    fn chemical_potential_atom(&self) -> Option<Atom> {
490        self.particle.chemical_potential_atom()
491    }
492}
493
494impl ParseEdge {
495    pub fn new(particle: impl Into<PossibleParticle>) -> Self {
496        ParseEdge {
497            is_dummy: false,
498            name: None,
499            particle: particle.into(),
500            dod: None,
501            lmb_id: None,
502            num: None,
503            is_cut: None,
504            momtrop_edge_power: None,
505            vakint_edge_power: None,
506        }
507    }
508
509    pub fn is_dummy(mut self) -> Self {
510        self.is_dummy = true;
511        self
512    }
513
514    pub fn with_label(mut self, label: String) -> Self {
515        self.name = Some(label);
516        self
517    }
518
519    pub fn with_dod(mut self, dod: i32) -> Self {
520        self.dod = Some(dod);
521        self
522    }
523
524    pub fn with_lmb_id(mut self, lmb_id: LoopIndex) -> Self {
525        self.lmb_id = Some(lmb_id);
526        self
527    }
528
529    pub fn with_num(mut self, num: Atom) -> Self {
530        self.num = Some(num);
531        self
532    }
533}
534
535impl ParseEdge {
536    pub(crate) fn localize_ainds(
537        atom: impl AtomCore,
538        eid: EdgeIndex,
539        hedge_pair: HedgePair,
540    ) -> Atom {
541        let a = atom
542            .replace(GS.edgeid)
543            .with(Atom::num(usize::from(eid) as i64))
544            .replace_map(|term, _, out| {
545                if let AtomView::Fun(f) = term
546                    && f.get_symbol() == GS.edgeid
547                    && f.get_nargs() == 1
548                    && let Ok(i) = i64::try_from(f.iter().next().unwrap())
549                    && let Ok(u) = u16::try_from(i)
550                {
551                    **out = eid.aind(u).into();
552                }
553            });
554
555        match hedge_pair {
556            HedgePair::Paired { source, sink } | HedgePair::Split { source, sink, .. } => a
557                .replace(GS.sink_id)
558                .with(Atom::num(sink.0 as i64))
559                .replace_map(|term, _, out| {
560                    if let AtomView::Fun(f) = term
561                        && f.get_symbol() == GS.sink_id
562                        && f.get_nargs() == 1
563                        && let Ok(i) = i64::try_from(f.iter().next().unwrap())
564                        && let Ok(u) = u16::try_from(i)
565                    {
566                        **out = sink.aind(u).into();
567                    }
568                })
569                .replace(GS.source_id)
570                .with(Atom::num(source.0 as i64))
571                .replace_map(|term, _, out| {
572                    if let AtomView::Fun(f) = term
573                        && f.get_symbol() == GS.source_id
574                        && f.get_nargs() == 1
575                        && let Ok(i) = i64::try_from(f.iter().next().unwrap())
576                        && let Ok(u) = u16::try_from(i)
577                    {
578                        **out = source.aind(u).into();
579                    }
580                }),
581            HedgePair::Unpaired { hedge, flow } => match flow {
582                Flow::Source => a
583                    .replace(GS.source_id)
584                    .with(Atom::num(hedge.0 as i64))
585                    .replace_map(|term, _, out| {
586                        if let AtomView::Fun(f) = term
587                            && f.get_symbol() == GS.source_id
588                            && f.get_nargs() == 1
589                            && let Ok(i) = i64::try_from(f.iter().next().unwrap())
590                            && let Ok(u) = u16::try_from(i)
591                        {
592                            **out = hedge.aind(u).into();
593                        }
594                    }),
595                Flow::Sink => a
596                    .replace(GS.sink_id)
597                    .with(Atom::num(hedge.0 as i64))
598                    .replace_map(|term, _, out| {
599                        if let AtomView::Fun(f) = term
600                            && f.get_symbol() == GS.sink_id
601                            && f.get_nargs() == 1
602                            && let Ok(i) = i64::try_from(f.iter().next().unwrap())
603                            && let Ok(u) = u16::try_from(i)
604                        {
605                            **out = hedge.aind(u).into();
606                        }
607                    }),
608            },
609        }
610    }
611    #[allow(clippy::type_complexity)]
612    pub(crate) fn parse<'a>(
613        model: &'a Model,
614    ) -> impl FnMut(
615        &'a HedgeGraph<DotEdgeData, DotVertexData, DotHedgeData>,
616        EdgeIndex,
617        HedgePair,
618        EdgeData<&'a DotEdgeData>,
619    ) -> Result<EdgeData<Self>> {
620        |_: &'a HedgeGraph<DotEdgeData, DotVertexData, DotHedgeData>,
621         eid: EdgeIndex,
622         p: HedgePair,
623         e_data: EdgeData<&'a DotEdgeData>| {
624            let e = e_data.data;
625            let label = e.get::<_, String>("name").transpose()?;
626
627            let lmb_id: Option<LoopIndex> = e
628                .get::<_, usize>("lmb_id")
629                .transpose()?
630                .map(LoopIndex::from);
631
632            let is_cut: Option<Hedge> = e.get::<_, usize>("is_cut").transpose()?.map(Hedge::from);
633
634            let dod = e
635                .get::<_, String>("dod")
636                .transpose()
637                .with_context(|| "Error parsing dod".to_string())?
638                .map(|a| a.strip_parse())
639                .transpose()?;
640            let is_dummy = e.get::<_, bool>("is_dummy").transpose()?.unwrap_or(false);
641
642            let num = e
643                .get::<_, String>("num")
644                .transpose()?
645                .map(|a| -> Result<Atom> {
646                    Ok(Self::localize_ainds(a.strip_parse::<Atom>()?, eid, p))
647                })
648                .transpose()?;
649
650            let mass = e
651                .get::<_, String>("mass")
652                .transpose()?
653                .map(|a| a.strip_parse::<Atom>())
654                .transpose()?;
655
656            let momtrop_edge_power: Option<Atom> = e
657                .get::<_, String>("momtrop_edge_power")
658                .transpose()?
659                .map(|a| parse!(&a));
660            let vakint_edge_power: Option<isize> =
661                e.get::<_, isize>("vakint_edge_power").transpose()?;
662
663            let particle: PossibleParticle = if let Some(v) = e.get::<_, isize>("pdg") {
664                model.try_get_particle_from_pdg(v?)?.into()
665            } else if let Some(v) = e.get::<_, String>("particle") {
666                let pname: String = v?.strip_parse()?;
667                model.try_get_particle(pname)?.into()
668            } else {
669                ().into()
670            };
671
672            let orientation = if e.local_statements.contains_key("dir") {
673                e_data.orientation
674            } else {
675                particle.orientation()
676            };
677            Ok(EdgeData::new(
678                ParseEdge {
679                    is_dummy,
680                    dod,
681                    lmb_id,
682                    particle: particle.override_mass(mass),
683                    num,
684                    is_cut,
685                    name: label,
686                    momtrop_edge_power,
687                    vakint_edge_power,
688                },
689                orientation,
690            ))
691        }
692    }
693}
gammalooprs/graph/edge.rs

gammalooprs/graph/
edge.rs
