ariadnetor-tensor 0.0.4

//! Convenience constructors and joined accessors for `BlockSparseTensorData<T, S>`.

use std::sync::Arc;

use aligned_vec::{AVec, ConstAlign};
use ariadnetor_core::backend::MemoryOrder;
use num_traits::Zero;
use rand::RngExt;

use super::{BlockCoord, BlockMeta, BlockSparseLayout, BlockSparseStorage, QNIndex};
use crate::{Sector, TensorData};

/// Backend-less BlockSparse tensor bundle =
/// `TensorData<BlockSparseStorage<T>, BlockSparseLayout<S>>`.
pub type BlockSparseTensorData<T, S> = TensorData<BlockSparseStorage<T>, BlockSparseLayout<S>>;

impl<T, S: Sector> BlockSparseTensorData<T, S> {
    /// Construct a zero-filled `BlockSparseTensorData` with all
    /// flux-allowed blocks.
    ///
    /// Cross-crate: the user-facing constructor is
    /// [`BlockSparseTensor::zeros`](crate::BlockSparseTensor::zeros),
    /// which pins memory order to the active backend. Direct callers
    /// (notably `ariadnetor-linalg` BSp kernel output-construction sites)
    /// that need an explicit `order` go through this helper or build
    /// `TensorData::new(storage, layout)` directly.
    pub fn zeros(indices: Vec<QNIndex<S>>, flux: S, order: MemoryOrder) -> Self
    where
        T: Clone + Zero,
    {
        let layout = BlockSparseLayout::new(indices, flux, order);
        let extent = <BlockSparseLayout<S> as crate::TensorLayout>::storage_extent(&layout);
        let mut data: AVec<T, ConstAlign<64>> = AVec::with_capacity(64, extent);
        data.resize(extent, T::zero());
        let storage = BlockSparseStorage::from_aligned(data);
        Self::new(storage, layout)
    }

    /// Construct by populating each flux-allowed block from a closure.
    ///
    /// The closure receives the block coordinate and its dense block
    /// shape (one entry per leg) and must return the block's flat data
    /// in the layout's memory `order`. Forbidden blocks are not
    /// queried. Block coordinates are visited in the layout's
    /// lexicographic enumeration order.
    ///
    /// Cross-crate: the user-facing constructor is
    /// [`BlockSparseTensor::from_block_fn`](crate::BlockSparseTensor::from_block_fn),
    /// which pins memory order to the active backend.
    ///
    /// # Panics
    ///
    /// Panics if the closure returns a `Vec<T>` whose length differs
    /// from `product(block_shape)` (the per-block element count).
    pub fn from_block_fn<F>(indices: Vec<QNIndex<S>>, flux: S, order: MemoryOrder, mut f: F) -> Self
    where
        T: Clone + Zero,
        F: FnMut(&BlockCoord, &[usize]) -> Vec<T>,
    {
        let layout = BlockSparseLayout::new(indices, flux, order);
        let extent = <BlockSparseLayout<S> as crate::TensorLayout>::storage_extent(&layout);
        let mut data: AVec<T, ConstAlign<64>> = AVec::with_capacity(64, extent);
        data.resize(extent, T::zero());
        for meta in layout.block_metas() {
            let block_shape = layout
                .block_shape(&meta.coord)
                .expect("BlockSparseLayout enumerated coord must resolve to a block shape");
            let block = f(&meta.coord, &block_shape);
            assert_eq!(
                block.len(),
                meta.size,
                "from_block_fn: closure returned {} elements for block {:?}, expected {}",
                block.len(),
                meta.coord,
                meta.size,
            );
            for (dst, src) in data[meta.offset..meta.offset + meta.size]
                .iter_mut()
                .zip(block)
            {
                *dst = src;
            }
        }
        let storage = BlockSparseStorage::from_aligned(data);
        Self::new(storage, layout)
    }

    /// Construct with all flux-allowed blocks filled with random
    /// values from the standard distribution.
    ///
    /// Cross-crate: the user-facing constructor is
    /// [`BlockSparseTensor::random`](crate::BlockSparseTensor::random),
    /// which pins memory order to the active backend.
    pub fn random<R: rand::Rng>(
        indices: Vec<QNIndex<S>>,
        flux: S,
        order: MemoryOrder,
        rng: &mut R,
    ) -> Self
    where
        rand::distr::StandardUniform: rand::distr::Distribution<T>,
    {
        let layout = BlockSparseLayout::new(indices, flux, order);
        let extent = <BlockSparseLayout<S> as crate::TensorLayout>::storage_extent(&layout);
        let mut data: AVec<T, ConstAlign<64>> = AVec::with_capacity(64, extent);
        for _ in 0..extent {
            data.push(rng.random());
        }
        let storage = BlockSparseStorage::from_aligned(data);
        Self::new(storage, layout)
    }

    /// Data slice for a block identified by coordinate.
    ///
    /// Returns `None` if the block is not stored (zero by symmetry).
    pub fn block_data(&self, coord: &BlockCoord) -> Option<&[T]> {
        let &idx = self.layout().block_index().get(coord)?;
        let meta = &self.layout().block_metas()[idx];
        Some(&self.storage().data()[meta.offset..meta.offset + meta.size])
    }

    /// Logical shape (total dimension per leg). Forwards to the layout.
    pub fn shape(&self) -> &[usize] {
        self.layout().shape()
    }

    /// Rank (number of legs). Forwards to the layout.
    pub fn rank(&self) -> usize {
        self.layout().rank()
    }

    /// Conserved flux (total quantum number). Forwards to the layout.
    pub fn flux(&self) -> &S {
        self.layout().flux()
    }

    /// Per-leg QN indices. Forwards to the layout.
    pub fn indices(&self) -> &[super::QNIndex<S>] {
        self.layout().indices()
    }

    /// Number of stored (non-zero) blocks. Forwards to the layout.
    pub fn num_blocks(&self) -> usize {
        self.layout().num_blocks()
    }

    /// Block metadata (sorted by coordinate). Forwards to the layout.
    pub fn block_metas(&self) -> &[BlockMeta] {
        self.layout().block_metas()
    }

    /// Check whether a block coordinate satisfies the flux conservation law.
    /// Forwards to the layout.
    pub fn is_allowed_block(&self, coord: &BlockCoord) -> bool {
        self.layout().is_allowed_block(coord)
    }

    /// Memory order the paired storage is laid out in. Forwards to the layout.
    pub fn order(&self) -> ariadnetor_core::backend::MemoryOrder {
        self.layout().order()
    }

    /// Mutable data slice for a block identified by coordinate
    /// (triggers CoW on the storage half if shared).
    pub fn block_data_mut(&mut self, coord: &BlockCoord) -> Option<&mut [T]>
    where
        T: Clone,
    {
        let &idx = self.layout().block_index().get(coord)?;
        let meta = &self.layout().block_metas()[idx];
        let offset = meta.offset;
        let size = meta.size;
        let arc = self.storage_mut().arc_mut();
        let data = Arc::make_mut(arc);
        Some(&mut data[offset..offset + size])
    }
}

impl<T, S: Sector> BlockSparseTensorData<T, S>
where
    T: ariadnetor_core::Scalar,
{
    /// Hermitian adjoint: element-wise conjugation of the data, flip
    /// of every QNIndex direction (Out↔In), and dualization of the
    /// flux.
    ///
    /// Block coordinates and packed offsets are preserved
    /// ([`BlockSparseLayout::dagger_layout`] reuses them). Involution:
    /// `x.dagger().dagger() == x`.
    pub fn dagger(&self) -> Self {
        let new_layout = self.layout().dagger_layout();
        let new_data: AVec<T, ConstAlign<64>> =
            AVec::from_iter(64, self.storage().data().iter().copied().map(|x| x.conj()));
        let storage = BlockSparseStorage::from_aligned(new_data);
        Self::new(storage, new_layout)
    }

    /// Element-wise complex conjugate. Layout (including directions
    /// and flux) is preserved; use [`dagger`](Self::dagger) when the
    /// adjoint structure is required for inner products.
    pub fn conj(&self) -> Self {
        let new_data: AVec<T, ConstAlign<64>> =
            AVec::from_iter(64, self.storage().data().iter().copied().map(|x| x.conj()));
        let storage = BlockSparseStorage::from_aligned(new_data);
        Self::new(storage, self.layout().clone())
    }

    /// Total number of stored elements across all blocks. Forwards
    /// to the storage half.
    pub fn stored_len(&self) -> usize {
        self.storage().stored_len()
    }

    /// Frobenius norm: √(Σ |element|²). Forwards to the storage half.
    pub fn norm_frobenius(&self) -> T::Real {
        self.storage().norm_frobenius()
    }

    /// Frobenius norm (alias for [`norm_frobenius`](Self::norm_frobenius)).
    pub fn norm(&self) -> T::Real {
        self.storage().norm()
    }

    /// Normalize to unit Frobenius norm in place. Returns the norm
    /// before normalization. Panics if the tensor has zero norm.
    pub fn normalize(&mut self) -> T::Real {
        self.storage_mut().normalize()
    }

    /// Normalize and return a new tensor (out-of-place). Returns
    /// `(normalized_tensor, original_norm)`. Panics if the tensor has
    /// zero norm.
    pub fn normalized(&self) -> (Self, T::Real) {
        let mut result = self.clone();
        let norm = result.normalize();
        (result, norm)
    }
}

impl<T, S: Sector> BlockSparseTensorData<T, S>
where
    T: Clone,
{
    /// Scale every stored element by a scalar factor in place
    /// (triggers CoW if shared).
    pub fn scale<F>(&mut self, factor: F)
    where
        T: std::ops::Mul<F, Output = T>,
        F: Clone,
    {
        self.storage_mut().scale(factor);
    }

    /// Scale every stored element and return a new tensor
    /// (out-of-place).
    pub fn scaled<F>(&self, factor: F) -> Self
    where
        T: std::ops::Mul<F, Output = T>,
        F: Clone,
    {
        let mut result = self.clone();
        result.scale(factor);
        result
    }
}