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#[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 )]
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 }
469}