ariadnetor_tensor/block_sparse/
tensor_data.rs

1//! Convenience constructors and joined accessors for `BlockSparseTensorData<T, S>`.
2
3use std::sync::Arc;
4
5use aligned_vec::{AVec, ConstAlign};
6use ariadnetor_core::backend::MemoryOrder;
7use num_traits::Zero;
8use rand::RngExt;
9
10use super::{BlockCoord, BlockMeta, BlockSparseLayout, BlockSparseStorage, QNIndex};
11use crate::{Sector, TensorData};
12
13/// Backend-less BlockSparse tensor bundle =
14/// `TensorData<BlockSparseStorage<T>, BlockSparseLayout<S>>`.
15pub type BlockSparseTensorData<T, S> = TensorData<BlockSparseStorage<T>, BlockSparseLayout<S>>;
16
17impl<T, S: Sector> BlockSparseTensorData<T, S> {
18    /// Construct a zero-filled `BlockSparseTensorData` with all
19    /// flux-allowed blocks.
20    ///
21    /// Cross-crate: the user-facing constructor is
22    /// [`BlockSparseTensor::zeros`](crate::BlockSparseTensor::zeros),
23    /// which pins memory order to the active backend. Direct callers
24    /// (notably `ariadnetor-linalg` BSp kernel output-construction sites)
25    /// that need an explicit `order` go through this helper or build
26    /// `TensorData::new(storage, layout)` directly.
27    pub fn zeros(indices: Vec<QNIndex<S>>, flux: S, order: MemoryOrder) -> Self
28    where
29        T: Clone + Zero,
30    {
31        let layout = BlockSparseLayout::new(indices, flux, order);
32        let extent = <BlockSparseLayout<S> as crate::TensorLayout>::storage_extent(&layout);
33        let mut data: AVec<T, ConstAlign<64>> = AVec::with_capacity(64, extent);
34        data.resize(extent, T::zero());
35        let storage = BlockSparseStorage::from_aligned(data);
36        Self::new(storage, layout)
37    }
38
39    /// Construct by populating each flux-allowed block from a closure.
40    ///
41    /// The closure receives the block coordinate and its dense block
42    /// shape (one entry per leg) and must return the block's flat data
43    /// in the layout's memory `order`. Forbidden blocks are not
44    /// queried. Block coordinates are visited in the layout's
45    /// lexicographic enumeration order.
46    ///
47    /// Cross-crate: the user-facing constructor is
48    /// [`BlockSparseTensor::from_block_fn`](crate::BlockSparseTensor::from_block_fn),
49    /// which pins memory order to the active backend.
50    ///
51    /// # Panics
52    ///
53    /// Panics if the closure returns a `Vec<T>` whose length differs
54    /// from `product(block_shape)` (the per-block element count).
55    pub fn from_block_fn<F>(indices: Vec<QNIndex<S>>, flux: S, order: MemoryOrder, mut f: F) -> Self
56    where
57        T: Clone + Zero,
58        F: FnMut(&BlockCoord, &[usize]) -> Vec<T>,
59    {
60        let layout = BlockSparseLayout::new(indices, flux, order);
61        let extent = <BlockSparseLayout<S> as crate::TensorLayout>::storage_extent(&layout);
62        let mut data: AVec<T, ConstAlign<64>> = AVec::with_capacity(64, extent);
63        data.resize(extent, T::zero());
64        for meta in layout.block_metas() {
65            let block_shape = layout
66                .block_shape(&meta.coord)
67                .expect("BlockSparseLayout enumerated coord must resolve to a block shape");
68            let block = f(&meta.coord, &block_shape);
69            assert_eq!(
70                block.len(),
71                meta.size,
72                "from_block_fn: closure returned {} elements for block {:?}, expected {}",
73                block.len(),
74                meta.coord,
75                meta.size,
76            );
77            for (dst, src) in data[meta.offset..meta.offset + meta.size]
78                .iter_mut()
79                .zip(block)
80            {
81                *dst = src;
82            }
83        }
84        let storage = BlockSparseStorage::from_aligned(data);
85        Self::new(storage, layout)
86    }
87
88    /// Construct with all flux-allowed blocks filled with random
89    /// values from the standard distribution.
90    ///
91    /// Cross-crate: the user-facing constructor is
92    /// [`BlockSparseTensor::random`](crate::BlockSparseTensor::random),
93    /// which pins memory order to the active backend.
94    pub fn random<R: rand::Rng>(
95        indices: Vec<QNIndex<S>>,
96        flux: S,
97        order: MemoryOrder,
98        rng: &mut R,
99    ) -> Self
100    where
101        rand::distr::StandardUniform: rand::distr::Distribution<T>,
102    {
103        let layout = BlockSparseLayout::new(indices, flux, order);
104        let extent = <BlockSparseLayout<S> as crate::TensorLayout>::storage_extent(&layout);
105        let mut data: AVec<T, ConstAlign<64>> = AVec::with_capacity(64, extent);
106        for _ in 0..extent {
107            data.push(rng.random());
108        }
109        let storage = BlockSparseStorage::from_aligned(data);
110        Self::new(storage, layout)
111    }
112
113    /// Data slice for a block identified by coordinate.
114    ///
115    /// Returns `None` if the block is not stored (zero by symmetry).
116    pub fn block_data(&self, coord: &BlockCoord) -> Option<&[T]> {
117        let &idx = self.layout().block_index().get(coord)?;
118        let meta = &self.layout().block_metas()[idx];
119        Some(&self.storage().data()[meta.offset..meta.offset + meta.size])
120    }
121
122    /// Logical shape (total dimension per leg). Forwards to the layout.
123    pub fn shape(&self) -> &[usize] {
124        self.layout().shape()
125    }
126
127    /// Rank (number of legs). Forwards to the layout.
128    pub fn rank(&self) -> usize {
129        self.layout().rank()
130    }
131
132    /// Conserved flux (total quantum number). Forwards to the layout.
133    pub fn flux(&self) -> &S {
134        self.layout().flux()
135    }
136
137    /// Per-leg QN indices. Forwards to the layout.
138    pub fn indices(&self) -> &[super::QNIndex<S>] {
139        self.layout().indices()
140    }
141
142    /// Number of stored (non-zero) blocks. Forwards to the layout.
143    pub fn num_blocks(&self) -> usize {
144        self.layout().num_blocks()
145    }
146
147    /// Block metadata (sorted by coordinate). Forwards to the layout.
148    pub fn block_metas(&self) -> &[BlockMeta] {
149        self.layout().block_metas()
150    }
151
152    /// Check whether a block coordinate satisfies the flux conservation law.
153    /// Forwards to the layout.
154    pub fn is_allowed_block(&self, coord: &BlockCoord) -> bool {
155        self.layout().is_allowed_block(coord)
156    }
157
158    /// Memory order the paired storage is laid out in. Forwards to the layout.
159    pub fn order(&self) -> ariadnetor_core::backend::MemoryOrder {
160        self.layout().order()
161    }
162
163    /// Mutable data slice for a block identified by coordinate
164    /// (triggers CoW on the storage half if shared).
165    pub fn block_data_mut(&mut self, coord: &BlockCoord) -> Option<&mut [T]>
166    where
167        T: Clone,
168    {
169        let &idx = self.layout().block_index().get(coord)?;
170        let meta = &self.layout().block_metas()[idx];
171        let offset = meta.offset;
172        let size = meta.size;
173        let arc = self.storage_mut().arc_mut();
174        let data = Arc::make_mut(arc);
175        Some(&mut data[offset..offset + size])
176    }
177}
178
179impl<T, S: Sector> BlockSparseTensorData<T, S>
180where
181    T: ariadnetor_core::Scalar,
182{
183    /// Hermitian adjoint: element-wise conjugation of the data, flip
184    /// of every QNIndex direction (Out↔In), and dualization of the
185    /// flux.
186    ///
187    /// Block coordinates and packed offsets are preserved
188    /// ([`BlockSparseLayout::dagger_layout`] reuses them). Involution:
189    /// `x.dagger().dagger() == x`.
190    pub fn dagger(&self) -> Self {
191        let new_layout = self.layout().dagger_layout();
192        let new_data: AVec<T, ConstAlign<64>> =
193            AVec::from_iter(64, self.storage().data().iter().copied().map(|x| x.conj()));
194        let storage = BlockSparseStorage::from_aligned(new_data);
195        Self::new(storage, new_layout)
196    }
197
198    /// Element-wise complex conjugate. Layout (including directions
199    /// and flux) is preserved; use [`dagger`](Self::dagger) when the
200    /// adjoint structure is required for inner products.
201    pub fn conj(&self) -> Self {
202        let new_data: AVec<T, ConstAlign<64>> =
203            AVec::from_iter(64, self.storage().data().iter().copied().map(|x| x.conj()));
204        let storage = BlockSparseStorage::from_aligned(new_data);
205        Self::new(storage, self.layout().clone())
206    }
207
208    /// Total number of stored elements across all blocks. Forwards
209    /// to the storage half.
210    pub fn stored_len(&self) -> usize {
211        self.storage().stored_len()
212    }
213
214    /// Frobenius norm: √(Σ |element|²). Forwards to the storage half.
215    pub fn norm_frobenius(&self) -> T::Real {
216        self.storage().norm_frobenius()
217    }
218
219    /// Frobenius norm (alias for [`norm_frobenius`](Self::norm_frobenius)).
220    pub fn norm(&self) -> T::Real {
221        self.storage().norm()
222    }
223
224    /// Normalize to unit Frobenius norm in place. Returns the norm
225    /// before normalization. Panics if the tensor has zero norm.
226    pub fn normalize(&mut self) -> T::Real {
227        self.storage_mut().normalize()
228    }
229
230    /// Normalize and return a new tensor (out-of-place). Returns
231    /// `(normalized_tensor, original_norm)`. Panics if the tensor has
232    /// zero norm.
233    pub fn normalized(&self) -> (Self, T::Real) {
234        let mut result = self.clone();
235        let norm = result.normalize();
236        (result, norm)
237    }
238}
239
240impl<T, S: Sector> BlockSparseTensorData<T, S>
241where
242    T: Clone,
243{
244    /// Scale every stored element by a scalar factor in place
245    /// (triggers CoW if shared).
246    pub fn scale<F>(&mut self, factor: F)
247    where
248        T: std::ops::Mul<F, Output = T>,
249        F: Clone,
250    {
251        self.storage_mut().scale(factor);
252    }
253
254    /// Scale every stored element and return a new tensor
255    /// (out-of-place).
256    pub fn scaled<F>(&self, factor: F) -> Self
257    where
258        T: std::ops::Mul<F, Output = T>,
259        F: Clone,
260    {
261        let mut result = self.clone();
262        result.scale(factor);
263        result
264    }
265}
ariadnetor_tensor/block_sparse/tensor_data.rs

ariadnetor_tensor/block_sparse/
tensor_data.rs
