algebra_sparse/

set.rs

// 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 crate::traits::IntoView;
use crate::{CsVecBuilder, CsrMatrixView, Real};

/// A collection of CSR matrices stored efficiently in a single data structure.
///
/// This structure is designed to store multiple CSR matrices with different dimensions
/// in a compact format. It's particularly useful for applications that need to manage
/// many sparse matrices with similar sparsity patterns or for systems that generate
/// multiple matrices during computation.
///
/// # Format
///
/// The set stores all matrices in contiguous arrays:
/// - `col_indices`: Column indices for all matrices concatenated
/// - `values`: Non-zero values for all matrices concatenated
/// - `row_offsets`: Row offsets for all matrices concatenated
/// - `ncols`: Number of columns for each individual matrix
/// - `partition`: Metadata to locate each matrix within the concatenated data
///
/// # Examples
///
/// ```rust
/// use algebra_sparse::CsrMatrixSet;
/// use algebra_sparse::traits::IntoView;
/// use algebra_sparse::CsrMatrixSetMethods;
///
/// let mut set: CsrMatrixSet<f64> = CsrMatrixSet::default();
///
/// // Add first matrix to the set
/// {
///     let mut builder1 = set.new_matrix(3, 1e-10);
///     builder1.new_row().push(0, 1.0);
/// } // builder1 is dropped here and matrix is added to set
///
/// // Add second matrix to the set
/// {
///     let mut builder2 = set.new_matrix(2, 1e-10);
///     builder2.new_row().push(1, 3.0);
/// } // builder2 is dropped here and matrix is added to set
///
/// println!("Set contains {} matrices", (&set).len());
/// ```
#[derive(Clone)]
pub struct CsrMatrixSet<T> {
    /// The column indices of the non-zero entries for all matrices.
    col_indices: Vec<usize>,
    /// The non-zero values for all matrices.
    values: Vec<T>,
    /// The offsets of each row in the `col_indices` and `values` arrays for all matrices.
    row_offsets: Vec<usize>,
    /// The number of columns for each matrix in the set.
    ncols: Vec<usize>,
    /// Partition information to locate individual matrices within the concatenated data.
    partition: Vec<Partition>,
}

impl<T> Default for CsrMatrixSet<T> {
    /// Creates a new empty CSR matrix set.
    #[inline]
    fn default() -> Self {
        Self {
            col_indices: Vec::new(),
            values: Vec::new(),
            row_offsets: vec![0],
            ncols: Vec::new(),
            partition: Vec::new(),
        }
    }
}

#[derive(Clone)]
struct Partition {
    pub value_offset: usize,
    pub value_len: usize,
    pub row_offset: usize,
    pub row_len: usize,
}

impl Partition {
    #[inline]
    pub fn value_range(&self) -> std::ops::Range<usize> {
        self.value_offset..self.value_offset + self.value_len
    }

    #[inline]
    pub fn row_offset_range(&self) -> std::ops::Range<usize> {
        self.row_offset..self.row_offset + self.row_len
    }
}

impl<T: Real> CsrMatrixSet<T> {
    /// Clears all matrices from the set.
    ///
    /// This removes all data and allows reuse of the set.
    pub fn clear(&mut self) {
        self.col_indices.clear();
        self.values.clear();
        self.row_offsets.clear();
        self.ncols.clear();
        self.partition.clear();
    }

    /// Creates a new matrix builder for this set.
    ///
    /// The matrix will be automatically added to the set when the builder is dropped.
    ///
    /// # Arguments
    /// * `ncol` - Number of columns for the new matrix
    /// * `zero_threshold` - Values below this threshold are filtered out
    ///
    /// # Returns
    /// A `CsrMatrixBuilder` for constructing the new matrix
    pub fn new_matrix(&mut self, ncol: usize, zero_threshold: T) -> CsrMatrixBuilder<T> {
        let value_start = self.values.len();
        let row_start = self.row_offsets.len();
        self.row_offsets.push(0);
        CsrMatrixBuilder {
            set: self,
            zero_threshold,
            value_start,
            row_start,
            ncol,
        }
    }
}

impl<T> CsrMatrixSet<T> {
    /// Returns the view of the matrix at the given index.
    ///
    /// # Arguments
    /// * `index` - Index of the matrix to retrieve
    ///
    /// # Returns
    /// A `CsrMatrixView` representing the requested matrix
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds
    #[inline]
    pub fn get(&self, index: usize) -> CsrMatrixView<T> {
        let partition = &self.partition[index];
        CsrMatrixView::from_parts_unchecked(
            &self.row_offsets[partition.row_offset_range()],
            &self.col_indices[partition.value_range()],
            &self.values[partition.value_range()],
            self.ncols[index],
        )
    }

    /// Returns a view of the entire matrix set.
    #[inline]
    pub fn as_view(&self) -> CsrMatrixSetView<'_, T> {
        CsrMatrixSetView {
            col_indices: &self.col_indices,
            values: &self.values,
            row_offsets: &self.row_offsets,
            ncols: &self.ncols,
            partition: &self.partition,
        }
    }
}

/// An immutable view of a CSR matrix set for efficient read-only access.
///
/// This structure provides zero-cost abstraction access to multiple CSR matrices stored
/// in a compact, consolidated format. It's designed for scenarios where you need to
/// read or process multiple sparse matrices without the overhead of copying data
/// or creating separate matrix objects.
///
/// # Use Cases
///
/// Views are particularly useful for:
///
/// ## Parallel Processing
/// ```rust
/// use algebra_sparse::CsrMatrixSet;
///
/// let mut set = CsrMatrixSet::default();
/// // Add some matrices to the set
/// {
///     let mut builder = set.new_matrix(3, 1e-10);
///     builder.new_row().push(0, 1.0);
/// }
/// {
///     let mut builder = set.new_matrix(2, 1e-10);
///     builder.new_row().push(1, 2.0);
/// }
///
/// let view = set.as_view();
/// assert_eq!(view.len(),2);
/// let (left, right) = view.split_at(view.len() / 2);
/// assert_eq!(left.len(),1);
/// assert_eq!(right.len(),1);
/// ```
#[derive(Clone, Copy)]
pub struct CsrMatrixSetView<'a, T> {
    col_indices: &'a [usize],
    values: &'a [T],
    row_offsets: &'a [usize],
    ncols: &'a [usize],
    partition: &'a [Partition],
}

impl<'a, T> CsrMatrixSetView<'a, T> {
    /// Returns the view of the matrix at the given index.
    ///
    /// # Arguments
    /// * `index` - Index of the matrix to retrieve
    ///
    /// # Returns
    /// A `CsrMatrixView` representing the requested matrix
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds
    #[inline]
    pub fn get(self, index: usize) -> CsrMatrixView<'a, T> {
        let partition = &self.partition[index];
        CsrMatrixView::from_parts_unchecked(
            &self.row_offsets[partition.row_offset_range()],
            &self.col_indices[partition.value_range()],
            &self.values[partition.value_range()],
            self.ncols[index],
        )
    }

    /// Returns the number of matrices in the set.
    #[inline]
    pub fn len(&self) -> usize {
        self.partition.len()
    }

    /// Returns true if the set contains no matrices.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.partition.is_empty()
    }

    /// Splits the matrix set view into two at the given index.
    ///
    /// This is a zero-cost operation that creates two independent views that reference
    /// the same underlying data but represent disjoint subsets of the matrices.
    /// The operation is O(1) and involves no copying or allocation of matrix data.
    ///
    /// # Arguments
    /// * `index` - The split position. The left view will contain matrices at indices `[0, index)`,
    ///   and the right view will contain matrices at indices `[index, len)`.
    ///
    /// # Returns
    /// A tuple of two views: `(left, right)` where:
    /// - `left` contains matrices `0..index`
    /// - `right` contains matrices `index..len`
    ///
    /// # Panics
    ///
    /// Panics if `index > len()`. Splitting at `index = 0` or `index = len()` is allowed
    /// and will return an empty view on one side.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use algebra_sparse::CsrMatrixSet;
    ///
    /// let mut set = CsrMatrixSet::default();
    /// // Add 3 matrices to the set
    /// {
    ///     let mut builder = set.new_matrix(2, 1e-10);
    ///     builder.new_row().push(0, 1.0);
    /// }
    /// {
    ///     let mut builder = set.new_matrix(2, 1e-10);
    ///     builder.new_row().push(1, 2.0);
    /// }
    /// {
    ///     let mut builder = set.new_matrix(2, 1e-10);
    ///     builder.new_row().push(0, 3.0);
    /// }
    ///
    /// let view = set.as_view();
    ///
    /// // Split in the middle
    /// let (left, right) = view.split_at(1);
    /// assert_eq!(left.len(), 1);  // matrices 0
    /// assert_eq!(right.len(), 2); // matrices 1, 2
    ///
    /// // Split at the beginning
    /// let (empty, all) = view.split_at(0);
    /// assert!(empty.is_empty());
    /// assert_eq!(all.len(), 3);
    ///
    /// // Split at the end
    /// let (all, empty) = view.split_at(3);
    /// assert_eq!(all.len(), 3);
    /// assert!(empty.is_empty());
    /// ```
    ///
    /// # Parallel Processing
    ///
    /// This method is particularly useful for parallel processing:
    ///
    /// ```rust
    /// # use algebra_sparse::CsrMatrixSet;
    /// # let mut set: CsrMatrixSet<f64> = CsrMatrixSet::default();
    /// let view = set.as_view();
    /// let midpoint = view.len() / 2;
    /// let (left, right) = view.split_at(midpoint);
    ///
    /// // Process halves independently (parallel processing example)
    /// // This is conceptual - actual parallel processing would use rayon or similar
    /// ```
    ///
    /// # Memory Efficiency
    ///
    /// Both resulting views share references to the same underlying data:
    /// - No matrix data is copied during the split
    /// - Both views have independent lifetimes
    #[inline]
    pub fn split_at(self, index: usize) -> (Self, Self) {
        let (left_partition, right_partition) = self.partition.split_at(index);
        let (left_ncols, right_ncols) = self.ncols.split_at(index);
        let left = CsrMatrixSetView {
            col_indices: self.col_indices,
            values: self.values,
            row_offsets: self.row_offsets,
            ncols: left_ncols,
            partition: left_partition,
        };
        let right = CsrMatrixSetView {
            col_indices: self.col_indices,
            values: self.values,
            row_offsets: self.row_offsets,
            ncols: right_ncols,
            partition: right_partition,
        };
        (left, right)
    }
}

impl<'a, T> IntoView for &'a CsrMatrixSet<T> {
    type View = CsrMatrixSetView<'a, T>;

    #[inline]
    fn into_view(self) -> Self::View {
        self.as_view()
    }
}

pub trait CsrMatrixSetMethods<V> {
    /// Returns the number of matrices in the set.
    fn len(&self) -> usize;

    /// Returns true if the set contains no matrices.
    #[inline]
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<'a, T, V> CsrMatrixSetMethods<V> for &'a T
where
    &'a T: IntoView<View = CsrMatrixSetView<'a, V>>,
    V: Real,
{
    #[inline]
    fn len(&self) -> usize {
        self.into_view().len()
    }
}

/// A builder for constructing CSR matrices within a `CsrMatrixSet`.
///
/// This builder allows efficient construction of CSR matrices that will be stored
/// in a matrix set. When the builder is dropped, the matrix is automatically
/// added to the parent set.
///
/// # Examples
///
/// ```rust
/// use algebra_sparse::CsrMatrixSet;
///
/// let mut set = CsrMatrixSet::default();
/// let mut builder = set.new_matrix(3, 1e-10);
///
/// let mut row_builder = builder.new_row();
/// row_builder.push(0, 1.0);
/// row_builder.push(2, 2.0);
/// // Matrix is automatically added to set when builder is dropped
/// ```
pub struct CsrMatrixBuilder<'a, T> {
    set: &'a mut CsrMatrixSet<T>,
    /// The value index start for this matrix in the set's `values` array.
    value_start: usize,
    /// The row index start for this matrix in the set's `row_offsets` array.
    row_start: usize,
    /// The number of columns of this matrix.
    ncol: usize,
    /// Values below this threshold will be ignored during construction.
    zero_threshold: T,
}

/// Automatically finalizes the CSR matrix and adds it to the set when the builder is dropped.
///
/// This implementation ensures that the matrix is properly added to the parent set
/// when the builder goes out of scope, including updating all necessary metadata.
impl<T> Drop for CsrMatrixBuilder<'_, T> {
    fn drop(&mut self) {
        self.set.ncols.push(self.ncol);
        let partition = Partition {
            value_offset: self.value_start,
            value_len: self.set.values.len() - self.value_start,
            row_offset: self.row_start,
            row_len: self.set.row_offsets.len() - self.row_start,
        };
        self.set.partition.push(partition);
    }
}

impl<T: Real> CsrMatrixBuilder<'_, T> {
    /// Returns the number of columns for the matrix being built.
    #[inline]
    pub fn ncol(&self) -> usize {
        self.ncol
    }

    /// Creates a new row builder for this matrix.
    ///
    /// The returned builder can be used to add non-zero elements to the next row.
    ///
    /// # Returns
    /// A `CsVecBuilder` for constructing a sparse row
    #[inline]
    pub fn new_row(&mut self) -> CsVecBuilder<T> {
        CsVecBuilder::from_parts_unchecked(
            &mut self.set.col_indices,
            &mut self.set.row_offsets,
            &mut self.set.values,
            self.value_start,
            self.zero_threshold,
        )
    }
}

#[cfg(test)]
mod tests {
    use approx::assert_relative_eq;

    use super::*;
    use crate::csm::CsrMatrixViewMethods;
    use crate::traits::IntoView;

    /// Helper function to create a test matrix set with multiple matrices
    fn create_test_matrix_set() -> CsrMatrixSet<f32> {
        let mut set = CsrMatrixSet::default();

        // First matrix: 3x3 sparse matrix
        {
            let mut builder = set.new_matrix(3, 1e-10);
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(0, 1.0), (2, 2.0)]);
            }
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(1, 3.0)]);
            }
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(0, 4.0), (1, 5.0)]);
            }
        }

        // Second matrix: 2x2 sparse matrix
        {
            let mut builder = set.new_matrix(2, 1e-10);
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(0, 6.0), (1, 7.0)]);
            }
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(1, 8.0)]);
            }
        }

        // Third matrix: 4x1 sparse matrix
        {
            let mut builder = set.new_matrix(1, 1e-10);
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(0, 9.0)]);
            }
        }

        // Fourth matrix: 2x3 sparse matrix
        {
            let mut builder = set.new_matrix(3, 1e-10);
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(1, 10.0)]);
            }
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(0, 11.0), (2, 12.0)]);
            }
        }

        set
    }

    #[test]
    fn test_split_at_beginning() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        // Split at index 0: left empty, right contains all matrices
        let (left, right) = view.split_at(0);

        assert_eq!(left.len(), 0);
        assert_eq!(right.len(), 4);

        // Verify right view contains all matrices in original order
        for i in 0..4 {
            let original_matrix = set.get(i);
            let split_matrix = right.get(i);
            assert_relative_eq!(original_matrix.to_dense(), split_matrix.to_dense());
        }
    }

    #[test]
    fn test_split_at_end() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        // Split at index 4: left contains all matrices, right empty
        let (left, right) = view.split_at(4);

        assert_eq!(left.len(), 4);
        assert_eq!(right.len(), 0);

        // Verify left view contains all matrices in original order
        for i in 0..4 {
            let original_matrix = set.get(i);
            let split_matrix = left.get(i);
            assert_relative_eq!(original_matrix.to_dense(), split_matrix.to_dense());
        }
    }

    #[test]
    fn test_split_at_middle() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        // Split at index 2: left contains matrices 0,1; right contains matrices 2,3
        let (left, right) = view.split_at(2);

        assert_eq!(left.len(), 2);
        assert_eq!(right.len(), 2);

        // Verify left view contains first two matrices
        for i in 0..2 {
            let original_matrix = set.get(i);
            let split_matrix = left.get(i);
            assert_relative_eq!(original_matrix.to_dense(), split_matrix.to_dense());
        }

        // Verify right view contains last two matrices
        for i in 0..2 {
            let original_matrix = set.get(i + 2);
            let split_matrix = right.get(i);
            assert_relative_eq!(original_matrix.to_dense(), split_matrix.to_dense());
        }
    }

    #[test]
    fn test_split_at_various_positions() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        // Test splitting at each possible position
        for split_index in 0..=4 {
            let (left, right) = view.split_at(split_index);

            assert_eq!(left.len(), split_index);
            assert_eq!(right.len(), 4 - split_index);

            // Verify all matrices are correctly distributed
            for i in 0..4 {
                let original_matrix = set.get(i);
                let split_matrix = if i < split_index {
                    left.get(i)
                } else {
                    right.get(i - split_index)
                };
                assert_relative_eq!(original_matrix.to_dense(), split_matrix.to_dense());
            }
        }
    }

    #[test]
    fn test_split_multiple_times() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        // First split at index 2
        let (left, right) = view.split_at(2);

        // Split left part at index 1
        let (left_left, left_right) = left.split_at(1);

        // Split right part at index 1
        let (right_left, right_right) = right.split_at(1);

        // Verify all parts have correct lengths
        assert_eq!(left_left.len(), 1);
        assert_eq!(left_right.len(), 1);
        assert_eq!(right_left.len(), 1);
        assert_eq!(right_right.len(), 1);

        // Verify matrices are correctly distributed
        assert_relative_eq!(set.get(0).to_dense(), left_left.get(0).to_dense());
        assert_relative_eq!(set.get(1).to_dense(), left_right.get(0).to_dense());
        assert_relative_eq!(set.get(2).to_dense(), right_left.get(0).to_dense());
        assert_relative_eq!(set.get(3).to_dense(), right_right.get(0).to_dense());
    }

    #[test]
    fn test_split_single_matrix() {
        let mut set = CsrMatrixSet::default();

        // Create a set with just one matrix
        {
            let mut builder = set.new_matrix(2, 1e-10);
            {
                let mut row = builder.new_row();
                row.extend_with_nonzeros(vec![(0, 1.0), (1, 2.0)]);
            }
        }

        let view = (&set).into_view();

        // Split at index 0
        let (left, right) = view.split_at(0);
        assert_eq!(left.len(), 0);
        assert_eq!(right.len(), 1);

        // Split at index 1
        let (left, right) = view.split_at(1);
        assert_eq!(left.len(), 1);
        assert_eq!(right.len(), 0);

        assert_relative_eq!(set.get(0).to_dense(), left.get(0).to_dense());
    }

    #[test]
    fn test_split_empty_view() {
        let set = CsrMatrixSet::<f32>::default();
        let view = set.as_view();

        // Split empty view at index 0
        let (left, right) = view.split_at(0);

        assert_eq!(left.len(), 0);
        assert_eq!(right.len(), 0);
        assert!(left.is_empty());
        assert!(right.is_empty());
    }

    #[test]
    fn test_split_view_data_integrity() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        let (left, right) = view.split_at(2);

        // Verify that all views share the same underlying data
        // by checking that modifications through one view are visible through others
        // (Note: since we're working with views, we can't modify the data,
        // but we can verify the data integrity by comparing matrices)

        let original_matrices: Vec<_> = (0..4).map(|i| set.get(i).to_dense()).collect();
        let left_matrices: Vec<_> = (0..2).map(|i| left.get(i).to_dense()).collect();
        let right_matrices: Vec<_> = (0..2).map(|i| right.get(i).to_dense()).collect();

        // Verify data integrity
        for (i, original) in original_matrices.iter().enumerate() {
            let split_matrix = if i < 2 {
                &left_matrices[i]
            } else {
                &right_matrices[i - 2]
            };
            assert_relative_eq!(original, split_matrix);
        }
    }

    #[test]
    fn test_split_view_independence() {
        let set = create_test_matrix_set();
        let view = set.as_view();

        let (left1, right1) = view.split_at(2);
        let (left2, right2) = view.split_at(2);

        // Both splits should produce identical results
        assert_eq!(left1.len(), left2.len());
        assert_eq!(right1.len(), right2.len());

        for i in 0..left1.len() {
            assert_relative_eq!(left1.get(i).to_dense(), left2.get(i).to_dense());
        }

        for i in 0..right1.len() {
            assert_relative_eq!(right1.get(i).to_dense(), right2.get(i).to_dense());
        }
    }
}