Clean up symbolic cholesky factorization

There were some inefficient allocations, and TODOs, this resolves them,
and adds some comments making it more clear why the code works as it does.
This commit is contained in:
julianknodt 2023-08-27 22:55:02 -07:00
parent f404bcbd6d
commit 493044e402
1 changed files with 37 additions and 52 deletions

View File

@ -2,7 +2,6 @@ use crate::csc::CscMatrix;
use crate::ops::serial::spsolve_csc_lower_triangular; use crate::ops::serial::spsolve_csc_lower_triangular;
use crate::ops::Op; use crate::ops::Op;
use crate::pattern::SparsityPattern; use crate::pattern::SparsityPattern;
use core::{iter, mem};
use nalgebra::{DMatrix, DMatrixView, DMatrixViewMut, RealField}; use nalgebra::{DMatrix, DMatrixView, DMatrixViewMut, RealField};
use std::fmt::{Display, Formatter}; use std::fmt::{Display, Formatter};
@ -291,56 +290,50 @@ impl<T: RealField> CscCholesky<T> {
} }
} }
/// For a given sparsity pattern, a specified row, and a precomputed elimination tree
/// marks is a buffer which indicates which nodes have been traversed, and is reset before each
/// use. `out` stores the row indices of the nonzero elements.
fn reach( fn reach(
pattern: &SparsityPattern, pattern: &SparsityPattern,
j: usize, row: usize,
max_j: usize, etree: &[usize],
tree: &[usize], marks: &mut [bool],
marks: &mut Vec<bool>,
out: &mut Vec<usize>, out: &mut Vec<usize>,
) { ) {
marks.clear(); assert_eq!(marks.len(), etree.len());
marks.resize(tree.len(), false); marks.fill(false);
// TODO: avoid all those allocations. let start_len = out.len();
let mut tmp = Vec::new();
let mut res = Vec::new();
for &irow in pattern.lane(j) { for mut curr in pattern.lane(row).iter().copied() {
let mut curr = irow; while curr != usize::MAX && curr <= row && !marks[curr] {
while curr != usize::max_value() && curr <= max_j && !marks[curr] {
marks[curr] = true; marks[curr] = true;
tmp.push(curr); out.push(curr);
curr = tree[curr]; curr = etree[curr];
} }
tmp.append(&mut res);
mem::swap(&mut tmp, &mut res);
} }
out[start_len..].sort_unstable();
res.sort_unstable();
out.append(&mut res);
} }
fn nonzero_pattern(m: &SparsityPattern) -> (SparsityPattern, SparsityPattern) { fn nonzero_pattern(m: &SparsityPattern) -> (SparsityPattern, SparsityPattern) {
let etree = elimination_tree(m); let etree = elimination_tree(m);
// Note: We assume CSC, therefore rows == minor and cols == major
let (nrows, ncols) = (m.minor_dim(), m.major_dim()); // note that m must be square.
let n = m.minor_dim();
let mut rows = Vec::with_capacity(m.nnz()); let mut rows = Vec::with_capacity(m.nnz());
let mut col_offsets = Vec::with_capacity(ncols + 1); let mut col_offsets = Vec::with_capacity(n + 1);
let mut marks = Vec::new(); col_offsets.push(0);
let mut marks = vec![false; etree.len()];
// NOTE: the following will actually compute the non-zero pattern of // NOTE: the following will actually compute the non-zero pattern of
// the transpose of l. // the transpose of l.
col_offsets.push(0); for i in 0..n {
for i in 0..nrows { reach(m, i, &etree, &mut marks, &mut rows);
reach(m, i, i, &etree, &mut marks, &mut rows);
col_offsets.push(rows.len()); col_offsets.push(rows.len());
} }
let u_pattern = let u_pattern = SparsityPattern::try_from_offsets_and_indices(n, n, col_offsets, rows).unwrap();
SparsityPattern::try_from_offsets_and_indices(nrows, ncols, col_offsets, rows).unwrap();
// TODO: Avoid this transpose? // TODO: Avoid this transpose?
let l_pattern = u_pattern.transpose(); let l_pattern = u_pattern.transpose();
@ -348,30 +341,22 @@ fn nonzero_pattern(m: &SparsityPattern) -> (SparsityPattern, SparsityPattern) {
(l_pattern, u_pattern) (l_pattern, u_pattern)
} }
/// Constructs the elimination tree for a given sparsity pattern.
/// The elimination tree is characterized as:
/// `parent[i] = min{ j > i | U[j,i] != 0 }`, where `U = L^T` is the transpose of the cholesky
/// matrix `L`.
fn elimination_tree(pattern: &SparsityPattern) -> Vec<usize> { fn elimination_tree(pattern: &SparsityPattern) -> Vec<usize> {
// Note: The pattern is assumed to of a CSC matrix, so the number of rows is // note that the tree is square so it doesn't matter if this is the major or minor dim
// given by the minor dimension let n = pattern.minor_dim();
let nrows = pattern.minor_dim(); let mut ancestor: Vec<_> = vec![usize::MAX; n];
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 (col, mut row) in pattern.entries() {
for &irow in pattern.lane(k) { while col > row {
let mut i = irow; let parent = ancestor[row];
ancestor[row] = ancestor[row].min(col);
while i < k { row = parent;
let i_ancestor = ancestor[i];
ancestor[i] = k;
if i_ancestor == usize::max_value() {
forest[i] = k;
break;
}
i = i_ancestor;
}
} }
} }
forest ancestor
} }