gammalooprs/graph/
hedge_data.rs

1use itertools::Itertools;
2use linnet::{
3    half_edge::involution::{Flow, Hedge},
4    parser::DotHedgeData,
5};
6use serde::{Deserialize, Serialize};
7use spenso::structure::{
8    OrderedStructure, PermutedStructure, TensorStructure,
9    representation::{LibraryRep, LibrarySlot},
10    slot::IsAbstractSlot,
11};
12use symbolica::atom::{Atom, FunctionBuilder};
13
14use crate::{graph::edge::PossibleParticle, numerator::aind::Aind};
15
16use super::{Autogen, parse::ParseGraph};
17
18use color_eyre::Result;
19
20#[derive(Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
21#[trait_decode(trait = symbolica::state::HasStateMap)]
22pub struct HedgeIndices {
23    pub vertex_indices: OrderedStructure<LibraryRep, Aind>,
24    pub edge_indices: OrderedStructure<LibraryRep, Aind>,
25}
26
27impl HedgeIndices {
28    pub(crate) fn new(edge_indices: OrderedStructure<LibraryRep, Aind>) -> Self {
29        Self {
30            vertex_indices: edge_indices.clone().dual(),
31            edge_indices,
32        }
33    }
34}
35
36#[derive(Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
37#[trait_decode(trait = symbolica::state::HasStateMap)]
38pub struct NumIndices {
39    pub color_indices: HedgeIndices,
40    pub spin_indices: HedgeIndices,
41}
42
43impl NumIndices {
44    fn from_particle(particle: &PossibleParticle, flow: Flow, h: Hedge) -> Self {
45        let creps = particle.color_reps(flow);
46        let sreps = particle.spin_reps();
47
48        let mut init = 0;
49        let mut last = None;
50        let color_structure: PermutedStructure<_> = creps
51            .external_reps_iter()
52            .map(|r| {
53                if let Some(l) = last {
54                    if l != r {
55                        last = Some(r);
56                        init = 0;
57                    } else {
58                        init += 1;
59                    }
60                } else {
61                    last = Some(r);
62                    init = 0;
63                }
64                r.slot(Aind::Hedge(h.0 as u16, init))
65            })
66            .collect();
67
68        let spin_structure: PermutedStructure<_> = sreps
69            .external_reps_iter()
70            .map(|r| {
71                if let Some(l) = last {
72                    if l != r {
73                        last = Some(r);
74                        init = 0;
75                    } else {
76                        init += 1;
77                    }
78                } else {
79                    last = Some(r);
80                    init = 0;
81                }
82                r.slot(Aind::Hedge(h.0 as u16, init))
83            })
84            .collect();
85
86        NumIndices {
87            color_indices: HedgeIndices::new(color_structure.structure),
88            spin_indices: HedgeIndices::new(spin_structure.structure),
89        }
90    }
91
92    pub(crate) fn parse<'a>(
93        graph: &'a ParseGraph,
94    ) -> impl FnMut(Hedge, &'a ParseHedgeData) -> Self {
95        |h, _| {
96            let eid = graph[&h];
97            let flow = graph.flow(h);
98
99            Self::from_particle(&graph[eid].particle, flow, h)
100        }
101    }
102}
103
104#[derive(Clone, bincode_trait_derive::Encode, bincode_trait_derive::Decode)]
105#[trait_decode(trait = symbolica::state::HasStateMap)]
106pub struct HedgeData {
107    pub num_indices: NumIndices,
108    pub ufo_order: Autogen<u8>,
109}
110
111impl HedgeData {
112    pub(crate) fn edge_color_slots<'a>(&'a self) -> impl Iterator<Item = LibrarySlot<Aind>> + 'a {
113        self.num_indices
114            .color_indices
115            .edge_indices
116            .external_structure_iter()
117    }
118
119    pub(crate) fn edge_spin_slots<'a>(&'a self) -> impl Iterator<Item = LibrarySlot<Aind>> + 'a {
120        self.num_indices
121            .spin_indices
122            .edge_indices
123            .external_structure_iter()
124    }
125
126    pub(crate) fn polarization(&self, mut builder: FunctionBuilder) -> Atom {
127        for s in self.edge_spin_slots() {
128            builder = builder.add_arg(s.to_atom())
129        }
130        builder.finish()
131    }
132
133    pub(crate) fn color_kronekers(&self, other: &Self) -> Atom {
134        let mut color = Atom::num(1);
135        for (i, j) in self.edge_color_slots().zip_eq(other.edge_color_slots()) {
136            if i.rep().matches(&j.rep()) {
137                color *= j.rep().id(i.aind, j.aind);
138            } else {
139                panic!("Should be the same rep found:{} and {}", i, j)
140            }
141        }
142
143        color
144    }
145}
146
147#[derive(Debug, Clone, Default, Serialize, Deserialize)]
148pub struct ParseHedgeData {
149    pub ufo_order: Option<u8>,
150}
151
152impl ParseHedgeData {
153    pub(crate) fn parse<'a>() -> impl FnMut((Hedge, &'a DotHedgeData)) -> Result<Self> {
154        |(_, h)| {
155            let Some(statement) = h
156                .statement
157                .as_deref()
158                .map(str::trim)
159                .filter(|s| !s.is_empty())
160            else {
161                return Ok(ParseHedgeData::default());
162            };
163
164            Ok(json5::from_str(statement)?)
165        }
166    }
167}
168
169impl From<&HedgeData> for DotHedgeData {
170    fn from(value: &HedgeData) -> Self {
171        let payload = ParseHedgeData {
172            ufo_order: (!value.ufo_order.autogenerated).then_some(value.ufo_order.value),
173        };
174
175        let statement = if payload.ufo_order.is_none() {
176            None
177        } else {
178            Some(json5::to_string(&payload).expect("serializing hedge payload should not fail"))
179        };
180
181        DotHedgeData::from(statement)
182    }
183}
184
185impl From<&ParseHedgeData> for DotHedgeData {
186    fn from(value: &ParseHedgeData) -> Self {
187        let statement = if value.ufo_order.is_none() {
188            None
189        } else {
190            Some(json5::to_string(value).expect("serializing hedge payload should not fail"))
191        };
192
193        DotHedgeData::from(statement)
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use linnet::{half_edge::involution::Hedge, parser::DotHedgeData};
200
201    use super::ParseHedgeData;
202
203    #[test]
204    fn parse_json5_hedge_payload() {
205        let dot = DotHedgeData::from(Some("{ufo_order: 2}".to_string()));
206        let parsed = ParseHedgeData::parse()((Hedge(0), &dot)).unwrap();
207
208        assert_eq!(parsed.ufo_order, Some(2));
209    }
210
211    #[test]
212    fn serialize_parse_hedge_data_to_json5_statement() {
213        let dot = DotHedgeData::from(Some("{ufo_order: 1}".to_string()));
214        let parsed = ParseHedgeData::parse()((Hedge(0), &dot)).unwrap();
215
216        let serialized: DotHedgeData = (&parsed).into();
217        let statement = serialized.statement.expect("statement should be present");
218
219        assert!(statement.contains("ufo_order"));
220    }
221}
gammalooprs/graph/hedge_data.rs

gammalooprs/graph/
hedge_data.rs
