From fe9bef1e3ffd29254c3254d8d163f92869cfe69b Mon Sep 17 00:00:00 2001 From: julianknodt Date: Wed, 23 Aug 2023 00:25:05 -0700 Subject: [PATCH] Implement SparseLU factorization Add `solve_upper_triangular` to `CsrMatrix` This allows a sparse matrix to be used for efficient solving with a dense LU decomposition. --- nalgebra-sparse/src/csc.rs | 83 ++++++++++++- nalgebra-sparse/src/csr.rs | 51 +++++++- nalgebra-sparse/src/factorization/lu.rs | 16 +++ nalgebra-sparse/src/factorization/mod.rs | 4 + nalgebra-sparse/src/lib.rs | 7 ++ nalgebra-sparse/src/pattern.rs | 142 +++++++++++++++++++++++ nalgebra-sparse/tests/sparsity.rs | 67 +++++++++++ nalgebra-sparse/tests/unit_tests/csc.rs | 48 ++++++++ nalgebra-sparse/tests/unit_tests/csr.rs | 32 +++++ 9 files changed, 448 insertions(+), 2 deletions(-) create mode 100644 nalgebra-sparse/src/factorization/lu.rs create mode 100644 nalgebra-sparse/tests/sparsity.rs diff --git a/nalgebra-sparse/src/csc.rs b/nalgebra-sparse/src/csc.rs index 0cd89f5c2..ab9207cb2 100644 --- a/nalgebra-sparse/src/csc.rs +++ b/nalgebra-sparse/src/csc.rs @@ -12,7 +12,7 @@ use crate::csr::CsrMatrix; use crate::pattern::{SparsityPattern, SparsityPatternFormatError, SparsityPatternIter}; use crate::{SparseEntry, SparseEntryMut, SparseFormatError, SparseFormatErrorKind}; -use nalgebra::Scalar; +use nalgebra::{RealField, Scalar}; use num_traits::One; use std::slice::{Iter, IterMut}; @@ -553,6 +553,87 @@ impl CscMatrix { self.filter(|i, j, _| i >= j) } + /// Solves a lower triangular system, `self` is a matrix of NxN, and `b` is a column vector of size N + /// Assuming that b is dense. + // TODO add an option here for assuming diagonal is one. + pub fn dense_lower_triangular_solve(&self, b: &[T], out: &mut [T], unit_diagonal: bool) + where + T: RealField + Copy, + { + assert_eq!(self.nrows(), self.ncols()); + assert_eq!(self.ncols(), b.len()); + assert_eq!(out.len(), b.len()); + out.copy_from_slice(b); + let n = b.len(); + + for i in 0..n { + let mul = out[i]; + let col = self.col(i); + for (&ri, &v) in col.row_indices().iter().zip(col.values().iter()) { + // ensure that only using the lower part + if ri == i { + if !unit_diagonal { + out[ri] /= v; + } + } else if ri > i { + out[ri] -= v * mul; + } + } + } + } + + /// Solves a sparse lower triangular system `Ax = b`, with both the matrix and vector + /// sparse. + /// sparsity_idxs should be precomputed using the sparse_lower_triangle. + /// Assumes that the diagonal of the sparse matrix is all 1. + pub fn sparse_lower_triangular_solve( + &self, + b_idxs: &[usize], + b: &[T], + // idx -> row + // for now, is permitted to be unsorted + // TODO maybe would be better to enforce sorted, but would have to sort internally. + out_sparsity_pattern: &[usize], + out: &mut [T], + assume_unit: bool, + ) where + T: RealField + Copy, + { + assert_eq!(self.nrows(), self.ncols()); + assert_eq!(b.len(), b_idxs.len()); + assert!(b_idxs.iter().all(|&bi| bi < self.ncols())); + + assert_eq!(out_sparsity_pattern.len(), out.len()); + assert!(out_sparsity_pattern.iter().all(|&i| i < self.ncols())); + + // initialize out with b + for (&bv, &bi) in b.iter().zip(b_idxs.iter()) { + let out_pos = out_sparsity_pattern.iter().position(|&p| p == bi).unwrap(); + out[out_pos] = bv; + } + + for (i, &row) in out_sparsity_pattern.iter().enumerate() { + let mul = out[i]; + let col = self.col(row); + for (ni, &nrow) in out_sparsity_pattern.iter().enumerate() { + if nrow < row { + continue; + } + // TODO in a sorted version may be able to iterate without + // having the cost of binary search at each iteration + let l_val = col + .get_entry(nrow) + .map(|v| v.into_value()) + .unwrap_or_else(T::zero); + if !assume_unit && (ni == i || nrow == row) { + out[i] /= l_val; + } else if nrow > row { + out[ni] -= l_val * mul; + } + } + } + } + /// Returns the diagonal of the matrix as a sparse matrix. #[must_use] pub fn diagonal_as_csc(&self) -> Self diff --git a/nalgebra-sparse/src/csr.rs b/nalgebra-sparse/src/csr.rs index 255e84047..8e0a3ae53 100644 --- a/nalgebra-sparse/src/csr.rs +++ b/nalgebra-sparse/src/csr.rs @@ -12,7 +12,7 @@ use crate::csc::CscMatrix; use crate::pattern::{SparsityPattern, SparsityPatternFormatError, SparsityPatternIter}; use crate::{SparseEntry, SparseEntryMut, SparseFormatError, SparseFormatErrorKind}; -use nalgebra::Scalar; +use nalgebra::{DMatrix, DMatrixView, RealField, Scalar}; use num_traits::One; use std::slice::{Iter, IterMut}; @@ -573,6 +573,55 @@ impl CsrMatrix { { CscMatrix::from(self).transpose_as_csr() } + + /// Solves the equation `Ax = b`, treating `self` as an upper triangular matrix. + /// If `A` is not upper triangular, elements in the lower triangle below the diagonal + /// will be ignored. + /// + /// If `m` has zeros along the diagonal, returns `None`. + /// Panics: + /// Panics if `A` and `b` have incompatible shapes, specifically if `b` + /// has a different number of rows and than `a`. + pub fn solve_upper_triangular<'a>(&self, b: impl Into>) -> Option> + where + T: RealField + Scalar, + { + // https://www.nicolasboumal.net/papers/MAT321_Lecture_notes_Boumal_2019.pdf + // page 48 + let b: DMatrixView<'a, T> = b.into(); + assert_eq!(b.nrows(), self.nrows()); + + let out_cols = b.ncols(); + let out_rows = self.nrows(); + + let mut out = DMatrix::zeros(out_rows, out_cols); + for r in (0..out_rows).rev() { + let row = self.row(r); + // only take upper triangle elements + let mut row_iter = row + .col_indices() + .iter() + .copied() + .zip(row.values().iter()) + .filter(|&(c, _)| c >= r); + + let (c, div) = row_iter.next()?; + // This implies there is a 0 on the diagonal + if c != r || div.is_zero() { + return None; + } + for c in 0..out_cols { + let numer = b.index((r, c)).clone(); + let numer = numer + - row_iter + .clone() + .map(|(a_col, val)| val.clone() * b.index((a_col, c)).clone()) + .fold(T::zero(), |acc, n| acc + n); + *out.index_mut((r, c)) = numer / div.clone(); + } + } + Some(out) + } } /// Convert pattern format errors into more meaningful CSR-specific errors. diff --git a/nalgebra-sparse/src/factorization/lu.rs b/nalgebra-sparse/src/factorization/lu.rs new file mode 100644 index 000000000..a6053fbc2 --- /dev/null +++ b/nalgebra-sparse/src/factorization/lu.rs @@ -0,0 +1,16 @@ +use crate::CscMatrix; +use nalgebra::RealField; + +pub struct LeftLookingLUFactorization { + /// A single matrix stores both the lower and upper triangular components + l_u: CscMatrix +} + +impl LeftLookingLUFactorization { + /// Construct a new sparse LU factorization + /// from a given CSC matrix. + pub fn new(a: &CscMatrix) -> Self { + assert_eq!(a.nrows(), a.ncols()); + todo!(); + } +} diff --git a/nalgebra-sparse/src/factorization/mod.rs b/nalgebra-sparse/src/factorization/mod.rs index b77a857c0..e12cb011c 100644 --- a/nalgebra-sparse/src/factorization/mod.rs +++ b/nalgebra-sparse/src/factorization/mod.rs @@ -4,3 +4,7 @@ mod cholesky; pub use cholesky::*; + +//mod lu; + +//pub use lu::*; diff --git a/nalgebra-sparse/src/lib.rs b/nalgebra-sparse/src/lib.rs index c62677c36..fc12509e7 100644 --- a/nalgebra-sparse/src/lib.rs +++ b/nalgebra-sparse/src/lib.rs @@ -279,4 +279,11 @@ impl<'a, T: Clone + Zero> SparseEntryMut<'a, T> { SparseEntryMut::Zero => T::zero(), } } + /// If the entry is nonzero, returns `Some(&mut value)`, otherwise returns `None`. + pub fn nonzero(self) -> Option<&'a mut T> { + match self { + SparseEntryMut::NonZero(v) => Some(v), + SparseEntryMut::Zero => None, + } + } } diff --git a/nalgebra-sparse/src/pattern.rs b/nalgebra-sparse/src/pattern.rs index c51945b7c..3440cfff0 100644 --- a/nalgebra-sparse/src/pattern.rs +++ b/nalgebra-sparse/src/pattern.rs @@ -61,6 +61,14 @@ impl SparsityPattern { minor_dim, } } + /// Creates the sparsity pattern of an identity matrix of size `n`. + pub fn identity(n: usize) -> Self { + Self { + major_offsets: (0..=n).collect(), + minor_indices: (0..n).collect(), + minor_dim: n, + } + } /// The offsets for the major dimension. #[inline] @@ -109,6 +117,13 @@ impl SparsityPattern { pub fn lane(&self, major_index: usize) -> &[usize] { self.get_lane(major_index).unwrap() } + /// Returns an iterator over all lanes in this sparsity pattern, in order. + /// Does not omit empty lanes. + #[inline] + #[must_use] + pub fn lanes(&self) -> impl Iterator { + (0..self.major_offsets.len() - 1).map(move |r| self.lane(r)) + } /// Get the lane at the given index, or `None` if out of bounds. #[inline] @@ -289,6 +304,133 @@ impl SparsityPattern { ) .expect("Internal error: Transpose should never fail.") } + + /// Computes the output sparsity pattern of `x` in `Ax = b`. + /// where A's nonzero pattern is given by `self` and the non-zero indices + /// of vector `b` are specified as a slice. + /// The output is not necessarily in sorted order, but is topological sort order. + /// Treats `self` as lower triangular, even if there are elements in the upper triangle. + /// Acts as if b is one major lane (i.e. CSC matrix and one column) + pub fn sparse_lower_triangular_solve(&self, b: &[usize]) -> Vec { + assert!(b.iter().all(|&i| i < self.major_dim())); + let mut out = vec![]; + + // From a given starting column, traverses and finds all reachable indices. + fn reach(sp: &SparsityPattern, j: usize, out: &mut Vec) { + // already traversed + if out.contains(&j) { + return; + } + + out.push(j); + for &i in sp.lane(j) { + reach(sp, i, out); + } + } + + for &i in b { + reach(&self, i, &mut out); + } + + out + } + + /// Left-looking Sparse LU decomposition from Gilbert/Peierls. + /// returns the sparsity pattern of the output + pub fn left_looking_lu_decomposition(&self) -> SparsityPattern { + assert_eq!(self.major_dim(), self.minor_dim()); + let n = self.minor_dim(); + let mut sp = SparsityPatternBuilder::new(n, n); + for col in 0..n { + let mut x = sp + .valid_partial() + .sparse_lower_triangular_solve(self.lane(col)); + x.sort_unstable(); + for row in x { + assert!(sp.insert(col, row).is_ok()); + } + } + sp.build() + } +} + +/// A builder that allows for constructing a sparsity pattern. +/// It requires elements to be added in sorted order. Specifically, +/// For each element the major must be >= the previous element's major. +/// If the major is the same, the minor must be in ascending order. +#[derive(Debug, Clone, PartialEq)] +pub struct SparsityPatternBuilder { + buf: SparsityPattern, + major_dim: usize, +} + +/// +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum BuilderInsertError { + /// + MajorTooLow, + /// + MinorTooLow, +} + +impl SparsityPatternBuilder { + /// Constructs a new empty builder. + pub fn new(major_dim: usize, minor_dim: usize) -> Self { + Self { + buf: SparsityPattern { + major_offsets: vec![0], + minor_indices: vec![], + minor_dim, + }, + major_dim, + } + } + /// Allows for general assignment of indices + pub fn insert(&mut self, maj: usize, min: usize) -> Result<(), BuilderInsertError> { + assert!(maj < self.major_dim); + assert!(min < self.buf.minor_dim); + + let curr_major = self.buf.major_dim(); + + // cannot go backwards in major + if maj < curr_major { + return Err(BuilderInsertError::MajorTooLow); + } + // cannot go backwards in minor + if maj == curr_major + && !self.buf.minor_indices.is_empty() + && min <= *self.buf.minor_indices.last().unwrap() + { + return Err(BuilderInsertError::MinorTooLow); + } + // add any advances in row. + for _ in curr_major..maj { + self.buf.major_offsets.push(self.buf.minor_indices.len()); + } + self.buf.minor_indices.push(min); + Ok(()) + } + /// Returns a valid partial sparsity pattern. + /// All the major lanes up to the current insertion will be completed. + pub(crate) fn valid_partial(&mut self) -> &SparsityPattern { + if *self.buf.major_offsets.last().unwrap() != self.buf.minor_indices.len() { + self.buf.major_offsets.push(self.buf.minor_indices.len()); + } + &self.buf + } + /// Consumes self and outputs the constructed `SparsityPattern`. + /// If elements were added to the last major, but `advance_major` + /// was not called, will implicitly call `advance_major` then + /// output the values. + #[inline] + pub fn build(mut self) -> SparsityPattern { + let curr_major = self.buf.major_dim(); + for _ in curr_major..self.major_dim { + self.buf.major_offsets.push(self.buf.minor_indices.len()); + } + assert_eq!(self.buf.major_dim(), self.major_dim); + self.buf + } } /// Error type for `SparsityPattern` format errors. diff --git a/nalgebra-sparse/tests/sparsity.rs b/nalgebra-sparse/tests/sparsity.rs new file mode 100644 index 000000000..d70f46237 --- /dev/null +++ b/nalgebra-sparse/tests/sparsity.rs @@ -0,0 +1,67 @@ +use nalgebra_sparse::pattern::{SparsityPattern, SparsityPatternBuilder}; + +#[test] +fn sparsity_identity() { + let n = 100; + let speye = SparsityPattern::identity(n); + for (i, j) in speye.entries() { + assert_eq!(i, j); + } + assert_eq!(speye.major_dim(), n); + assert_eq!(speye.minor_dim(), n); +} + +#[test] +fn lower_sparse_solve() { + // just a smaller number so it's easier to debug + let n = 8; + let speye = SparsityPattern::identity(n); + assert_eq!(speye.sparse_lower_triangular_solve(&[0, 5]), vec![0, 5]); + + // test case from + // https://www.youtube.com/watch?v=1dGRTOwBkQs&ab_channel=TimDavis + let mut builder = SparsityPatternBuilder::new(14, 14); + // building CscMatrix, so it will be col, row + #[rustfmt::skip] + let indices = vec![ + (0, 0), (0, 2), + (1, 1), (1, 3), (1, 6), (1, 8), + (2,2), (2,4), (2,7), + (3,3), (3,8), + (4,4), (4,7), + (5,5), (5,8), (5,9), + (6,6), (6,9), (6,10), + (7,7), (7,9), + (8,8), (8,11), (8,12), + (9,9), (9,10), (9, 12), (9, 13), + (10,10), (10,11), (10,12), + (11,11), (11,12), + (12,12), (12,13), + (13,13), + ]; + for (maj, min) in indices.iter().copied() { + assert!(builder.insert(maj, min).is_ok()); + } + let sp = builder.build(); + assert_eq!(sp.major_dim(), 14); + assert_eq!(sp.minor_dim(), 14); + assert_eq!(sp.nnz(), indices.len()); + for ij in sp.entries() { + assert!(indices.contains(&ij)); + } + assert_eq!( + sp.sparse_lower_triangular_solve(&[3, 5]), + vec![3, 8, 11, 12, 13, 5, 9, 10] + ); +} + +#[test] +fn test_builder() { + let mut builder = SparsityPatternBuilder::new(2, 2); + assert!(builder.insert(0, 0).is_ok()); + assert!(builder.insert(0, 0).is_err()); + assert!(builder.insert(0, 1).is_ok()); + assert!(builder.insert(0, 1).is_err()); + assert!(builder.insert(1, 0).is_ok()); + assert!(builder.insert(1, 0).is_err()); +} diff --git a/nalgebra-sparse/tests/unit_tests/csc.rs b/nalgebra-sparse/tests/unit_tests/csc.rs index 1554b8a66..271f46636 100644 --- a/nalgebra-sparse/tests/unit_tests/csc.rs +++ b/nalgebra-sparse/tests/unit_tests/csc.rs @@ -721,3 +721,51 @@ proptest! { prop_assert_eq!(DMatrix::from(&csc), DMatrix::identity(n, n)); } } + +#[test] +fn test_dense_lower_triangular_solve() { + let mut a = CscMatrix::identity(3); + let v = [1., 2., 3.]; + let mut out = [0.; 3]; + a.dense_lower_triangular_solve(&v, &mut out, true); + assert_eq!(out, v); + a.dense_lower_triangular_solve(&v, &mut out, false); + assert_eq!(out, v); + if let SparseEntryMut::NonZero(e) = a.index_entry_mut(0, 0) { + *e = 2.; + }; + a.dense_lower_triangular_solve(&v, &mut out, false); + assert_eq!(out, [0.5, 2., 3.]); + a.dense_lower_triangular_solve(&v, &mut out, true); + assert_eq!(out, v); +} + +#[test] +fn test_sparse_lower_triangular_solve() { + let a = CscMatrix::identity(3); + let vi = [0, 1, 2]; + let v = [1., 2., 3.]; + let mut out = [0.; 3]; + a.sparse_lower_triangular_solve(&vi, &v, &vi, &mut out, true); + assert_eq!(out, v); + a.sparse_lower_triangular_solve(&vi, &v, &vi, &mut out, false); + assert_eq!(out, v); + + // now test with large identity matrix + let mut a = CscMatrix::identity(1000); + let vi = [0, 999]; + let v = [3., 2.]; + let mut out = [0.; 2]; + a.sparse_lower_triangular_solve(&vi, &v, &vi, &mut out, true); + assert_eq!(out, v); + a.sparse_lower_triangular_solve(&vi, &v, &vi, &mut out, false); + assert_eq!(out, v); + + if let SparseEntryMut::NonZero(e) = a.index_entry_mut(0, 0) { + *e = 2.; + }; + a.sparse_lower_triangular_solve(&vi, &v, &vi, &mut out, false); + assert_eq!(out, [1.5, 2.]); + a.sparse_lower_triangular_solve(&vi, &v, &vi, &mut out, true); + assert_eq!(out, v); +} diff --git a/nalgebra-sparse/tests/unit_tests/csr.rs b/nalgebra-sparse/tests/unit_tests/csr.rs index a00470d59..5574760a4 100644 --- a/nalgebra-sparse/tests/unit_tests/csr.rs +++ b/nalgebra-sparse/tests/unit_tests/csr.rs @@ -506,6 +506,38 @@ fn csr_matrix_get_index_entry() { } } +#[test] +fn csr_upper_triangle_solve_ok() { + const N: usize = 10; + let eye = CsrMatrix::::identity(N); + let b = DMatrix::from_row_slice(N, 1, &[0.5; N]); + assert_eq!(eye.solve_upper_triangular(&b), Some(b)); + + let mut eye = CsrMatrix::::identity(N); + let v = if let SparseEntryMut::NonZero(v) = eye.index_entry_mut(0, 0) { + v + } else { + unreachable!(); + }; + *v = 2.0; + let b = DMatrix::from_row_slice(N, 1, &[1.0; N]); + let mut x = b.clone(); + x[0] = 0.5; + assert_eq!(eye.solve_upper_triangular(&b), Some(x)); + + #[rustfmt::skip] + let dense = DMatrix::from_row_slice(3, 3, &[ + 1., 1., 1., + 0., 1., 1., + 0., 0., 1., + ]); + let csr = CsrMatrix::from(&dense); + + let b = DMatrix::from_row_slice(3, 1, &[1.0; 3]); + let x = DMatrix::from_row_slice(3, 1, &[-1., 0., 1.]); + assert_eq!(csr.solve_upper_triangular(&b), Some(x)); +} + #[test] fn csr_matrix_row_iter() { #[rustfmt::skip]