Skip to main content

gammalooprs/numerator/
aind.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use linnet::half_edge::{
4    NodeIndex,
5    involution::{EdgeIndex, Hedge},
6};
7use spenso::{
8    structure::slot::{AbsInd, DummyAind, ParseableAind, SlotError},
9    utils::{to_subscript, to_superscript},
10};
11use symbolica::{
12    atom::{Atom, AtomView, Symbol, representation::FunView},
13    coefficient::CoefficientView,
14    function, symbol,
15};
16use thiserror::Error;
17
18use crate::utils::GS;
19
20static DUMMYCOUNTER: AtomicUsize = AtomicUsize::new(0);
21
22/// A type that represents the name of an index in a tensor.
23#[derive(
24    Debug,
25    Copy,
26    Clone,
27    Ord,
28    PartialOrd,
29    Eq,
30    PartialEq,
31    Hash,
32    bincode_trait_derive::Encode,
33    bincode_trait_derive::Decode,
34    // bincode_trait_derive::BorrowDecodeFromDecode,
35)]
36#[trait_decode(trait = symbolica::state::HasStateMap)]
37pub enum Aind {
38    Normal(usize),
39    Hedge(u16, u16),
40    Edge(u16, u16),
41    Vertex(u16, u16),
42    UVTerm(u16, u16),
43    Symbol(Symbol),
44    Dummy(usize),
45}
46
47pub trait NewAind {
48    fn aind(self, local: u16) -> Aind;
49}
50
51impl NewAind for EdgeIndex {
52    fn aind(self, local: u16) -> Aind {
53        Aind::Edge(usize::from(self) as u16, local)
54    }
55}
56
57impl NewAind for NodeIndex {
58    fn aind(self, local: u16) -> Aind {
59        Aind::Vertex(usize::from(self) as u16, local)
60    }
61}
62impl NewAind for Hedge {
63    fn aind(self, local: u16) -> Aind {
64        Aind::Hedge(self.0 as u16, local)
65    }
66}
67
68impl AbsInd for Aind {}
69
70impl DummyAind for Aind {
71    fn new_dummy() -> Self {
72        let index = DUMMYCOUNTER.fetch_add(1, Ordering::Relaxed);
73        crate::debug_tags!(#generation, #inspect;
74            aind_dummy = true,
75            stage = "aind_new_dummy",
76            dummy_index = index,
77            "Aind dummy allocated"
78        );
79        Aind::Dummy(index)
80    }
81
82    fn is_dummy(&self) -> bool {
83        matches!(self, Aind::Dummy(_))
84    }
85
86    fn new_dummy_at(i: usize) -> Self {
87        Aind::Dummy(i)
88    }
89}
90
91impl ParseableAind for Aind {
92    type Error = AindError;
93
94    fn from_view(view: AtomView<'_>) -> Result<Self, Self::Error> {
95        view.try_into()
96    }
97
98    fn to_atom(&self) -> Atom {
99        (*self).into()
100    }
101}
102
103impl std::fmt::Display for Aind {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        match self {
106            Aind::Symbol(v) => {
107                if f.sign_minus() {
108                    write!(f, "_{}", v)
109                } else if f.sign_plus() {
110                    write!(f, "^{}", v)
111                } else {
112                    write!(f, "{}", v)
113                }
114            }
115            Aind::Normal(v) => {
116                if f.sign_minus() {
117                    write!(f, "{}", to_subscript(*v as isize))
118                } else if f.sign_plus() {
119                    write!(f, "{}", to_superscript(*v as isize))
120                } else {
121                    write!(f, "{}", v)
122                }
123            }
124            Aind::Hedge(i, j) => {
125                if f.sign_minus() {
126                    write!(
127                        f,
128                        "h{}.{}",
129                        to_subscript(*i as isize),
130                        to_subscript(*j as isize)
131                    )
132                } else if f.sign_plus() {
133                    write!(
134                        f,
135                        "h{}'{}",
136                        to_superscript(*i as isize),
137                        to_superscript(*j as isize)
138                    )
139                } else {
140                    write!(f, "h{}-{}", i, j)
141                }
142            }
143            Aind::Edge(i, j) => {
144                if f.sign_minus() {
145                    write!(
146                        f,
147                        "e{}.{}",
148                        to_subscript(*i as isize),
149                        to_subscript(*j as isize)
150                    )
151                } else if f.sign_plus() {
152                    write!(
153                        f,
154                        "e{}'{}",
155                        to_superscript(*i as isize),
156                        to_superscript(*j as isize)
157                    )
158                } else {
159                    write!(f, "e{}-{}", i, j)
160                }
161            }
162            Aind::Vertex(i, j) => {
163                if f.sign_minus() {
164                    write!(
165                        f,
166                        "v{}.{}",
167                        to_subscript(*i as isize),
168                        to_subscript(*j as isize)
169                    )
170                } else if f.sign_plus() {
171                    write!(
172                        f,
173                        "v{}'{}",
174                        to_superscript(*i as isize),
175                        to_superscript(*j as isize)
176                    )
177                } else {
178                    write!(f, "v{}-{}", i, j)
179                }
180            }
181
182            Aind::Dummy(v) => {
183                if f.sign_minus() {
184                    write!(f, "d{}", to_subscript(*v as isize))
185                } else if f.sign_plus() {
186                    write!(f, "d{}", to_superscript(*v as isize))
187                } else {
188                    write!(f, "d{}", v)
189                }
190            }
191            Aind::UVTerm(i, j) => {
192                if f.sign_minus() {
193                    write!(
194                        f,
195                        "u{}.{}",
196                        to_subscript(*i as isize),
197                        to_subscript(*j as isize)
198                    )
199                } else if f.sign_plus() {
200                    write!(
201                        f,
202                        "u{}'{}",
203                        to_superscript(*i as isize),
204                        to_superscript(*j as isize)
205                    )
206                } else {
207                    write!(f, "u{}-{}", i, j)
208                }
209            }
210        }
211    }
212}
213
214impl From<usize> for Aind {
215    fn from(value: usize) -> Self {
216        Aind::Normal(value)
217    }
218}
219
220impl From<Aind> for Atom {
221    fn from(value: Aind) -> Self {
222        match value {
223            Aind::Symbol(s) => Atom::var(s),
224            Aind::Dummy(i) => function!(GS.dummyaind, i as i64),
225            Aind::Normal(i) => Atom::num(i as i64),
226            Aind::UVTerm(i, j) => {
227                if j != 0 {
228                    function!(GS.uvaind, i as i64, j as i64)
229                } else {
230                    function!(GS.uvaind, i as i64)
231                }
232            }
233            Aind::Edge(i, j) => {
234                if j != 0 {
235                    function!(GS.edgeaind, i as i64, j as i64)
236                } else {
237                    function!(GS.edgeaind, i as i64)
238                }
239            }
240            Aind::Hedge(i, j) => {
241                if j != 0 {
242                    function!(GS.hedgeaind, i as i64, j as i64)
243                } else {
244                    function!(GS.hedgeaind, i as i64)
245                }
246            }
247            Aind::Vertex(i, j) => {
248                if j != 0 {
249                    function!(GS.vertexaind, i as i64, j as i64)
250                } else {
251                    function!(GS.vertexaind, i as i64)
252                }
253            }
254        }
255    }
256}
257impl<'a> From<Aind> for symbolica::atom::AtomOrView<'a> {
258    fn from(value: Aind) -> Self {
259        symbolica::atom::AtomOrView::Atom(Atom::from(value))
260    }
261}
262
263#[derive(Error, Debug)]
264pub enum AindError {
265    #[error("Argument is not a natural number")]
266    NotNatural,
267    #[error("Argument  {0} is not a valid index")]
268    NotIndex(String),
269    #[error("parsing error")]
270    ParsingError(String),
271}
272
273impl From<AindError> for SlotError {
274    fn from(value: AindError) -> Self {
275        SlotError::Any(value.into())
276    }
277}
278
279fn parse_natural_i64(arg: AtomView<'_>) -> Result<i64, AindError> {
280    let index =
281        i64::try_from(arg).map_err(|_| AindError::NotIndex(format!("Invalid index {arg}")))?;
282    if index >= 0 {
283        Ok(index)
284    } else {
285        Err(AindError::NotIndex(format!("Negative index {index}")))
286    }
287}
288
289fn parse_dummy_aind(f: FunView<'_>) -> Result<Aind, AindError> {
290    if f.get_nargs() != 1 {
291        return Err(AindError::ParsingError(format!(
292            "Incorrect number of arguments to dummy:{}",
293            f.as_view()
294        )));
295    }
296
297    let i = parse_natural_i64(f.iter().next().unwrap())?;
298    let i = usize::try_from(i).map_err(|e| AindError::NotIndex(e.to_string()))?;
299    DUMMYCOUNTER.fetch_max(i + 1, Ordering::Relaxed);
300    crate::debug_tags!(#generation, #inspect;
301        aind_dummy = true,
302        stage = "aind_parse_dummy",
303        dummy_index = i,
304        "Aind dummy parsed"
305    );
306    Ok(Aind::Dummy(i))
307}
308
309fn parse_u16_aind_pair(f: FunView<'_>, name: &str) -> Result<(u16, u16), AindError> {
310    if !(1..=2).contains(&f.get_nargs()) {
311        return Err(AindError::ParsingError(format!(
312            "Incorrect number of arguments to {name}:{}",
313            f.as_view()
314        )));
315    }
316
317    let mut iter = f.iter();
318    let i = parse_natural_i64(iter.next().unwrap())?;
319    let j = iter.next().map(parse_natural_i64).unwrap_or(Ok(0))?;
320
321    let i = u16::try_from(i).map_err(|e| AindError::NotIndex(e.to_string()))?;
322    let j = u16::try_from(j).map_err(|e| AindError::NotIndex(e.to_string()))?;
323
324    Ok((i, j))
325}
326
327impl TryFrom<AtomView<'_>> for Aind {
328    type Error = AindError;
329
330    fn try_from(view: AtomView<'_>) -> Result<Self, Self::Error> {
331        match view {
332            AtomView::Var(v) => Ok(Aind::Symbol(v.get_symbol())),
333            AtomView::Num(n) => match n.get_coeff_view() {
334                CoefficientView::Natural(n, 1, _, _) => Ok(Aind::Normal(
335                    usize::try_from(n).map_err(|e| AindError::NotIndex(e.to_string()))?,
336                )),
337                _ => Err(AindError::NotNatural),
338            },
339            AtomView::Fun(f) => {
340                let symbol = f.get_symbol();
341                if symbol == GS.dummyaind || symbol == symbol!("dummy") {
342                    parse_dummy_aind(f)
343                } else if symbol == GS.edgeaind || symbol == symbol!("edge") {
344                    let (i, j) = parse_u16_aind_pair(f, "edge")?;
345                    Ok(Aind::Edge(i, j))
346                } else if symbol == GS.hedgeaind || symbol == symbol!("hedge") {
347                    let (i, j) = parse_u16_aind_pair(f, "hedge")?;
348                    Ok(Aind::Hedge(i, j))
349                } else if symbol == GS.vertexaind {
350                    let (i, j) = parse_u16_aind_pair(f, "vertex")?;
351                    Ok(Aind::Vertex(i, j))
352                } else if symbol == GS.uvaind {
353                    let (i, j) = parse_u16_aind_pair(f, "uv")?;
354                    Ok(Aind::UVTerm(i, j))
355                } else {
356                    Err(AindError::NotIndex(format!(
357                        "Invalid index {}",
358                        f.as_view()
359                    )))
360                }
361            }
362
363            _ => Err(AindError::NotIndex(view.to_string())),
364        }
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use idenso::{
371        dirac::AGS,
372        shorthands::metric::MetricSimplifier,
373        tensor::{SymbolicNetParse, SymbolicTensor},
374    };
375    use spenso::{
376        network::parsing::{ParseSettings, ShadowedStructure, StructureFromAtom},
377        structure::{TensorStructure, permuted::Perm},
378    };
379    use symbolica::{
380        atom::{Atom, AtomCore},
381        function, parse_lit, symbol,
382    };
383
384    use crate::initialisation::initialise;
385
386    use super::*;
387
388    #[test]
389    fn parses_gammaloop_owned_aind_symbols() {
390        initialise().unwrap();
391
392        let cases = [
393            (Aind::Dummy(111), function!(GS.dummyaind, 111i64)),
394            (Aind::Edge(1, 0), function!(GS.edgeaind, 1i64)),
395            (Aind::Edge(1, 2), function!(GS.edgeaind, 1i64, 2i64)),
396            (Aind::Hedge(3, 0), function!(GS.hedgeaind, 3i64)),
397            (Aind::Hedge(3, 4), function!(GS.hedgeaind, 3i64, 4i64)),
398            (Aind::Vertex(5, 0), function!(GS.vertexaind, 5i64)),
399            (Aind::Vertex(5, 6), function!(GS.vertexaind, 5i64, 6i64)),
400            (Aind::UVTerm(7, 0), function!(GS.uvaind, 7i64)),
401            (Aind::UVTerm(7, 8), function!(GS.uvaind, 7i64, 8i64)),
402        ];
403
404        for (expected, atom) in cases {
405            assert_eq!(Aind::try_from(atom.as_view()).unwrap(), expected);
406        }
407    }
408
409    #[test]
410    fn parses_indexed_gammaloop_tensors_without_compact_rewrite() {
411        initialise().unwrap();
412
413        let bis0 = function!(symbol!("spenso::bis"), 4i64, Atom::from(Aind::Hedge(0, 0)));
414        let bis1 = function!(symbol!("spenso::bis"), 4i64, Atom::from(Aind::Hedge(1, 0)));
415        let mink = function!(symbol!("spenso::mink"), 4i64, Atom::from(Aind::Edge(2, 1)));
416
417        let expr = function!(GS.ubar, 1i64, bis0.clone())
418            * function!(GS.u, 1i64, bis1.clone())
419            * function!(GS.emr_mom, 2i64, mink.clone())
420            * function!(AGS.gamma, bis0, bis1, mink);
421
422        let net = expr
423            .parse_to_symbolic_net::<Aind>(&ParseSettings::default())
424            .unwrap();
425        let dangling = net.graph.dangling_indices();
426
427        assert!(
428            dangling.is_empty(),
429            "indexed GammaLoop tensors were parsed with dangling slots: {dangling:?}"
430        );
431    }
432
433    #[test]
434    fn test_structure_parsing() {
435        initialise().unwrap();
436        let expr = parse_lit!(gamma(
437            spenso::mink(4, edge(1, 1)),
438            spenso::mink(4, edge(1)),
439            spenso::mink(4, hedge(1, 1)),
440            spenso::mink(4, hedge(1)),
441            spenso::mink(4, vertex(2, 1)),
442            spenso::mink(4, vertex(2)),
443            spenso::mink(4, dummy(111)),
444            spenso::mink(4, 1)
445        ));
446        let structure = ShadowedStructure::<Aind>::parse(expr.as_view());
447
448        match structure {
449            Ok(s) => {
450                let pexpr = s
451                    .clone()
452                    .map_structure(|a| SymbolicTensor::from_named(&a).unwrap())
453                    .permute_inds()
454                    .expression
455                    .simplify_metrics();
456                assert_eq!(s.structure.order(), 8);
457                assert_eq!(
458                    expr,
459                    pexpr,
460                    "{}\n not equal to\n{}",
461                    expr.to_canonical_string(),
462                    pexpr.to_canonical_string()
463                );
464            }
465            Err(e) => println!("Error parsing structure: {}", e),
466        }
467        // assert_eq!(structure.to_string(), "s₁₂");
468    }
469}