use crate::SparseVec;
use crate::factor::{Factors, Lu, VACANT, idx32};
use crate::real::{Real, is_exact_zero, lt};
use crate::update::BasisUpdate;
pub const HYPER_HISTORY_WEIGHT: f64 = 0.05;
pub(crate) const NO_ETA: u32 = u32::MAX;
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct SolveStats {
pub calls: usize,
pub hyper_calls: usize,
pub rhs_density_sum: f64,
pub result_density_sum: f64,
}
impl SolveStats {
pub fn mean_rhs_density(&self) -> f64 {
self.mean(self.rhs_density_sum)
}
pub fn mean_result_density(&self) -> f64 {
self.mean(self.result_density_sum)
}
pub fn hyper_fraction(&self) -> f64 {
#[allow(clippy::cast_precision_loss)] let h = self.hyper_calls as f64;
self.mean(h)
}
fn mean(&self, sum: f64) -> f64 {
if self.calls == 0 {
0.0
} else {
#[allow(clippy::cast_precision_loss)]
let c = self.calls as f64;
sum / c
}
}
fn record(&mut self, hyper: bool, rhs_density: f64, result_density: f64) {
self.calls += 1;
self.hyper_calls += usize::from(hyper);
self.rhs_density_sum += rhs_density;
self.result_density_sum += result_density;
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Kind {
Ftran = 0,
Btran = 1,
FtranUpdate = 2,
}
#[derive(Clone, Debug)]
pub(crate) struct Workspace<T: Real> {
pub work: Vec<T>,
pub work2: Vec<T>,
pub mark: Vec<u32>,
pub stamp: u32,
pub stack: Vec<(u32, u32)>,
pub list: Vec<u32>,
pub post: Vec<u32>,
pub queued: Vec<bool>,
pub flag: Vec<bool>,
}
impl<T: Real> Default for Workspace<T> {
fn default() -> Self {
Self {
work: Vec::new(),
work2: Vec::new(),
mark: Vec::new(),
stamp: 0,
stack: Vec::new(),
list: Vec::new(),
post: Vec::new(),
queued: Vec::new(),
flag: Vec::new(),
}
}
}
impl<T: Real> Workspace<T> {
pub fn resize(&mut self, m: usize) {
if self.work.len() != m {
*self = Self {
work: vec![T::ZERO; m],
work2: vec![T::ZERO; m],
mark: vec![0; m],
flag: vec![false; m],
..Self::default()
};
}
}
pub fn new_stamp(&mut self) -> u32 {
if self.stamp == u32::MAX {
self.mark.fill(0);
self.stamp = 0;
}
self.stamp += 1;
self.stamp
}
}
fn reach<'g>(
mark: &mut [u32],
stamp: u32,
stack: &mut Vec<(u32, u32)>,
seeds: &[u32],
post: &mut Vec<u32>,
adj: impl Fn(usize) -> &'g [u32],
) {
for &s in seeds {
if mark[s as usize] == stamp {
continue;
}
mark[s as usize] = stamp;
stack.push((s, 0));
while let Some(top) = stack.last_mut() {
let (v, k) = *top;
let succ = adj(v as usize);
if let Some(&w) = succ.get(k as usize) {
top.1 = k + 1;
if mark[w as usize] != stamp {
mark[w as usize] = stamp;
stack.push((w, 0));
}
} else {
stack.pop();
post.push(v);
}
}
}
}
impl<T: Real, U: BasisUpdate<T>> Lu<T, U> {
pub(crate) fn ftran_kind(&mut self, rhs: &mut SparseVec<T>, kind: Kind) {
self.factors.assert_ready(rhs);
let rhs_density = rhs.density();
let hyper = self.use_hyper(rhs, kind);
let mut ws = std::mem::take(&mut self.factors.ws);
let for_update = kind == Kind::FtranUpdate;
if hyper {
self.factors
.ftran_hyper(&self.update, rhs, &mut ws, for_update);
} else {
self.factors
.ftran_dense(&self.update, rhs, &mut ws, for_update);
}
self.factors.ws = ws;
if for_update {
self.factors.spike_valid = true;
}
self.record(kind, hyper, rhs_density, rhs.density());
}
pub(crate) fn btran_impl(&mut self, rhs: &mut SparseVec<T>) {
self.factors.assert_ready(rhs);
let rhs_density = rhs.density();
let hyper = self.use_hyper(rhs, Kind::Btran);
let mut ws = std::mem::take(&mut self.factors.ws);
if hyper {
self.factors.btran_hyper(&self.update, rhs, &mut ws);
} else {
self.factors.btran_dense(&self.update, rhs, &mut ws);
}
self.factors.ws = ws;
self.record(Kind::Btran, hyper, rhs_density, rhs.density());
}
fn use_hyper(&self, rhs: &SparseVec<T>, kind: Kind) -> bool {
rhs.is_sparse()
&& !lt(self.opts.hyper_rhs_density, rhs.density())
&& !lt(
self.opts.hyper_result_density,
self.expected_density[kind as usize],
)
}
fn record(&mut self, kind: Kind, hyper: bool, rhs_density: f64, result_density: f64) {
let ema = &mut self.expected_density[kind as usize];
*ema = (1.0 - HYPER_HISTORY_WEIGHT) * *ema + HYPER_HISTORY_WEIGHT * result_density;
let st = &mut self.factors.stats;
let s = match kind {
Kind::Ftran => &mut st.ftran,
Kind::Btran => &mut st.btran,
Kind::FtranUpdate => &mut st.ftran_update,
};
s.record(hyper, rhs_density, result_density);
}
pub fn expected_densities(&self) -> [f64; 3] {
self.expected_density
}
}
impl<T: Real> Factors<T> {
fn assert_ready(&self, rhs: &SparseVec<T>) {
assert!(
self.valid,
"solve with a factor that holds no valid factorization"
);
assert_eq!(rhs.size(), self.m, "right-hand side has the wrong size");
}
fn record_spike(&mut self, y: &[T], rows: Option<&[u32]>) {
for &i in &self.spike_idx {
self.spike[i as usize] = T::ZERO;
}
self.spike_idx.clear();
match rows {
Some(rows) => {
for &i in rows {
self.spike[i as usize] = y[i as usize];
}
self.spike_idx.extend_from_slice(rows);
}
None => {
for (i, (&v, s)) in y.iter().zip(self.spike.iter_mut()).enumerate() {
if !is_exact_zero(v) {
*s = v;
self.spike_idx.push(idx32(i));
}
}
}
}
}
fn ftran_dense<U: BasisUpdate<T>>(
&mut self,
upd: &U,
rhs: &mut SparseVec<T>,
ws: &mut Workspace<T>,
for_update: bool,
) {
let y = rhs.values_mut();
self.apply_l(y);
upd.apply_ftran(y);
if for_update {
self.record_spike(y, None);
}
self.solve_u(y, &mut ws.work2);
rhs.rebuild_pattern();
}
fn btran_dense<U: BasisUpdate<T>>(
&self,
upd: &U,
rhs: &mut SparseVec<T>,
ws: &mut Workspace<T>,
) {
let c = rhs.values_mut();
let w = &mut ws.work2;
self.solve_ut(c, w);
upd.apply_btran(w);
self.apply_lt(w);
c.swap_with_slice(w);
rhs.rebuild_pattern();
}
pub(crate) fn apply_l(&self, y: &mut [T]) {
for (e, &p) in self.l_pivot.iter().enumerate() {
let y_p = y[p as usize];
if is_exact_zero(y_p) {
continue;
}
let r = self.l_start[e]..self.l_start[e + 1];
for (&i, &l) in self.l_index[r.clone()].iter().zip(&self.l_value[r]) {
y[i as usize] -= l * y_p;
}
}
}
fn solve_u(&self, y: &mut [T], x: &mut [T]) {
for s in (0..self.step_row.len()).rev() {
let (p, q) = (self.step_row[s], self.step_col[s]);
if p == VACANT {
continue;
}
let (p, q) = (p as usize, q as usize);
let x_q = y[p] / self.diag[p];
x[q] = x_q;
if is_exact_zero(x_q) {
continue;
}
let (idx, val) = self.u_cols.line(q);
for (&i, &u) in idx.iter().zip(val) {
y[i as usize] -= u * x_q;
}
}
y.copy_from_slice(x);
x.fill(T::ZERO);
}
fn solve_ut(&self, c: &mut [T], w: &mut [T]) {
for s in 0..self.step_row.len() {
let (p, q) = (self.step_row[s], self.step_col[s]);
if p == VACANT {
continue;
}
let (p, q) = (p as usize, q as usize);
let w_p = c[q] / self.diag[p];
c[q] = T::ZERO;
w[p] = w_p;
if is_exact_zero(w_p) {
continue;
}
let (idx, val) = self.u_rows.line(p);
for (&j, &u) in idx.iter().zip(val) {
c[j as usize] -= u * w_p;
}
}
}
fn apply_lt(&self, w: &mut [T]) {
for (e, &p) in self.l_pivot.iter().enumerate().rev() {
let r = self.l_start[e]..self.l_start[e + 1];
let s: T = self.l_index[r.clone()]
.iter()
.zip(&self.l_value[r])
.map(|(&i, &l)| l * w[i as usize])
.sum();
w[p as usize] -= s;
}
}
#[inline]
fn l_succ(&self, p: usize) -> &[u32] {
match self.l_eta_of_row[p] {
NO_ETA => &[],
e => &self.l_index[self.l_start[e as usize]..self.l_start[e as usize + 1]],
}
}
#[inline]
fn q_of(&self, p: usize) -> usize {
self.step_col[self.row_step[p]] as usize
}
#[inline]
fn p_of(&self, q: usize) -> usize {
self.step_row[self.col_step[q]] as usize
}
fn ftran_hyper<U: BasisUpdate<T>>(
&mut self,
upd: &U,
rhs: &mut SparseVec<T>,
ws: &mut Workspace<T>,
for_update: bool,
) {
let mut seeds = std::mem::take(&mut ws.list);
seeds.clear();
for &i in rhs.pattern().unwrap_or(&[]) {
let v = rhs.get(i);
if !is_exact_zero(v) {
ws.work[i as usize] = v;
seeds.push(i);
}
}
rhs.clear();
let stamp = ws.new_stamp();
ws.post.clear();
reach(
&mut ws.mark,
stamp,
&mut ws.stack,
&seeds,
&mut ws.post,
|p| self.l_succ(p),
);
let y = &mut ws.work;
for &p in ws.post.iter().rev() {
let e = self.l_eta_of_row[p as usize];
let y_p = y[p as usize];
if e == NO_ETA || is_exact_zero(y_p) {
continue;
}
let r = self.l_start[e as usize]..self.l_start[e as usize + 1];
for (&i, &l) in self.l_index[r.clone()].iter().zip(&self.l_value[r]) {
y[i as usize] -= l * y_p;
}
}
let mut list = seeds;
list.clear();
list.extend(ws.post.iter().rev());
upd.apply_ftran_hyper(
&mut ws.work,
&mut list,
&mut ws.mark,
ws.stamp,
&mut ws.queued,
);
if for_update {
self.record_spike(&ws.work, Some(&list));
}
let stamp = ws.new_stamp();
ws.post.clear();
reach(
&mut ws.mark,
stamp,
&mut ws.stack,
&list,
&mut ws.post,
|p| self.u_cols.line(self.q_of(p)).0,
);
let y = &mut ws.work;
for &p in ws.post.iter().rev() {
let p = p as usize;
let q = self.q_of(p);
let x_q = y[p] / self.diag[p];
y[p] = T::ZERO;
rhs.set(idx32(q), x_q);
if is_exact_zero(x_q) {
continue;
}
let (idx, val) = self.u_cols.line(q);
for (&i, &u) in idx.iter().zip(val) {
y[i as usize] -= u * x_q;
}
}
ws.list = list;
}
fn btran_hyper<U: BasisUpdate<T>>(
&self,
upd: &U,
rhs: &mut SparseVec<T>,
ws: &mut Workspace<T>,
) {
let mut list = std::mem::take(&mut ws.list);
list.clear();
for &q in rhs.pattern().unwrap_or(&[]) {
let v = rhs.get(q);
if !is_exact_zero(v) {
ws.work[q as usize] = v;
list.push(q);
}
}
rhs.clear();
let stamp = ws.new_stamp();
ws.post.clear();
reach(
&mut ws.mark,
stamp,
&mut ws.stack,
&list,
&mut ws.post,
|q| self.u_rows.line(self.p_of(q)).0,
);
list.clear();
for &q in ws.post.iter().rev() {
let q = q as usize;
let p = self.p_of(q);
let w_p = ws.work[q] / self.diag[p];
ws.work[q] = T::ZERO;
ws.work2[p] = w_p;
list.push(idx32(p));
if is_exact_zero(w_p) {
continue;
}
let (idx, val) = self.u_rows.line(p);
for (&j, &u) in idx.iter().zip(val) {
ws.work[j as usize] -= u * w_p;
}
}
let stamp = ws.new_stamp();
for &p in &list {
ws.mark[p as usize] = stamp;
}
upd.apply_btran_hyper(
&mut ws.work2,
&mut list,
&mut ws.mark,
ws.stamp,
&mut ws.queued,
);
let stamp = ws.new_stamp();
ws.post.clear();
reach(
&mut ws.mark,
stamp,
&mut ws.stack,
&list,
&mut ws.post,
|i| &self.lt_pivot[self.lt_start[i]..self.lt_start[i + 1]],
);
let w = &mut ws.work2;
for &i in ws.post.iter().rev() {
let i = i as usize;
let w_i = w[i];
if is_exact_zero(w_i) {
continue;
}
let r = self.lt_start[i]..self.lt_start[i + 1];
for (&p, &l) in self.lt_pivot[r.clone()].iter().zip(&self.lt_value[r]) {
w[p as usize] -= l * w_i;
}
}
for &i in ws.post.iter().rev() {
rhs.set(i, w[i as usize]);
w[i as usize] = T::ZERO;
}
ws.list = list;
}
}