use crate::index::dim_to_u32;
use crate::real::Real;
use crate::{CscMatrix, SparseError};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Perm {
pub fwd: Vec<u32>,
pub inv: Vec<u32>,
}
impl Perm {
pub fn identity(n: usize) -> Self {
assert!(dim_to_u32("permutation", n).is_ok(), "n exceeds u32::MAX");
let fwd: Vec<u32> = (0..n).map(idx32).collect();
Self {
inv: fwd.clone(),
fwd,
}
}
pub fn from_fwd(fwd: Vec<u32>) -> Result<Self, SparseError> {
let n = fwd.len();
dim_to_u32("permutation", n)?;
let mut inv = vec![u32::MAX; n];
for (k, &i) in fwd.iter().enumerate() {
let i = i as usize;
if i >= n {
return Err(SparseError::IndexOutOfRange {
what: "permutation entry",
index: i,
bound: n,
});
}
if inv[i] != u32::MAX {
return Err(SparseError::Malformed(format!(
"(P1): index {i} appears twice in the permutation"
)));
}
inv[i] = idx32(k);
}
let p = Self { fwd, inv };
debug_assert_eq!(p.check_invariants(), Ok(()));
Ok(p)
}
#[inline]
pub fn len(&self) -> usize {
self.fwd.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.fwd.is_empty()
}
#[must_use]
pub fn inverse(&self) -> Self {
Self {
fwd: self.inv.clone(),
inv: self.fwd.clone(),
}
}
pub fn is_identity(&self) -> bool {
self.fwd.iter().enumerate().all(|(k, &i)| i as usize == k)
}
pub fn check_invariants(&self) -> Result<(), SparseError> {
let n = self.fwd.len();
if self.inv.len() != n {
return Err(SparseError::Dimension {
what: "inverse permutation",
expected: n,
found: self.inv.len(),
});
}
for (k, &i) in self.fwd.iter().enumerate() {
let i = i as usize;
if i >= n {
return Err(SparseError::IndexOutOfRange {
what: "permutation entry",
index: i,
bound: n,
});
}
if self.inv[i] as usize != k {
return Err(SparseError::Malformed(format!(
"(P1): inv[fwd[{k}]] = {} ≠ {k}",
self.inv[i]
)));
}
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Pattern {
pub n_rows: usize,
pub col_ptr: Vec<usize>,
pub row_idx: Vec<u32>,
}
impl Pattern {
pub fn from_raw_parts(
n_rows: usize,
col_ptr: Vec<usize>,
row_idx: Vec<u32>,
) -> Result<Self, SparseError> {
let p = Self {
n_rows,
col_ptr,
row_idx,
};
p.check_invariants()?;
Ok(p)
}
pub fn empty(n_rows: usize, n_cols: usize) -> Self {
assert!(dim_to_u32("rows", n_rows).is_ok() && dim_to_u32("columns", n_cols).is_ok());
Self {
n_rows,
col_ptr: vec![0; n_cols + 1],
row_idx: Vec::new(),
}
}
#[inline]
pub fn n_rows(&self) -> usize {
self.n_rows
}
#[inline]
pub fn n_cols(&self) -> usize {
self.col_ptr.len().saturating_sub(1)
}
#[inline]
pub fn nnz(&self) -> usize {
self.row_idx.len()
}
#[inline]
pub fn col(&self, j: usize) -> &[u32] {
&self.row_idx[self.col_ptr[j]..self.col_ptr[j + 1]]
}
#[must_use]
pub fn transpose(&self) -> Self {
let mut ptr = vec![0_usize; self.n_rows + 1];
for &i in &self.row_idx {
ptr[i as usize + 1] += 1;
}
for i in 0..self.n_rows {
ptr[i + 1] += ptr[i];
}
let mut next = ptr[..self.n_rows].to_vec();
let mut idx = vec![0_u32; self.nnz()];
for j in 0..self.n_cols() {
for &i in self.col(j) {
idx[next[i as usize]] = idx32(j);
next[i as usize] += 1;
}
}
let t = Self {
n_rows: self.n_cols(),
col_ptr: ptr,
row_idx: idx,
};
debug_assert_eq!(t.check_invariants(), Ok(()));
t
}
pub fn symmetric(&self) -> Result<Self, SparseError> {
let n = self.n_cols();
if self.n_rows != n {
return Err(SparseError::Dimension {
what: "pattern (must be square to symmetrize)",
expected: n,
found: self.n_rows,
});
}
let t = self.transpose();
let mut col_ptr = Vec::with_capacity(n + 1);
let mut row_idx = Vec::with_capacity(2 * self.nnz());
col_ptr.push(0);
for j in 0..n {
let (a, b) = (self.col(j), t.col(j));
let (mut p, mut q) = (0, 0);
while p < a.len() || q < b.len() {
let x = a.get(p).copied().unwrap_or(u32::MAX);
let y = b.get(q).copied().unwrap_or(u32::MAX);
let next = x.min(y);
p += usize::from(x == next);
q += usize::from(y == next);
row_idx.push(next);
}
col_ptr.push(row_idx.len());
}
let s = Self {
n_rows: n,
col_ptr,
row_idx,
};
debug_assert_eq!(s.check_invariants(), Ok(()));
Ok(s)
}
#[must_use]
pub fn ata(&self) -> Self {
let n = self.n_cols();
let t = self.transpose();
let mut mark = vec![usize::MAX; n];
let mut col_ptr = Vec::with_capacity(n + 1);
let mut row_idx = Vec::new();
col_ptr.push(0);
for j in 0..n {
let start = row_idx.len();
for &i in self.col(j) {
for &k in t.col(i as usize) {
if mark[k as usize] != j {
mark[k as usize] = j;
row_idx.push(k);
}
}
}
row_idx[start..].sort_unstable();
col_ptr.push(row_idx.len());
}
let s = Self {
n_rows: n,
col_ptr,
row_idx,
};
debug_assert_eq!(s.check_invariants(), Ok(()));
s
}
pub fn cholesky_nnz(&self, perm: &Perm) -> Result<usize, SparseError> {
let n = self.n_cols();
if self.n_rows != n || perm.len() != n {
return Err(SparseError::Dimension {
what: "cholesky_nnz (square pattern and permutation of its order)",
expected: n,
found: if self.n_rows == n {
perm.len()
} else {
self.n_rows
},
});
}
perm.check_invariants()?;
let (_, counts) = etree_counts(&self.symmetric()?, perm);
Ok(counts.iter().sum())
}
#[cfg_attr(not(feature = "faer"), allow(dead_code))] pub(crate) fn etree_postorder(&self, perm: &Perm) -> Perm {
let n = perm.len();
debug_assert_eq!(self.n_cols(), n);
let (parent, counts) = etree_counts(self, perm);
let none = usize::MAX;
let mut kids: Vec<(usize, usize, usize)> = (0..n)
.filter(|&k| parent[k] != none)
.map(|k| (parent[k], counts[k], k))
.collect();
kids.sort_unstable();
let mut start = vec![0_usize; n + 1];
for &(p, _, _) in &kids {
start[p + 1] += 1;
}
for k in 0..n {
start[k + 1] += start[k];
}
let mut fwd = Vec::with_capacity(n);
let mut stack: Vec<(usize, usize)> = Vec::new();
for root in (0..n).filter(|&k| parent[k] == none) {
stack.push((root, start[root]));
while let Some(top) = stack.last_mut() {
let (x, next) = *top;
if next < start[x + 1] {
top.1 += 1;
let c = kids[next].2;
stack.push((c, start[c]));
} else {
stack.pop();
fwd.push(perm.fwd[x]);
}
}
}
let post = Perm::from_fwd(fwd).expect("(P4): a postorder is a permutation");
debug_assert!(
(0..n).all(|k| parent[k] == none
|| post.inv[perm.fwd[k] as usize] < post.inv[perm.fwd[parent[k]] as usize]),
"(P4): a node before its parent"
);
post
}
pub fn check_invariants(&self) -> Result<(), SparseError> {
dim_to_u32("rows", self.n_rows)?;
dim_to_u32("columns", self.n_cols())?;
if self.col_ptr.is_empty() {
return Err(SparseError::Malformed("(P2): empty pointer array".into()));
}
if self.col_ptr[0] != 0 || self.col_ptr[self.n_cols()] != self.row_idx.len() {
return Err(SparseError::Malformed(format!(
"(P2): pointers must run from 0 to nnz = {}",
self.row_idx.len()
)));
}
for j in 0..self.n_cols() {
if self.col_ptr[j] > self.col_ptr[j + 1] {
return Err(SparseError::Malformed(format!(
"(P2): pointer of column {j} decreases"
)));
}
let col = self.col(j);
for (p, &i) in col.iter().enumerate() {
if i as usize >= self.n_rows {
return Err(SparseError::IndexOutOfRange {
what: "pattern row index",
index: i as usize,
bound: self.n_rows,
});
}
if p > 0 && col[p - 1] >= i {
return Err(SparseError::Malformed(format!(
"(P2): rows of column {j} not strictly increasing"
)));
}
}
}
Ok(())
}
}
impl<T: Real> CscMatrix<T> {
pub fn pattern(&self) -> Pattern {
Pattern {
n_rows: self.n_rows,
col_ptr: self.col_ptr.clone(),
row_idx: self.row_idx.clone(),
}
}
pub fn symmetric_pattern(&self) -> Result<Pattern, SparseError> {
self.pattern().symmetric()
}
}
#[inline]
fn etree_counts(g: &Pattern, perm: &Perm) -> (Vec<usize>, Vec<usize>) {
let n = perm.len();
let none = usize::MAX;
let mut parent = vec![none; n];
let mut mark = vec![none; n];
let mut counts = vec![1_usize; n];
for k in 0..n {
mark[k] = k;
for &i in g.col(perm.fwd[k] as usize) {
let mut x = perm.inv[i as usize] as usize;
if x >= k {
continue;
}
while mark[x] != k {
mark[x] = k;
counts[x] += 1;
if parent[x] == none {
parent[x] = k;
break;
}
x = parent[x];
}
}
}
(parent, counts)
}
pub(crate) fn idx32(i: usize) -> u32 {
u32::try_from(i).expect("index exceeds u32::MAX")
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
fn pattern_of(m: usize, n: usize, entries: &[(usize, usize)]) -> Pattern {
let mut cols: Vec<Vec<u32>> = vec![Vec::new(); n];
for &(i, j) in entries {
cols[j % n].push(idx32(i % m));
}
let mut col_ptr = vec![0];
let mut row_idx = Vec::new();
for mut c in cols {
c.sort_unstable();
c.dedup();
row_idx.extend(c);
col_ptr.push(row_idx.len());
}
Pattern::from_raw_parts(m, col_ptr, row_idx).unwrap()
}
fn cholesky_nnz_dense(s: &Pattern, perm: &Perm) -> usize {
let n = s.n_cols();
let mut adj = vec![vec![false; n]; n];
for j in 0..n {
for &i in s.col(j) {
let (a, b) = (perm.inv[i as usize] as usize, perm.inv[j] as usize);
if a != b {
adj[a][b] = true;
adj[b][a] = true;
}
}
}
let mut count = n;
for k in 0..n {
let later: Vec<usize> = (k + 1..n).filter(|&l| adj[k][l]).collect();
count += later.len();
for &a in &later {
for &b in &later {
if a != b {
adj[a][b] = true;
}
}
}
}
count
}
#[test]
fn perm_constructors_and_errors() {
let p = Perm::identity(4);
assert!(p.is_identity() && p.check_invariants().is_ok() && p.len() == 4);
let q = Perm::from_fwd(vec![2, 0, 3, 1]).unwrap();
assert_eq!(q.inv, vec![1, 3, 0, 2]);
assert_eq!(q.inverse().inverse(), q);
assert!(!q.is_identity() && !q.is_empty());
assert!(Perm::from_fwd(vec![0, 0]).is_err());
assert!(Perm::from_fwd(vec![0, 2]).is_err());
let bad = Perm {
fwd: vec![0, 1],
inv: vec![1, 0],
};
assert!(bad.check_invariants().is_err());
assert!(Perm::identity(0).is_empty());
}
#[test]
fn pattern_validation_and_access() {
assert!(Pattern::from_raw_parts(2, vec![0, 2], vec![1, 0]).is_err());
assert!(Pattern::from_raw_parts(2, vec![0, 1], vec![2]).is_err());
assert!(Pattern::from_raw_parts(2, vec![0, 2, 1], vec![0, 1]).is_err());
let p = Pattern::from_raw_parts(3, vec![0, 2, 2, 3], vec![0, 2, 1]).unwrap();
assert_eq!((p.n_rows(), p.n_cols(), p.nnz()), (3, 3, 3));
assert_eq!(p.col(0), &[0, 2]);
assert_eq!(p.col(1), []);
assert!(Pattern::empty(2, 3).check_invariants().is_ok());
let a = CscMatrix::<f64>::from_raw_parts(
3,
vec![0, 2, 2, 3],
vec![0, 2, 1],
vec![1.0, 2.0, 3.0],
)
.unwrap();
assert_eq!(a.pattern(), p);
let s = a.symmetric_pattern().unwrap();
assert_eq!(s.col(0), &[0, 2]);
assert_eq!(s.col(1), &[2]);
assert_eq!(s.col(2), &[0, 1]);
assert!(Pattern::empty(2, 3).symmetric().is_err());
assert!(p.cholesky_nnz(&Perm::identity(2)).is_err());
}
proptest! {
#[test]
fn pattern_algebra(m in 1_usize..12, n in 1_usize..12,
e in prop::collection::vec((0_usize..12, 0_usize..12), 0..60)) {
let s = pattern_of(m, n, &e);
prop_assert_eq!(s.transpose().transpose(), s.clone());
let g = s.ata();
prop_assert_eq!(g.transpose(), g.clone());
for j in 0..n {
if !s.col(j).is_empty() {
prop_assert!(g.col(j).contains(&idx32(j)));
}
}
let q = pattern_of(n, n, &e);
let sym = q.symmetric().unwrap();
prop_assert_eq!(sym.transpose(), sym.clone());
for j in 0..n {
for &i in q.col(j) {
prop_assert!(sym.col(j).contains(&i));
}
}
}
#[test]
fn cholesky_count_matches_dense_elimination(
n in 1_usize..25,
e in prop::collection::vec((0_usize..25, 0_usize..25), 0..80),
keys in prop::collection::vec(any::<u32>(), 25),
) {
let s = pattern_of(n, n, &e);
let mut fwd: Vec<u32> = (0..n).map(idx32).collect();
fwd.sort_by_key(|&i| keys[i as usize]);
let perm = Perm::from_fwd(fwd).unwrap();
prop_assert_eq!(s.cholesky_nnz(&perm).unwrap(), cholesky_nnz_dense(&s, &perm));
}
#[test]
fn etree_postorder_keeps_the_fill(
n in 1_usize..25,
e in prop::collection::vec((0_usize..25, 0_usize..25), 0..80),
keys in prop::collection::vec(any::<u32>(), 25),
) {
let s = pattern_of(n, n, &e).symmetric().unwrap();
let mut fwd: Vec<u32> = (0..n).map(idx32).collect();
fwd.sort_by_key(|&i| keys[i as usize]);
let perm = Perm::from_fwd(fwd).unwrap();
let post = s.etree_postorder(&perm);
prop_assert_eq!(post.check_invariants(), Ok(()));
prop_assert_eq!(s.cholesky_nnz(&post).unwrap(), s.cholesky_nnz(&perm).unwrap());
let (parent, counts) = etree_counts(&s, &post);
let mut size = vec![1_usize; n];
for k in 0..n {
if parent[k] != usize::MAX {
prop_assert!(parent[k] > k);
size[parent[k]] += size[k];
}
}
for (k, &p) in parent.iter().enumerate() {
if p != usize::MAX {
prop_assert!(k + size[p] > p);
}
}
for p in 0..n {
let kids: Vec<usize> = (0..n).filter(|&k| parent[k] == p).collect();
if let Some(&last) = kids.iter().max() {
prop_assert_eq!(last + 1, p);
let heaviest = kids.iter().map(|&k| counts[k]).max().unwrap();
prop_assert_eq!(counts[last], heaviest);
}
}
}
}
}