Skip to main content

ocas_atom/tensor/
spec.rs

1//! Tensor symmetry specifications for canonicalisation.
2//!
3//! A [`TensorSpec`] declares which function heads represent tensors and how
4//! their slots behave under permutation (symmetric, antisymmetric subsets,
5//! cyclic).  Used by [`super::canon`] to encode tensor expressions into the
6//! graph-isomorphism engine.
7
8use std::collections::HashMap;
9
10use crate::Symbol;
11
12/// Slot symmetry for a tensor: a set of slot-index subsets that are
13/// symmetric, antisymmetric, or form a cycle.
14#[derive(Debug, Clone, Default)]
15pub struct SymmetrySpec {
16    /// Slot subsets whose members are interchangeable (symmetric).
17    pub symmetric_subsets: Vec<Vec<usize>>,
18    /// Slot subsets whose members are antisymmetric (swap flips sign).
19    pub antisymmetric_subsets: Vec<Vec<usize>>,
20    /// Cyclic permutation on a subset of slots.
21    pub cyclic: Option<Vec<usize>>,
22}
23
24impl SymmetrySpec {
25    /// No symmetry at all — every slot is independent.
26    pub fn none() -> Self {
27        Self {
28            symmetric_subsets: Vec::new(),
29            antisymmetric_subsets: Vec::new(),
30            cyclic: None,
31        }
32    }
33
34    /// All slots are fully symmetric.
35    pub fn fully_symmetric(rank: usize) -> Self {
36        Self {
37            symmetric_subsets: vec![(0..rank).collect()],
38            antisymmetric_subsets: Vec::new(),
39            cyclic: None,
40        }
41    }
42
43    /// All slots are fully antisymmetric.
44    pub fn fully_antisymmetric(rank: usize) -> Self {
45        Self {
46            symmetric_subsets: Vec::new(),
47            antisymmetric_subsets: vec![(0..rank).collect()],
48            cyclic: None,
49        }
50    }
51
52    /// Check whether this spec implies the given slot position should
53    /// have its location "hidden" in the graph encoding (i.e. not
54    /// participate in canonicalisation comparison).
55    pub fn is_slot_hidden(&self, pos: usize) -> bool {
56        for subset in &self.symmetric_subsets {
57            if subset.contains(&pos) {
58                return true;
59            }
60        }
61        self.cyclic.as_ref().is_some_and(|c| c.contains(&pos))
62    }
63}
64
65/// Complete tensor specification: which function heads are tensors and
66/// what symmetries their slots have.
67#[derive(Debug, Clone, Default)]
68pub struct TensorRegistry {
69    specs: HashMap<Symbol, SymmetrySpec>,
70    /// Index group assignment: symbol → group identifier.
71    /// Different groups prevent dummy index renaming across dimensions.
72    index_groups: HashMap<Symbol, u64>,
73}
74
75impl TensorRegistry {
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    /// Register a tensor with its slot symmetry spec.
81    pub fn register(&mut self, name: Symbol, spec: SymmetrySpec) {
82        self.specs.insert(name, spec);
83    }
84
85    /// Set the index group for a label (e.g. "mu" → 1 for spacetime,
86    /// "i" → 2 for internal).
87    pub fn set_index_group(&mut self, label: Symbol, group: u64) {
88        self.index_groups.insert(label, group);
89    }
90
91    /// Look up a tensor's symmetry spec.
92    pub fn spec(&self, name: Symbol) -> Option<&SymmetrySpec> {
93        self.specs.get(&name)
94    }
95
96    /// Look up an index label's group (0 = ungrouped/default).
97    pub fn index_group(&self, label: Symbol) -> u64 {
98        self.index_groups.get(&label).copied().unwrap_or(0)
99    }
100}