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}