Skip to main content

gammalooprs/graph/
vertex.rs

1use itertools::Itertools;
2use linnet::{
3    half_edge::{HedgeGraph, NodeIndex, nodestore::BitVecNeighborIter},
4    parser::DotVertexData,
5};
6
7use symbolica::atom::Atom;
8
9use crate::{
10    GammaLoopContext,
11    feyngen::diagram_generator::NodeColorWithVertexRule,
12    model::{ArcParticle, ArcVertexRule, Model},
13};
14use color_eyre::Result;
15use eyre::{Context, eyre};
16
17use super::{
18    Autogen,
19    edge::ParseEdge,
20    global::ParseData,
21    hedge_data::ParseHedgeData,
22    parse::{StripParse, ToQuoted, extract_oriented_particles_from_vertex_hedges},
23};
24
25#[derive(Debug, Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
26#[trait_decode(trait = GammaLoopContext)]
27pub struct Vertex {
28    // #[bincode(with_serde)]
29    pub name: Autogen<String>,
30    pub vertex_rule: Option<ArcVertexRule>,
31    pub num: Autogen<Atom>,
32    // pub num_spin: ParamTensor<OrderedStructure<Euclidean, Aind>>,
33    // pub num_color: ParamTensor<OrderedStructure<Euclidean, Aind>>,
34    pub dod: Autogen<i32>,
35}
36
37impl Vertex {
38    pub(crate) fn get_num(&self) -> Atom {
39        self.num.value.clone() //.kill_color()
40    }
41
42    pub(crate) fn to_dot_data(&self, include_autogenerated_fields: bool) -> DotVertexData {
43        let mut v = DotVertexData::empty();
44        if include_autogenerated_fields || !self.name.autogenerated {
45            v.name = Some(self.name.value.clone());
46        }
47        if let Some(vertex_rule) = &self.vertex_rule {
48            v.add_statement("int_id", vertex_rule.name.as_str());
49        }
50        if include_autogenerated_fields || !self.dod.autogenerated {
51            v.add_statement("dod", self.dod.value);
52        }
53        if include_autogenerated_fields || !self.num.autogenerated {
54            v.add_statement("num", self.num.value.to_quoted());
55        }
56        if include_autogenerated_fields {
57            if self.name.autogenerated {
58                v.add_statement("name_autogen", true);
59            }
60            if self.dod.autogenerated {
61                v.add_statement("dod_autogen", true);
62            }
63            if self.num.autogenerated {
64                v.add_statement("num_autogen", true);
65            }
66        }
67        v
68    }
69}
70
71impl From<&Vertex> for DotVertexData {
72    fn from(value: &Vertex) -> Self {
73        value.to_dot_data(false)
74    }
75}
76impl Vertex {}
77
78#[derive(Debug, Clone)]
79pub struct ParseVertex {
80    pub name: Option<String>,
81    // pub strict: bool,
82    pub vertex_rule: Option<ArcVertexRule>,
83    pub num: Option<Atom>,
84    pub dod: Option<i32>,
85}
86
87impl From<&ParseVertex> for DotVertexData {
88    fn from(value: &ParseVertex) -> Self {
89        let mut v = DotVertexData::empty();
90        if let Some(name) = &value.name {
91            v.name = Some(name.clone());
92        }
93        if let Some(vertex_rule) = &value.vertex_rule {
94            v.add_statement("int_id", vertex_rule.name.as_str());
95        }
96        if let Some(dod) = value.dod {
97            v.add_statement("dod", dod);
98        }
99        if let Some(num) = &value.num {
100            v.add_statement("num", num.to_quoted());
101        }
102        v
103    }
104}
105
106impl ParseVertex {
107    pub fn with_num(mut self, num: Atom) -> Self {
108        self.num = Some(num);
109        self
110    }
111
112    pub fn with_label(mut self, label: String) -> Self {
113        self.name = Some(label);
114        self
115    }
116}
117
118impl From<ArcVertexRule> for ParseVertex {
119    fn from(vertex_rule: ArcVertexRule) -> Self {
120        ParseVertex {
121            name: None,
122            // strict: false,
123            vertex_rule: Some(vertex_rule),
124            dod: None,
125            num: None,
126        }
127    }
128}
129
130pub trait ParticleEdge {
131    fn particle(&self) -> Option<ArcParticle>;
132    fn is_dummy(&self) -> bool;
133}
134
135impl From<&NodeColorWithVertexRule> for ParseVertex {
136    fn from(value: &NodeColorWithVertexRule) -> Self {
137        value.vertex_rule.clone().into()
138    }
139}
140impl ParseVertex {
141    pub(crate) fn parse<'a>(
142        model: &'a Model,
143        parse_data: &'a ParseData,
144    ) -> impl FnMut(
145        &'a HedgeGraph<ParseEdge, &'a DotVertexData, ParseHedgeData>,
146        BitVecNeighborIter<'a>,
147        &'a &'a DotVertexData,
148    ) -> Result<Self> {
149        move |g, n, v| {
150            let name = v.name().map(|id| id.to_string());
151
152            let dod = v
153                .get::<_, String>("dod")
154                .transpose()
155                .with_context(|| "Error parsing vertex dod".to_string())?
156                .map(|a| a.strip_parse())
157                .transpose()?;
158
159            if let Some(num) = v.get::<_, String>("num") {
160                let num = num?;
161                Ok(ParseVertex {
162                    dod,
163                    name,
164                    vertex_rule: None,
165                    num: Some(num.strip_parse().with_context(|| {
166                        format!(
167                            "Error parsing vertex num {num} of graph {}",
168                            parse_data.name
169                        )
170                    })?),
171                })
172            } else if let Some(n) = v.get::<_, String>("int_id") {
173                let vertex_rule = Some(model.get_vertex_rule(n.unwrap()));
174
175                Ok(ParseVertex {
176                    dod,
177                    name,
178                    vertex_rule,
179                    num: None,
180                })
181            } else {
182                let node_id = n
183                    .clone()
184                    .next()
185                    .map(|h| g.node_id(h))
186                    .unwrap_or(NodeIndex(0));
187                let mut particles = extract_oriented_particles_from_vertex_hedges(g, n, model);
188                particles.sort();
189
190                let res = model.particle_set_to_vertex_rules_map.get(&particles);
191                if let Some(res) = res {
192                    if res.len() == 1 {
193                        Ok(ParseVertex {
194                            dod,
195                            name,
196                            vertex_rule: Some(res[0].clone()),
197                            num: None,
198                        })
199                    } else {
200                        Err(eyre!("Multiple vertex rules for {:?}", particles))
201                    }
202                } else {
203                    let particles = particles.iter().map(|p| p.name.as_str()).collect_vec();
204                    Err(eyre!(
205                        "Failed to find vertex rule for particles: {:?} for node {node_id} in graph {}",
206                        particles,
207                        parse_data.name,
208                    ))
209                }
210            }
211        }
212    }
213}