Skip to main content

ariadnetor_tensor/tensor/
block_sparse_ops.rs

1//! Pass-through inherent accessors on `BlockSparseTensor<T, S>`.
2//!
3//! The joined-form `Tensor` wraps a `TensorData<St, L>`; for block-sparse
4//! tensors the same per-block / per-leg introspection that
5//! `BlockSparseTensorData::block_data` and `BlockSparseLayout::indices`
6//! already expose is needed by downstream callers without forcing them to
7//! drill through `.data().layout().…`. The pass-throughs below thread the
8//! joined form's accessors back up to the `Tensor` surface so consumers
9//! can write `t.indices()`, `t.block_data(&coord)`, etc.
10
11use std::ops::Mul;
12
13use ariadnetor_core::Scalar;
14use num_traits::Float;
15
16use super::Tensor;
17use crate::{BlockCoord, BlockMeta, BlockSparseLayout, BlockSparseStorage, QNIndex, Sector};
18
19impl<T, S> Tensor<BlockSparseStorage<T>, BlockSparseLayout<S>>
20where
21    S: Sector,
22{
23    /// Memory order this tensor's flat block data is laid out in.
24    ///
25    /// Mirror of `DenseTensor::order` — saves block-sparse callers
26    /// from reaching through `.data().layout().order()` for a basic
27    /// layout property.
28    pub fn order(&self) -> ariadnetor_core::backend::MemoryOrder {
29        self.data.layout().order()
30    }
31
32    /// Per-leg `QNIndex` metadata (one entry per axis).
33    pub fn indices(&self) -> &[QNIndex<S>] {
34        self.data.layout().indices()
35    }
36
37    /// Overall flux label of the tensor.
38    pub fn flux(&self) -> &S {
39        self.data.layout().flux()
40    }
41
42    /// Number of flux-allowed blocks.
43    pub fn num_blocks(&self) -> usize {
44        self.data.layout().num_blocks()
45    }
46
47    /// Block metadata table (one entry per flux-allowed block).
48    pub fn block_metas(&self) -> &[BlockMeta] {
49        self.data.layout().block_metas()
50    }
51
52    /// Per-leg block shape for a stored block coordinate.
53    ///
54    /// Returns `None` if the coord is outside the layout's enumerated
55    /// block set.
56    pub fn block_shape(&self, coord: &BlockCoord) -> Option<Vec<usize>> {
57        self.data.layout().block_shape(coord)
58    }
59
60    /// Data slice for a block identified by coordinate.
61    ///
62    /// Returns `None` if the block is not stored (zero by symmetry).
63    pub fn block_data(&self, coord: &BlockCoord) -> Option<&[T]> {
64        self.data.block_data(coord)
65    }
66
67    /// Mutable data slice for a block identified by coordinate
68    /// (CoW-aware).
69    pub fn block_data_mut(&mut self, coord: &BlockCoord) -> Option<&mut [T]>
70    where
71        T: Clone,
72    {
73        self.data.block_data_mut(coord)
74    }
75}
76
77impl<T, S> Tensor<BlockSparseStorage<T>, BlockSparseLayout<S>>
78where
79    T: Scalar,
80    S: Sector,
81{
82    /// Frobenius norm: `sqrt(Σ |x|^2)` over the packed flat buffer.
83    pub fn norm(&self) -> T::Real {
84        let mut sq = <T::Real as num_traits::Zero>::zero();
85        for &x in self.data.storage().data() {
86            let a = x.abs();
87            sq = sq + a * a;
88        }
89        <T::Real as Float>::sqrt(sq)
90    }
91
92    /// Hermitian adjoint: element-wise conjugation, leg-direction flip,
93    /// and flux dualization.
94    pub fn dagger(&self) -> Self {
95        let td = self.data.dagger();
96        Self { data: td }
97    }
98
99    /// Element-wise complex conjugate.
100    pub fn conj(&self) -> Self {
101        let td = self.data.conj();
102        Self { data: td }
103    }
104}
105
106impl<T, S> Tensor<BlockSparseStorage<T>, BlockSparseLayout<S>>
107where
108    T: Clone,
109    S: Sector,
110{
111    /// Scale every stored element by a factor (in-place).
112    pub fn scale<F>(&mut self, factor: F)
113    where
114        T: Mul<F, Output = T>,
115        F: Clone,
116    {
117        self.data.scale(factor);
118    }
119
120    /// Scale every stored element by a factor (out-of-place).
121    pub fn scaled<F>(&self, factor: F) -> Self
122    where
123        T: Mul<F, Output = T>,
124        F: Clone,
125    {
126        Self {
127            data: self.data.scaled(factor),
128        }
129    }
130}