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

gammalooprs/graph/
edge.rs
