From b0ffd559628d82d98a19127d1be23ce62962d0f7 Mon Sep 17 00:00:00 2001 From: Andreas Longva Date: Fri, 17 Jul 2020 09:52:09 +0200 Subject: [PATCH] Initial CSR and SparsityPattern impls (WIP) --- nalgebra-sparse/src/csr.rs | 195 ++++++++++++++++++++++++ nalgebra-sparse/src/pattern.rs | 167 ++++++++++++++++++++ nalgebra-sparse/tests/unit_tests/coo.rs | 14 +- 3 files changed, 369 insertions(+), 7 deletions(-) create mode 100644 nalgebra-sparse/src/csr.rs create mode 100644 nalgebra-sparse/src/pattern.rs diff --git a/nalgebra-sparse/src/csr.rs b/nalgebra-sparse/src/csr.rs new file mode 100644 index 00000000..139c3499 --- /dev/null +++ b/nalgebra-sparse/src/csr.rs @@ -0,0 +1,195 @@ +use crate::{SparsityPattern, SparseFormatError}; +use crate::iter::SparsityPatternIter; + +use std::sync::Arc; +use std::slice::{IterMut, Iter}; + +/// A CSR representation of a sparse matrix. +/// +/// The Compressed Row Storage (CSR) format is well-suited as a general-purpose storage format +/// for many sparse matrix applications. +/// +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CsrMatrix { + // Rows are major, cols are minor in the sparsity pattern + sparsity_pattern: Arc, + values: Vec, +} + +impl CsrMatrix { + /// Create a zero CSR matrix with no explicitly stored entries. + pub fn new(nrows: usize, ncols: usize) -> Self { + Self { + sparsity_pattern: Arc::new(SparsityPattern::new(nrows, ncols)), + values: vec![], + } + } + + /// The number of rows in the matrix. + #[inline(always)] + pub fn nrows(&self) -> usize { + self.sparsity_pattern.major_dim() + } + + /// The number of columns in the matrix. + #[inline(always)] + pub fn ncols(&self) -> usize { + self.sparsity_pattern.minor_dim() + } + + /// The number of non-zeros in the matrix. + /// + /// Note that this corresponds to the number of explicitly stored entries, *not* the actual + /// number of algebraically zero entries in the matrix. Explicitly stored entries can still + /// be zero. Corresponds to the number of entries in the sparsity pattern. + #[inline(always)] + pub fn nnz(&self) -> usize { + self.sparsity_pattern.nnz() + } + + /// The row offsets defining part of the CSR format. + #[inline(always)] + pub fn row_offsets(&self) -> &[usize] { + self.sparsity_pattern.major_offsets() + } + + /// The column indices defining part of the CSR format. + #[inline(always)] + pub fn column_indices(&self) -> &[usize] { + self.sparsity_pattern.minor_indices() + } + + /// The non-zero values defining part of the CSR format. + #[inline(always)] + pub fn values(&self) -> &[T] { + &self.values + } + + /// Mutable access to the non-zero values. + #[inline(always)] + pub fn values_mut(&mut self) -> &mut [T] { + &mut self.values + } + + /// Try to construct a CSR matrix from raw CSR data. + /// + /// It is assumed that each row contains unique and sorted column indices that are in + /// bounds with respect to the number of columns in the matrix. If this is not the case, + /// an error is returned to indicate the failure. + /// + /// Panics + /// ------ + /// Panics if the lengths of the provided arrays are not compatible with the CSR format. + pub fn try_from_csr_data( + num_rows: usize, + num_cols: usize, + row_offsets: Vec, + col_indices: Vec, + values: Vec, + ) -> Result { + assert_eq!(col_indices.len(), values.len(), + "Number of values and column indices must be the same"); + let pattern = SparsityPattern::try_from_offsets_and_indices( + num_rows, num_cols, row_offsets, col_indices)?; + Ok(Self { + sparsity_pattern: Arc::new(pattern), + values, + }) + } + + /// An iterator over non-zero triplets (i, j, v). + /// + /// The iteration happens in row-major fashion, meaning that i increases monotonically, + /// and j increases monotonically within each row. + /// + /// Examples + /// -------- + /// ``` + /// # use nalgebra_sparse::CsrMatrix; + /// let row_offsets = vec![0, 2, 3, 4]; + /// let col_indices = vec![0, 2, 1, 0]; + /// let values = vec![1, 2, 3, 4]; + /// let mut csr = CsrMatrix::try_from_csr_data(3, 4, row_offsets, col_indices, values) + /// .unwrap(); + /// + /// let triplets: Vec<_> = csr.triplet_iter().map(|(i, j, v)| (i, j, *v)).collect(); + /// assert_eq!(triplets, vec![(0, 0, 1), (0, 2, 2), (1, 1, 3), (2, 0, 4)]); + /// ``` + pub fn triplet_iter(&self) -> CsrTripletIter { + CsrTripletIter { + pattern_iter: self.sparsity_pattern.entries(), + values_iter: self.values.iter() + } + } + + /// A mutable iterator over non-zero triplets (i, j, v). + /// + /// Iteration happens in the same order as for [triplet_iter](#method.triplet_iter). + /// + /// Examples + /// -------- + /// ``` + /// # use nalgebra_sparse::CsrMatrix; + /// # let row_offsets = vec![0, 2, 3, 4]; + /// # let col_indices = vec![0, 2, 1, 0]; + /// # let values = vec![1, 2, 3, 4]; + /// // Using the same data as in the `triplet_iter` example + /// let mut csr = CsrMatrix::try_from_csr_data(3, 4, row_offsets, col_indices, values) + /// .unwrap(); + /// + /// // Zero out lower-triangular terms + /// csr.triplet_iter_mut() + /// .filter(|(i, j, _)| j < i) + /// .for_each(|(_, _, v)| *v = 0); + /// + /// let triplets: Vec<_> = csr.triplet_iter().map(|(i, j, v)| (i, j, *v)).collect(); + /// assert_eq!(triplets, vec![(0, 0, 1), (0, 2, 2), (1, 1, 3), (2, 0, 0)]); + /// ``` + pub fn triplet_iter_mut(&mut self) -> CsrTripletIterMut { + CsrTripletIterMut { + pattern_iter: self.sparsity_pattern.entries(), + values_mut_iter: self.values.iter_mut() + } + } +} + +#[derive(Debug)] +pub struct CsrTripletIter<'a, T> { + pattern_iter: SparsityPatternIter<'a>, + values_iter: Iter<'a, T> +} + +impl<'a, T> Iterator for CsrTripletIter<'a, T> { + type Item = (usize, usize, &'a T); + + fn next(&mut self) -> Option { + let next_entry = self.pattern_iter.next(); + let next_value = self.values_iter.next(); + + match (next_entry, next_value) { + (Some((i, j)), Some(v)) => Some((i, j, v)), + _ => None + } + } +} + +#[derive(Debug)] +pub struct CsrTripletIterMut<'a, T> { + pattern_iter: SparsityPatternIter<'a>, + values_mut_iter: IterMut<'a, T> +} + +impl<'a, T> Iterator for CsrTripletIterMut<'a, T> { + type Item = (usize, usize, &'a mut T); + + #[inline(always)] + fn next(&mut self) -> Option { + let next_entry = self.pattern_iter.next(); + let next_value = self.values_mut_iter.next(); + + match (next_entry, next_value) { + (Some((i, j)), Some(v)) => Some((i, j, v)), + _ => None + } + } +} \ No newline at end of file diff --git a/nalgebra-sparse/src/pattern.rs b/nalgebra-sparse/src/pattern.rs new file mode 100644 index 00000000..bb9006e6 --- /dev/null +++ b/nalgebra-sparse/src/pattern.rs @@ -0,0 +1,167 @@ +use crate::SparseFormatError; + +/// A representation of the sparsity pattern of a CSR or COO matrix. +#[derive(Debug, Clone, PartialEq, Eq)] +// TODO: Make SparsityPattern parametrized by index type +// (need a solid abstraction for index types though) +pub struct SparsityPattern { + major_offsets: Vec, + minor_indices: Vec, + minor_dim: usize, +} + +impl SparsityPattern { + /// Create a sparsity pattern of the given dimensions without explicitly stored entries. + pub fn new(major_dim: usize, minor_dim: usize) -> Self { + Self { + major_offsets: vec![0; major_dim + 1], + minor_indices: vec![], + minor_dim, + } + } + + /// The offsets for the major dimension. + #[inline(always)] + pub fn major_offsets(&self) -> &[usize] { + &self.major_offsets + } + + /// The indices for the minor dimension. + #[inline(always)] + pub fn minor_indices(&self) -> &[usize] { + &self.minor_indices + } + + /// The major dimension. + #[inline(always)] + pub fn major_dim(&self) -> usize { + assert!(self.major_offsets.len() > 0); + self.major_offsets.len() - 1 + } + + /// The minor dimension. + #[inline(always)] + pub fn minor_dim(&self) -> usize { + self.minor_dim + } + + /// The number of "non-zeros", i.e. explicitly stored entries in the pattern. + #[inline(always)] + pub fn nnz(&self) -> usize { + self.minor_indices.len() + } + + /// Get the lane at the given index. + #[inline(always)] + pub fn lane(&self, major_index: usize) -> Option<&[usize]> { + let offset_begin = *self.major_offsets().get(major_index)?; + let offset_end = *self.major_offsets().get(major_index + 1)?; + Some(&self.minor_indices()[offset_begin..offset_end]) + } + + /// Try to construct a sparsity pattern from the given dimensions, major offsets + /// and minor indices. + /// + /// Returns an error if the data does not conform to the requirements. + /// + /// TODO: Maybe we should not do any assertions in any of the construction functions + pub fn try_from_offsets_and_indices( + major_dim: usize, + minor_dim: usize, + major_offsets: Vec, + minor_indices: Vec, + ) -> Result { + assert_eq!(major_offsets.len(), major_dim + 1); + assert_eq!(*major_offsets.last().unwrap(), minor_indices.len()); + Ok(Self { + major_offsets, + minor_indices, + minor_dim, + }) + } + + /// An iterator over the explicitly stored "non-zero" entries (i, j). + /// + /// The iteration happens in a lane-major fashion, meaning that the lane index i + /// increases monotonically. and the minor index j increases monotonically within each + /// lane i. + /// + /// Examples + /// -------- + /// + /// ``` + /// # use nalgebra_sparse::{SparsityPattern}; + /// let offsets = vec![0, 2, 3, 4]; + /// let minor_indices = vec![0, 2, 1, 0]; + /// let pattern = SparsityPattern::try_from_offsets_and_indices(3, 4, offsets, minor_indices) + /// .unwrap(); + /// + /// let entries: Vec<_> = pattern.entries().collect(); + /// assert_eq!(entries, vec![(0, 0), (0, 2), (1, 1), (2, 0)]); + /// ``` + /// + pub fn entries(&self) -> SparsityPatternIter { + SparsityPatternIter::from_pattern(self) + } +} + +#[derive(Debug, Clone)] +pub struct SparsityPatternIter<'a> { + // See implementation of Iterator::next for an explanation of how these members are used + major_offsets: &'a [usize], + minor_indices: &'a [usize], + current_lane_idx: usize, + remaining_minors_in_lane: &'a [usize], +} + +impl<'a> SparsityPatternIter<'a> { + fn from_pattern(pattern: &'a SparsityPattern) -> Self { + let first_lane_end = pattern.major_offsets().get(1).unwrap_or(&0); + let minors_in_first_lane = &pattern.minor_indices()[0 .. *first_lane_end]; + Self { + major_offsets: pattern.major_offsets(), + minor_indices: pattern.minor_indices(), + current_lane_idx: 0, + remaining_minors_in_lane: minors_in_first_lane + } + } +} + +impl<'a> Iterator for SparsityPatternIter<'a> { + type Item = (usize, usize); + + #[inline(always)] + fn next(&mut self) -> Option { + // We ensure fast iteration across each lane by iteratively "draining" a slice + // corresponding to the remaining column indices in the particular lane. + // When we reach the end of this slice, we are at the end of a lane, + // and we must do some bookkeeping for preparing the iteration of the next lane + // (or stop iteration if we're through all lanes). + // This way we can avoid doing unnecessary bookkeeping on every iteration, + // instead paying a small price whenever we jump to a new lane. + if let Some(minor_idx) = self.remaining_minors_in_lane.first() { + let item = Some((self.current_lane_idx, *minor_idx)); + self.remaining_minors_in_lane = &self.remaining_minors_in_lane[1..]; + item + } else { + loop { + // Keep skipping lanes until we found a non-empty lane or there are no more lanes + if self.current_lane_idx + 2 >= self.major_offsets.len() { + // We've processed all lanes, so we're at the end of the iterator + // (note: keep in mind that offsets.len() == major_dim() + 1, hence we need +2) + return None; + } else { + // Bump lane index and check if the lane is non-empty + self.current_lane_idx += 1; + let lower = self.major_offsets[self.current_lane_idx]; + let upper = self.major_offsets[self.current_lane_idx + 1]; + if upper > lower { + self.remaining_minors_in_lane = &self.minor_indices[(lower + 1) .. upper]; + return Some((self.current_lane_idx, self.minor_indices[lower])) + } + } + } + } + } +} + diff --git a/nalgebra-sparse/tests/unit_tests/coo.rs b/nalgebra-sparse/tests/unit_tests/coo.rs index b7504d51..1ffa286b 100644 --- a/nalgebra-sparse/tests/unit_tests/coo.rs +++ b/nalgebra-sparse/tests/unit_tests/coo.rs @@ -1,4 +1,4 @@ -use nalgebra_sparse::{CooMatrix, SparsePatternError}; +use nalgebra_sparse::{CooMatrix, SparseFormatError}; use nalgebra::DMatrix; use crate::assert_panics; @@ -91,25 +91,25 @@ fn coo_try_from_triplets_reports_out_of_bounds_indices() { { // 0x0 matrix let result = CooMatrix::::try_from_triplets(0, 0, vec![0], vec![0], vec![2]); - assert!(matches!(result, Err(SparsePatternError::IndexOutOfBounds(_)))); + assert!(matches!(result, Err(SparseFormatError::IndexOutOfBounds(_)))); } { // 1x1 matrix, row out of bounds let result = CooMatrix::::try_from_triplets(1, 1, vec![1], vec![0], vec![2]); - assert!(matches!(result, Err(SparsePatternError::IndexOutOfBounds(_)))); + assert!(matches!(result, Err(SparseFormatError::IndexOutOfBounds(_)))); } { // 1x1 matrix, col out of bounds let result = CooMatrix::::try_from_triplets(1, 1, vec![0], vec![1], vec![2]); - assert!(matches!(result, Err(SparsePatternError::IndexOutOfBounds(_)))); + assert!(matches!(result, Err(SparseFormatError::IndexOutOfBounds(_)))); } { // 1x1 matrix, row and col out of bounds let result = CooMatrix::::try_from_triplets(1, 1, vec![1], vec![1], vec![2]); - assert!(matches!(result, Err(SparsePatternError::IndexOutOfBounds(_)))); + assert!(matches!(result, Err(SparseFormatError::IndexOutOfBounds(_)))); } { @@ -118,7 +118,7 @@ fn coo_try_from_triplets_reports_out_of_bounds_indices() { let j = vec![0, 2, 1, 3, 3]; let v = vec![2, 3, 7, 3, 1]; let result = CooMatrix::::try_from_triplets(3, 5, i, j, v); - assert!(matches!(result, Err(SparsePatternError::IndexOutOfBounds(_)))); + assert!(matches!(result, Err(SparseFormatError::IndexOutOfBounds(_)))); } { @@ -127,7 +127,7 @@ fn coo_try_from_triplets_reports_out_of_bounds_indices() { let j = vec![0, 2, 1, 5, 3]; let v = vec![2, 3, 7, 3, 1]; let result = CooMatrix::::try_from_triplets(3, 5, i, j, v); - assert!(matches!(result, Err(SparsePatternError::IndexOutOfBounds(_)))); + assert!(matches!(result, Err(SparseFormatError::IndexOutOfBounds(_)))); } }