From 9bf1d0280d53a36770b6853708dd9e23e4dd64b3 Mon Sep 17 00:00:00 2001 From: sebcrozet Date: Tue, 30 Oct 2018 17:29:32 +0100 Subject: [PATCH] Fix cholesky computation. --- src/sparse/cs_matrix.rs | 34 +++ src/sparse/cs_matrix_analysis.rs | 184 ------------ src/sparse/cs_matrix_cholesky.rs | 330 +++++++++++++++++++++ src/sparse/mod.rs | 6 +- tests/sparse/cs_cholesky.rs | 55 ++++ tests/sparse/{cs_linalg.rs => cs_solve.rs} | 0 tests/sparse/mod.rs | 3 +- 7 files changed, 425 insertions(+), 187 deletions(-) delete mode 100644 src/sparse/cs_matrix_analysis.rs create mode 100644 tests/sparse/cs_cholesky.rs rename tests/sparse/{cs_linalg.rs => cs_solve.rs} (100%) diff --git a/src/sparse/cs_matrix.rs b/src/sparse/cs_matrix.rs index 1a28932c..ec4cc5b0 100644 --- a/src/sparse/cs_matrix.rs +++ b/src/sparse/cs_matrix.rs @@ -52,6 +52,15 @@ where pub(crate) vals: Vec, } +impl CsVecStorage +where + DefaultAllocator: Allocator, +{ + pub fn values(&self) -> &[N] { + &self.vals + } +} + impl CsVecStorage where DefaultAllocator: Allocator {} impl<'a, N: Scalar, R: Dim, C: Dim> CsStorageIter<'a, N, R, C> for CsVecStorage @@ -187,10 +196,35 @@ where } impl> CsMatrix { + pub fn from_data(data: S) -> Self { + CsMatrix { + data, + _phantoms: PhantomData, + } + } + pub fn len(&self) -> usize { self.data.len() } + pub fn nrows(&self) -> usize { + self.data.shape().0.value() + } + + pub fn ncols(&self) -> usize { + self.data.shape().1.value() + } + + pub fn shape(&self) -> (usize, usize) { + let (nrows, ncols) = self.data.shape(); + (nrows.value(), ncols.value()) + } + + pub fn is_square(&self) -> bool { + let (nrows, ncols) = self.data.shape(); + nrows.value() == ncols.value() + } + pub fn transpose(&self) -> CsMatrix where DefaultAllocator: Allocator, diff --git a/src/sparse/cs_matrix_analysis.rs b/src/sparse/cs_matrix_analysis.rs deleted file mode 100644 index 629904a4..00000000 --- a/src/sparse/cs_matrix_analysis.rs +++ /dev/null @@ -1,184 +0,0 @@ -use alga::general::{ClosedAdd, ClosedMul}; -use num::{One, Zero}; -use std::iter; -use std::marker::PhantomData; -use std::ops::{Add, Mul, Range}; -use std::slice; - -use allocator::Allocator; -use constraint::{AreMultipliable, DimEq, SameNumberOfRows, ShapeConstraint}; -use sparse::{CsMatrix, CsStorage, CsVector}; -use storage::{Storage, StorageMut}; -use {DefaultAllocator, Dim, Matrix, MatrixMN, Real, Scalar, Vector, VectorN, U1}; - -pub struct SymbolicAnalysis { - pinv: Vec, - q: Vec, - elimination_tree: Vec, - cp: Vec, - leftmost: Vec, - m2: usize, - lnz: usize, - unz: usize, -} - -#[derive(Copy, Clone, Debug)] -pub struct EliminationTreeNode { - parent: usize, -} - -impl EliminationTreeNode { - pub fn root() -> Self { - EliminationTreeNode { - parent: usize::max_value(), - } - } - - pub fn with_parent(parent: usize) -> Self { - EliminationTreeNode { parent } - } - - pub fn is_root(&self) -> bool { - self.parent == usize::max_value() - } - - pub fn parent(&self) -> usize { - self.parent - } -} - -impl> CsMatrix { - fn elimination_tree(&self) -> Vec { - let (nrows, ncols) = self.data.shape(); - assert_eq!( - nrows.value(), - ncols.value(), - "The matrix `self` must be square to compute its elimination tree." - ); - - let mut forest: Vec<_> = iter::repeat(EliminationTreeNode::root()) - .take(nrows.value()) - .collect(); - let mut ancestor: Vec<_> = iter::repeat(usize::max_value()) - .take(nrows.value()) - .collect(); - - for k in 0..nrows.value() { - for irow in self.data.column_row_indices(k) { - let mut i = irow; - - while i < k { - let i_ancestor = ancestor[i]; - ancestor[i] = k; - - if i_ancestor == usize::max_value() { - forest[i] = EliminationTreeNode::with_parent(k); - break; - } - - i = i_ancestor; - } - } - } - - forest - } - - fn reach( - &self, - j: usize, - max_j: usize, - tree: &[EliminationTreeNode], - marks: &mut Vec, - out: &mut Vec, - ) { - marks.clear(); - marks.resize(tree.len(), false); - - for irow in self.data.column_row_indices(j) { - let mut curr = irow; - while curr != usize::max_value() && curr <= max_j && !marks[curr] { - marks[curr] = true; - out.push(curr); - curr = tree[curr].parent; - } - } - } - - fn column_counts(&self, tree: &[EliminationTreeNode]) -> Vec { - let len = self.data.shape().0.value(); - let mut counts: Vec<_> = iter::repeat(0).take(len).collect(); - let mut reach = Vec::new(); - let mut marks = Vec::new(); - - for i in 0..len { - self.reach(i, i, tree, &mut marks, &mut reach); - - for j in reach.drain(..) { - counts[j] += 1; - } - } - - counts - } - - fn tree_postorder(tree: &[EliminationTreeNode]) -> Vec { - // FIXME: avoid all those allocations? - let mut first_child: Vec<_> = iter::repeat(usize::max_value()).take(tree.len()).collect(); - let mut other_children: Vec<_> = - iter::repeat(usize::max_value()).take(tree.len()).collect(); - - // Build the children list from the parent list. - // The set of children of the node `i` is given by the linked list - // starting at `first_child[i]`. The nodes of this list are then: - // { first_child[i], other_children[first_child[i]], other_children[other_children[first_child[i]], ... } - for (i, node) in tree.iter().enumerate() { - if !node.is_root() { - let brother = first_child[node.parent]; - first_child[node.parent] = i; - other_children[i] = brother; - } - } - - let mut stack = Vec::with_capacity(tree.len()); - let mut postorder = Vec::with_capacity(tree.len()); - - for (i, node) in tree.iter().enumerate() { - if node.is_root() { - Self::dfs( - i, - &mut first_child, - &other_children, - &mut stack, - &mut postorder, - ) - } - } - - postorder - } - - fn dfs( - i: usize, - first_child: &mut [usize], - other_children: &[usize], - stack: &mut Vec, - result: &mut Vec, - ) { - stack.clear(); - stack.push(i); - - while let Some(n) = stack.pop() { - let child = first_child[n]; - - if child == usize::max_value() { - // No children left. - result.push(n); - } else { - stack.push(n); - stack.push(child); - first_child[n] = other_children[child]; - } - } - } -} diff --git a/src/sparse/cs_matrix_cholesky.rs b/src/sparse/cs_matrix_cholesky.rs index 8b137891..2826b555 100644 --- a/src/sparse/cs_matrix_cholesky.rs +++ b/src/sparse/cs_matrix_cholesky.rs @@ -1 +1,331 @@ +use alga::general::{ClosedAdd, ClosedMul}; +use num::{One, Zero}; +use std::iter; +use std::marker::PhantomData; +use std::mem; +use std::ops::{Add, Mul, Range}; +use std::slice; +use allocator::Allocator; +use constraint::{AreMultipliable, DimEq, SameNumberOfRows, ShapeConstraint}; +use sparse::{CsMatrix, CsStorage, CsStorageIter, CsVecStorage, CsVector}; +use storage::{Storage, StorageMut}; +use {DefaultAllocator, Dim, Matrix, MatrixMN, Real, Scalar, Vector, VectorN, U1}; + +pub struct CsCholesky +where + DefaultAllocator: Allocator + Allocator, +{ + // Non-zero pattern of the original matrix upper-triangular part. + // Unlike the original matrix, the `original_p` array does contain the last sentinel value + // equal to `original_i.len()` at the end. + original_p: Vec, + original_i: Vec, + original_len: usize, // Number of elements on the numerical value vector of the original matrix. + // Decomposition result. + l: CsMatrix, + // Used only for the pattern. + // FIXME: store only the nonzero pattern instead. + u: CsMatrix, + ok: bool, + // Workspaces. + work_x: VectorN, + work_c: VectorN, +} + +impl CsCholesky +where + DefaultAllocator: Allocator + Allocator, +{ + /// Computes the cholesky decomposition of the sparse matrix `m`. + pub fn new(m: &CsMatrix) -> Self { + let mut me = Self::new_symbolic(m); + let _ = me.decompose(&m.data.vals); + me + } + /// Perform symbolic analysis for the given matrix. + /// + /// This does not access the numerical values of `m`. + pub fn new_symbolic(m: &CsMatrix) -> Self { + assert!( + m.is_square(), + "The matrix `m` must be square to compute its elimination tree." + ); + + let (l, u) = Self::nonzero_pattern(m); + + // Workspaces. + let work_x = unsafe { VectorN::new_uninitialized_generic(m.data.shape().0, U1) }; + let work_c = unsafe { VectorN::new_uninitialized_generic(m.data.shape().1, U1) }; + let mut original_p = m.data.p.as_slice().to_vec(); + original_p.push(m.data.i.len()); + + CsCholesky { + original_p, + original_i: m.data.i.clone(), + original_len: m.data.i.len(), + l, + u, + ok: false, + work_x, + work_c, + } + } + + pub fn l(&self) -> Option<&CsMatrix> { + if self.ok { + Some(&self.l) + } else { + None + } + } + + pub fn unwrap_l(self) -> Option> { + if self.ok { + Some(self.l) + } else { + None + } + } + + // Performs the numerical Cholesky decomposition given the set of numerical values. + pub fn decompose(&mut self, values: &[N]) -> bool { + assert!( + values.len() >= self.original_len, + "The set of values is too small." + ); + + // Reset `work_c` to the column pointers of `l`. + self.work_c.copy_from(&self.l.data.p); + + // Perform the decomposition. + for k in 0..self.l.nrows() { + // Scatter the k-th column of the original matrix with the values provided. + let column_range = self.original_p[k]..self.original_p[k + 1]; + + self.work_x[k] = N::zero(); + for p in column_range.clone() { + let irow = self.original_i[p]; + + if irow <= k { + self.work_x[irow] = values[p]; + } + } + + let mut diag = self.work_x[k]; + self.work_x[k] = N::zero(); + + // Triangular solve. + for irow in self.u.data.column_row_indices(k) { + if irow >= k { + continue; + } + + let lki = self.work_x[irow] / self.l.data.vals[self.l.data.p[irow]]; + self.work_x[irow] = N::zero(); + + for p in self.l.data.p[irow] + 1..self.work_c[irow] { + self.work_x[self.l.data.i[p]] -= self.l.data.vals[p] * lki; + } + + diag -= lki * lki; + let p = self.work_c[irow]; + self.work_c[irow] += 1; + self.l.data.i[p] = k; + self.l.data.vals[p] = lki; + } + + if diag <= N::zero() { + self.ok = false; + return false; + } + + // Deal with the diagonal element. + let p = self.work_c[k]; + self.work_c[k] += 1; + self.l.data.i[p] = k; + self.l.data.vals[p] = diag.sqrt(); + } + + self.ok = true; + true + } + + fn elimination_tree>(m: &CsMatrix) -> Vec { + let nrows = m.nrows(); + let mut forest: Vec<_> = iter::repeat(usize::max_value()).take(nrows).collect(); + let mut ancestor: Vec<_> = iter::repeat(usize::max_value()).take(nrows).collect(); + + for k in 0..nrows { + for irow in m.data.column_row_indices(k) { + let mut i = irow; + + while i < k { + let i_ancestor = ancestor[i]; + ancestor[i] = k; + + if i_ancestor == usize::max_value() { + forest[i] = k; + break; + } + + i = i_ancestor; + } + } + } + + forest + } + + fn reach>( + m: &CsMatrix, + j: usize, + max_j: usize, + tree: &[usize], + marks: &mut Vec, + out: &mut Vec, + ) { + marks.clear(); + marks.resize(tree.len(), false); + + // FIXME: avoid all those allocations. + let mut tmp = Vec::new(); + let mut res = Vec::new(); + + for irow in m.data.column_row_indices(j) { + let mut curr = irow; + while curr != usize::max_value() && curr <= max_j && !marks[curr] { + marks[curr] = true; + tmp.push(curr); + curr = tree[curr]; + } + + tmp.append(&mut res); + mem::swap(&mut tmp, &mut res); + } + + out.append(&mut res); + } + + fn nonzero_pattern>( + m: &CsMatrix, + ) -> (CsMatrix, CsMatrix) { + let etree = Self::elimination_tree(m); + let (nrows, ncols) = m.data.shape(); + let mut rows = Vec::with_capacity(m.len()); + let mut cols = unsafe { VectorN::new_uninitialized_generic(m.data.shape().0, U1) }; + let mut marks = Vec::new(); + + // NOTE: the following will actually compute the non-zero pattern of + // the transpose of l. + for i in 0..nrows.value() { + cols[i] = rows.len(); + Self::reach(m, i, i, &etree, &mut marks, &mut rows); + } + + let mut vals = Vec::with_capacity(rows.len()); + unsafe { + vals.set_len(rows.len()); + } + vals.shrink_to_fit(); + + let data = CsVecStorage { + shape: (nrows, ncols), + p: cols, + i: rows, + vals, + }; + + let u = CsMatrix::from_data(data); + // XXX: avoid this transpose. + let l = u.transpose(); + + (l, u) + } + + /* + * + * NOTE: All the following methods are untested and currently unused. + * + * + fn column_counts>( + m: &CsMatrix, + tree: &[usize], + ) -> Vec { + let len = m.data.shape().0.value(); + let mut counts: Vec<_> = iter::repeat(0).take(len).collect(); + let mut reach = Vec::new(); + let mut marks = Vec::new(); + + for i in 0..len { + Self::reach(m, i, i, tree, &mut marks, &mut reach); + + for j in reach.drain(..) { + counts[j] += 1; + } + } + + counts + } + + fn tree_postorder(tree: &[usize]) -> Vec { + // FIXME: avoid all those allocations? + let mut first_child: Vec<_> = iter::repeat(usize::max_value()).take(tree.len()).collect(); + let mut other_children: Vec<_> = + iter::repeat(usize::max_value()).take(tree.len()).collect(); + + // Build the children list from the parent list. + // The set of children of the node `i` is given by the linked list + // starting at `first_child[i]`. The nodes of this list are then: + // { first_child[i], other_children[first_child[i]], other_children[other_children[first_child[i]], ... } + for (i, parent) in tree.iter().enumerate() { + if *parent != usize::max_value() { + let brother = first_child[*parent]; + first_child[*parent] = i; + other_children[i] = brother; + } + } + + let mut stack = Vec::with_capacity(tree.len()); + let mut postorder = Vec::with_capacity(tree.len()); + + for (i, node) in tree.iter().enumerate() { + if *node == usize::max_value() { + Self::dfs( + i, + &mut first_child, + &other_children, + &mut stack, + &mut postorder, + ) + } + } + + postorder + } + + fn dfs( + i: usize, + first_child: &mut [usize], + other_children: &[usize], + stack: &mut Vec, + result: &mut Vec, + ) { + stack.clear(); + stack.push(i); + + while let Some(n) = stack.pop() { + let child = first_child[n]; + + if child == usize::max_value() { + // No children left. + result.push(n); + } else { + stack.push(n); + stack.push(child); + first_child[n] = other_children[child]; + } + } + } + */ +} diff --git a/src/sparse/mod.rs b/src/sparse/mod.rs index 320d76a0..6ce898e5 100644 --- a/src/sparse/mod.rs +++ b/src/sparse/mod.rs @@ -1,7 +1,9 @@ -pub use self::cs_matrix::{CsMatrix, CsStorage, CsStorageMut, CsVector}; +pub use self::cs_matrix::{ + CsMatrix, CsStorage, CsStorageIter, CsStorageMut, CsVecStorage, CsVector, +}; +pub use self::cs_matrix_cholesky::CsCholesky; mod cs_matrix; -mod cs_matrix_analysis; mod cs_matrix_cholesky; mod cs_matrix_conversion; mod cs_matrix_ops; diff --git a/tests/sparse/cs_cholesky.rs b/tests/sparse/cs_cholesky.rs new file mode 100644 index 00000000..9c199737 --- /dev/null +++ b/tests/sparse/cs_cholesky.rs @@ -0,0 +1,55 @@ +#![cfg_attr(rustfmt, rustfmt_skip)] + +use na::{CsMatrix, CsVector, CsCholesky, Cholesky, Matrix5, Vector5}; + +#[test] +fn cs_cholesky() { + let mut a = Matrix5::new( + 40.0, 0.0, 0.0, 0.0, 0.0, + 2.0, 60.0, 0.0, 0.0, 0.0, + 1.0, 0.0, 11.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 50.0, 0.0, + 1.0, 0.0, 0.0, 4.0, 10.0 + ); + a.fill_upper_triangle_with_lower_triangle(); + test_cholesky(a); + + let a = Matrix5::from_diagonal(&Vector5::new(40.0, 60.0, 11.0, 50.0, 10.0)); + test_cholesky(a); + + let mut a = Matrix5::new( + 40.0, 0.0, 0.0, 0.0, 0.0, + 2.0, 60.0, 0.0, 0.0, 0.0, + 1.0, 0.0, 11.0, 0.0, 0.0, + 1.0, 0.0, 0.0, 50.0, 0.0, + 0.0, 0.0, 0.0, 4.0, 10.0 + ); + a.fill_upper_triangle_with_lower_triangle(); + test_cholesky(a); + + let mut a = Matrix5::new( + 2.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 2.0, 0.0, 0.0, 0.0, + 1.0, 1.0, 2.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 2.0, 0.0, + 1.0, 1.0, 0.0, 0.0, 2.0 + ); + a.fill_upper_triangle_with_lower_triangle(); + test_cholesky(a); +} + + +fn test_cholesky(a: Matrix5) { + let cs_a: CsMatrix<_, _, _> = a.into(); + + let chol_a = Cholesky::new(a).unwrap(); + let chol_cs_a = CsCholesky::new(&cs_a); + let l = chol_a.l(); + println!("{:?}", chol_cs_a.l()); + let cs_l: Matrix5<_> = chol_cs_a.unwrap_l().unwrap().into(); + + println!("{}", l); + println!("{}", cs_l); + + assert_eq!(l, cs_l); +} diff --git a/tests/sparse/cs_linalg.rs b/tests/sparse/cs_solve.rs similarity index 100% rename from tests/sparse/cs_linalg.rs rename to tests/sparse/cs_solve.rs diff --git a/tests/sparse/mod.rs b/tests/sparse/mod.rs index 6c4d6d45..0e772c99 100644 --- a/tests/sparse/mod.rs +++ b/tests/sparse/mod.rs @@ -1,5 +1,6 @@ +mod cs_cholesky; mod cs_construction; mod cs_conversion; mod cs_matrix; mod cs_ops; -mod cs_linalg; \ No newline at end of file +mod cs_solve;