algebra-sparse 0.4.0-beta.1

// Copyright (C) 2020-2025 algebra-sparse authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use na::{DMatrix, DMatrixView, RealField};

use crate::traits::IntoView;

/// A sparse matrix composed of block diagonal matrices.
///
/// This structure efficiently stores matrices that have a block diagonal structure,
/// where non-zero elements are only present in square blocks along the main diagonal.
/// This is commonly found in physical simulations, finite element methods, and
/// systems with localized interactions.
///
/// # Format
///
/// The matrix stores all block data in three arrays:
/// - `values`: All block elements stored contiguously in column-major order
/// - `block_row_offsets`: Starting row index for each block
/// - `block_element_offsets`: Starting index in `values` array for each block
///
/// # Examples
///
/// ```rust
/// use algebra_sparse::DiagonalBlockMatrix;
///
/// // Create a block diagonal matrix with blocks of size 2 and 3
/// // 2x2 block [1 2; 3 4] stored as [1, 3, 2, 4] in column-major
/// // 3x3 block [5 6 7; 8 9 10; 11 12 13] stored as [5, 8, 11, 6, 9, 12, 7, 10, 13]
/// let block_values = vec![1.0, 3.0, 2.0, 4.0,  // 2x2 block
///                        5.0, 8.0, 11.0, 6.0, 9.0, 12.0, 7.0, 10.0, 13.0]; // 3x3 block
/// let block_sizes = [2, 3];
/// let matrix = DiagonalBlockMatrix::from_block_values(block_values, &block_sizes);
///
/// println!("Matrix shape: {}x{}", matrix.nrows(), matrix.ncols());
/// ```
pub struct DiagonalBlockMatrix<T> {
    /// All values for matrix blocks stored in a single vector.
    values: Vec<T>,
    /// Row offset for each block in the matrix. Size = `num_blocks` + 1.
    /// One extra element makes it easier to calculate the size of the last block.
    block_row_offsets: Vec<usize>,
    /// The offsets of the first element in the `values` array for each block.
    block_element_offsets: Vec<usize>,
}

impl<T> DiagonalBlockMatrix<T> {
    /// Creates a block diagonal matrix from block values and sizes.
    ///
    /// # Arguments
    /// * `values` - All block elements stored contiguously in column-major order
    /// * `block_sizes` - Size of each square block along the diagonal
    ///
    /// # Examples
    ///
    /// ```rust
    /// use algebra_sparse::DiagonalBlockMatrix;
    ///
    /// // Create blocks: [1 2; 3 4] and [5 6 7; 8 9 10; 11 12 13]
    /// // Stored in column-major order: [1, 3, 2, 4] for first block
    /// let values = vec![1.0, 3.0, 2.0, 4.0, 5.0, 8.0, 11.0, 6.0, 9.0, 12.0, 7.0, 10.0, 13.0];
    /// let block_sizes = [2, 3]; // 2x2 block and 3x3 block
    /// let matrix = DiagonalBlockMatrix::from_block_values(values, &block_sizes);
    /// ```
    #[inline]
    pub fn from_block_values(values: Vec<T>, block_sizes: &[usize]) -> Self {
        let mut block_row_offsets = Vec::with_capacity(block_sizes.len() + 1);
        let mut block_element_offsets = Vec::with_capacity(block_sizes.len() + 1);
        let mut block_row_offset = 0;
        let mut block_element_offset = 0;

        for size in block_sizes.iter().copied() {
            block_row_offsets.push(block_row_offset);
            block_element_offsets.push(block_element_offset);
            block_row_offset += size;
            block_element_offset += size * size;
        }
        block_row_offsets.push(block_row_offset);
        block_element_offsets.push(block_element_offset);

        Self {
            values,
            block_row_offsets,
            block_element_offsets,
        }
    }

    #[inline]
    pub fn nrows(&self) -> usize {
        self.block_row_offsets.last().copied().unwrap_or(0)
    }

    #[inline]
    pub fn ncols(&self) -> usize {
        self.nrows()
    }
}

impl<'a, T> IntoView for &'a DiagonalBlockMatrix<T> {
    type View = DiagonalBlockMatrixView<'a, T>;

    #[inline]
    fn into_view(self) -> Self::View {
        DiagonalBlockMatrixView {
            values: &self.values,
            block_element_offsets: &self.block_element_offsets,
            block_row_offsets: &self.block_row_offsets,
        }
    }
}

/// An immutable view of a block diagonal matrix.
///
/// This provides efficient read-only access to block diagonal matrix data without allocation.
/// Views are commonly used for matrix operations and can be easily created from
/// block diagonal matrices.
///
/// # Examples
///
/// ```rust
/// use algebra_sparse::DiagonalBlockMatrix;
/// use algebra_sparse::traits::IntoView;
///
/// let matrix = DiagonalBlockMatrix::from_block_values(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
/// let view = matrix.into_view();
///
/// println!("Number of blocks: {}", view.num_blocks());
/// println!("Matrix shape: {}x{}", view.nrows(), view.ncols());
/// ```
#[derive(Clone, Copy)]
pub struct DiagonalBlockMatrixView<'a, T> {
    /// All block elements stored contiguously in column-major order.
    values: &'a [T],
    /// Row offset for each block in the matrix.
    block_row_offsets: &'a [usize],
    /// The offsets of the first element in the `values` array for each block.
    block_element_offsets: &'a [usize],
}

impl<'a, T> DiagonalBlockMatrixView<'a, T>
where
    T: RealField,
{
    /// Creates a block diagonal matrix view from raw parts without checking validity.
    ///
    /// # Safety
    /// This function does not validate that the provided parts form a valid block diagonal matrix.
    /// Invalid parts may cause undefined behavior when accessing the matrix.
    ///
    /// # Arguments
    /// * `values` - All block elements stored contiguously
    /// * `block_row_offsets` - Row offset array for each block
    /// * `block_element_offsets` - Element offset array for each block
    #[inline]
    pub fn from_parts_unchecked(
        values: &'a [T],
        block_row_offsets: &'a [usize],
        block_element_offsets: &'a [usize],
    ) -> Self {
        let slf = Self {
            values,
            block_row_offsets,
            block_element_offsets,
        };
        #[cfg(debug_assertions)]
        slf.assert_valid();
        slf
    }

    #[cfg(debug_assertions)]
    fn assert_valid(&self) {
        let mut expected_num_values = 0;
        let num_blocks = self.num_blocks();
        for block_index in 0..num_blocks {
            assert!(
                self.block_row_offsets[block_index] < self.block_row_offsets[block_index + 1],
                "Block sizes must be positive"
            );
            let block_size =
                self.block_row_offsets[block_index + 1] - self.block_row_offsets[block_index];

            let num_block_elements = self.block_element_offsets[block_index + 1]
                - self.block_element_offsets[block_index];
            assert!(
                num_block_elements == block_size * block_size,
                "Block element offsets do not match the expected size based on block sizes, expected {}, got {}, block_index = {}",
                block_size * block_size,
                num_block_elements,
                block_index
            );

            expected_num_values += block_size * block_size;
        }

        assert!(
            expected_num_values == self.values.len(),
            "Number of values does not match the expected number based on block sizes , expected {}, got {}",
            expected_num_values,
            self.values.len()
        );
    }

    #[inline]
    pub fn values(&self) -> &[T] {
        self.values
    }

    #[inline]
    pub fn num_blocks(&self) -> usize {
        self.block_row_offsets.len() - 1
    }

    #[inline]
    pub fn nrows(&self) -> usize {
        self.block_row_offsets.last().copied().unwrap_or(0)
    }

    #[inline]
    pub fn ncols(&self) -> usize {
        self.nrows()
    }

    /// Get the size of a block
    ///
    /// # Panics
    ///
    /// Panics if the `block_index` is out of bounds.
    #[inline]
    pub fn get_block_size(&self, block_index: usize) -> usize {
        debug_assert!(
            (block_index < self.num_blocks()),
            "Block index out of bounds"
        );
        self.block_row_offsets[block_index + 1] - self.block_row_offsets[block_index]
    }

    #[inline]
    pub fn get_block_row_start(&self, block_index: usize) -> usize {
        self.block_row_offsets[block_index]
    }

    /// Get the range of rows for the given block index.
    #[inline]
    pub fn get_block_row_range(&self, block_index: usize) -> std::ops::Range<usize> {
        debug_assert!(
            (block_index < self.num_blocks()),
            "Block index out of bounds"
        );
        let start = self.block_row_offsets[block_index];
        let end = self.block_row_offsets[block_index + 1];
        start..end
    }

    /// Get the block matrix at the given index as a view.
    ///
    /// # Panics
    ///
    /// Panics if the `block_index` is out of bounds.
    pub fn view_block(&self, block_index: usize) -> DMatrixView<T> {
        debug_assert!(
            (block_index < self.num_blocks()),
            "Block index out of bounds"
        );
        let start = self.block_element_offsets[block_index];
        let end = self.block_element_offsets[block_index + 1];
        let size = self.block_row_offsets[block_index + 1] - self.block_row_offsets[block_index];
        debug_assert!(size > 0);
        DMatrixView::from_slice(&self.values[start..end], size, size)
    }

    pub fn to_dense(&self) -> DMatrix<T> {
        let nrows = self.nrows();
        let ncols = self.ncols();
        let mut dense = DMatrix::zeros(nrows, ncols);
        for b_idx in 0..self.num_blocks() {
            let block = self.view_block(b_idx);
            let range = self.get_block_row_range(b_idx);
            let mut dst = dense.view_range_mut(range.clone(), range);
            dst.copy_from(&block);
        }
        dense
    }
}