use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::fmt::Debug;
use crate::error::LuError;
use crate::factor::{Factors, LuOptions, VACANT, count_f64, idx32};
use crate::real::{Real, is_exact_zero, is_zero, lt};
use crate::solve::Workspace;
use crate::{Row, SparseError};
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum RefactorReason {
UpdateLimit,
FillGrowth,
Growth,
Unstable,
}
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum UpdateOutcome {
Updated,
RefactorNeeded(RefactorReason),
}
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct UpdateStats {
pub updates: usize,
pub nnz_eta: usize,
pub n_etas: usize,
pub update_growth: f64,
}
pub trait BasisUpdate<T: Real>: Default + Clone + Debug {
fn update(
&mut self,
factors: &mut Factors<T>,
opts: &LuOptions,
r: usize,
alpha_r: T,
) -> Result<UpdateOutcome, LuError>;
fn apply_ftran(&self, y: &mut [T]);
fn apply_btran(&self, w: &mut [T]);
fn apply_ftran_hyper(
&self,
y: &mut [T],
list: &mut Vec<u32>,
mark: &mut [u32],
stamp: u32,
queued: &mut [bool],
);
fn apply_btran_hyper(
&self,
w: &mut [T],
list: &mut Vec<u32>,
mark: &mut [u32],
stamp: u32,
queued: &mut [bool],
);
fn reset(&mut self, m: usize, update_growth: f64);
fn stats(&self) -> &UpdateStats;
fn name(&self) -> &'static str;
}
#[derive(Clone, Debug)]
pub struct ForrestTomlin<T: Real> {
eta_start: Vec<usize>,
eta_target: Vec<u32>,
eta_index: Vec<u32>,
eta_value: Vec<T>,
eta_readers: Vec<Vec<u32>>,
eta_targets: Vec<Vec<u32>>,
eta_rows: Vec<u32>,
stats: UpdateStats,
}
impl<T: Real> Default for ForrestTomlin<T> {
fn default() -> Self {
Self {
eta_start: vec![0],
eta_target: Vec::new(),
eta_index: Vec::new(),
eta_value: Vec::new(),
eta_readers: Vec::new(),
eta_targets: Vec::new(),
eta_rows: Vec::new(),
stats: UpdateStats::default(),
}
}
}
impl<T: Real> BasisUpdate<T> for ForrestTomlin<T> {
fn update(
&mut self,
f: &mut Factors<T>,
opts: &LuOptions,
r: usize,
alpha_r: T,
) -> Result<UpdateOutcome, LuError> {
if !f.valid {
return Err(LuError::NotFactorized);
}
if r >= f.m {
return Err(SparseError::IndexOutOfRange {
what: "leaving basis position",
index: r,
bound: f.m,
}
.into());
}
if !f.spike_valid {
return Err(LuError::NoSpike);
}
let tau = f.col_step[r];
let p = f.step_row[tau] as usize;
let u_old = f.diag[p];
let mut ws = std::mem::take(&mut f.ws);
let (eta, d, zeta) = eliminate_spike_row(f, opts, tau, p, &mut ws);
f.ws = ws;
if !lt(T::from_f64(opts.abs_pivot_tol) * zeta.max(T::ONE), d.abs()) {
return Err(LuError::Singular {
rank_deficiency: 1,
cols: vec![r],
rows: vec![Row::from_usize(p)],
});
}
let expected = alpha_r * u_old;
let scale = d.abs().max(expected.abs());
if lt(
T::from_f64(opts.update_stability_tol) * scale,
(d - expected).abs(),
) {
return Ok(UpdateOutcome::RefactorNeeded(RefactorReason::Unstable));
}
self.commit(f, opts, r, tau, p, d, eta);
debug_assert!(f.check_invariants().is_ok(), "{:?}", f.check_invariants());
Ok(self.policy(f, opts))
}
#[inline]
fn apply_ftran(&self, y: &mut [T]) {
for e in 0..self.eta_target.len() {
self.apply_eta(e, y);
}
}
#[inline]
fn apply_btran(&self, w: &mut [T]) {
for e in (0..self.eta_target.len()).rev() {
self.apply_eta_t(e, w);
}
}
fn apply_ftran_hyper(
&self,
y: &mut [T],
list: &mut Vec<u32>,
mark: &mut [u32],
stamp: u32,
queued: &mut [bool],
) {
if self.eta_target.is_empty() {
return;
}
let mut heap: BinaryHeap<Reverse<u32>> = BinaryHeap::new();
for &i in list.iter() {
for &e in &self.eta_readers[i as usize] {
if !queued[e as usize] {
queued[e as usize] = true;
heap.push(Reverse(e));
}
}
}
while let Some(Reverse(e)) = heap.pop() {
queued[e as usize] = false;
self.apply_eta(e as usize, y);
let t = self.eta_target[e as usize] as usize;
if mark[t] != stamp {
mark[t] = stamp;
list.push(idx32(t));
let readers = &self.eta_readers[t];
let after = readers.partition_point(|&e2| e2 <= e);
for &e2 in &readers[after..] {
if !queued[e2 as usize] {
queued[e2 as usize] = true;
heap.push(Reverse(e2));
}
}
}
}
}
fn apply_btran_hyper(
&self,
w: &mut [T],
list: &mut Vec<u32>,
mark: &mut [u32],
stamp: u32,
queued: &mut [bool],
) {
if self.eta_target.is_empty() {
return;
}
let mut heap: BinaryHeap<u32> = BinaryHeap::new();
for &p in list.iter() {
for &e in &self.eta_targets[p as usize] {
if !queued[e as usize] {
queued[e as usize] = true;
heap.push(e);
}
}
}
while let Some(e) = heap.pop() {
queued[e as usize] = false;
self.apply_eta_t(e as usize, w);
let r = self.eta_start[e as usize]..self.eta_start[e as usize + 1];
for &k in &self.eta_index[r] {
let k = k as usize;
if mark[k] == stamp {
continue;
}
mark[k] = stamp;
list.push(idx32(k));
let targets = &self.eta_targets[k];
let before = targets.partition_point(|&e2| e2 < e);
for &e2 in &targets[..before] {
if !queued[e2 as usize] {
queued[e2 as usize] = true;
heap.push(e2);
}
}
}
}
}
fn reset(&mut self, m: usize, update_growth: f64) {
self.eta_start = vec![0];
self.eta_target.clear();
self.eta_index.clear();
self.eta_value.clear();
if self.eta_readers.len() == m {
for &k in &self.eta_rows {
self.eta_readers[k as usize].clear();
self.eta_targets[k as usize].clear();
}
} else {
self.eta_readers = vec![Vec::new(); m];
self.eta_targets = vec![Vec::new(); m];
}
self.eta_rows.clear();
self.stats = UpdateStats {
update_growth,
..UpdateStats::default()
};
}
#[inline]
fn stats(&self) -> &UpdateStats {
&self.stats
}
fn name(&self) -> &'static str {
"forrest-tomlin"
}
}
impl<T: Real> ForrestTomlin<T> {
#[inline]
fn apply_eta(&self, e: usize, y: &mut [T]) {
let r = self.eta_start[e]..self.eta_start[e + 1];
let s: T = self.eta_index[r.clone()]
.iter()
.zip(&self.eta_value[r])
.map(|(&k, &mu)| mu * y[k as usize])
.sum();
y[self.eta_target[e] as usize] -= s;
}
#[inline]
fn apply_eta_t(&self, e: usize, w: &mut [T]) {
let w_t = w[self.eta_target[e] as usize];
if is_exact_zero(w_t) {
return;
}
let r = self.eta_start[e]..self.eta_start[e + 1];
for (&k, &mu) in self.eta_index[r.clone()].iter().zip(&self.eta_value[r]) {
w[k as usize] -= mu * w_t;
}
}
#[allow(clippy::too_many_arguments)] fn commit(
&mut self,
f: &mut Factors<T>,
opts: &LuOptions,
r: usize,
tau: usize,
p: usize,
d: T,
eta: Vec<(u32, T)>,
) {
let drop_tol = T::from_f64(opts.drop_tol);
let (r32, p32) = (idx32(r), idx32(p));
let old_col: Vec<u32> = f.u_cols.line(r).0.to_vec();
for i in old_col {
f.u_rows.remove(i as usize, r32);
}
f.u_cols.clear(r);
let old_row: Vec<u32> = f.u_rows.line(p).0.to_vec();
for j in old_row {
f.u_cols.remove(j as usize, p32);
}
f.u_rows.clear(p);
let mut max_u = f.stats.max_u.max(d.abs().to_f64());
for &i in &f.spike_idx {
let (iu, s_i) = (i as usize, f.spike[i as usize]);
if iu == p || is_zero(s_i, drop_tol) {
continue;
}
f.u_cols.push(r, i, s_i);
f.u_rows.push(iu, r32, s_i);
max_u = max_u.max(s_i.abs().to_f64());
}
f.diag[p] = d;
f.step_row[tau] = VACANT;
f.step_col[tau] = VACANT;
let last = f.step_row.len();
f.step_row.push(p32);
f.step_col.push(r32);
f.row_step[p] = last;
f.col_step[r] = last;
if !eta.is_empty() {
let e = idx32(self.eta_target.len());
self.eta_target.push(p32);
self.eta_targets[p].push(e);
self.eta_rows.push(p32);
for (k, mu) in eta {
self.eta_index.push(k);
self.eta_value.push(mu);
self.eta_readers[k as usize].push(e);
self.eta_rows.push(k);
}
self.eta_start.push(self.eta_index.len());
f.ws.queued.push(false);
}
f.spike_valid = false;
f.stats.max_u = max_u;
self.stats.updates += 1;
self.stats.nnz_eta = self.eta_value.len();
self.stats.n_etas = self.eta_target.len();
self.stats.update_growth = f.refresh_stats(opts.growth_limit, self.stats.nnz_eta);
}
fn policy(&self, f: &Factors<T>, opts: &LuOptions) -> UpdateOutcome {
let s = &f.stats;
let u = &self.stats;
let nnz_now = s.nnz_l + s.nnz_u + u.nnz_eta;
if u.updates >= opts.refactor_interval {
UpdateOutcome::RefactorNeeded(RefactorReason::UpdateLimit)
} else if lt(
opts.fill_factor * count_f64(f.nnz_factor0),
count_f64(nnz_now),
) {
UpdateOutcome::RefactorNeeded(RefactorReason::FillGrowth)
} else if lt(opts.growth_limit, u.update_growth) {
UpdateOutcome::RefactorNeeded(RefactorReason::Growth)
} else {
UpdateOutcome::Updated
}
}
}
fn eliminate_spike_row<T: Real>(
f: &Factors<T>,
opts: &LuOptions,
tau: usize,
p: usize,
ws: &mut Workspace<T>,
) -> (Vec<(u32, T)>, T, T) {
let drop_tol = T::from_f64(opts.drop_tol);
let w = &mut ws.work;
let flag = &mut ws.flag;
let mut heap: BinaryHeap<Reverse<usize>> = BinaryHeap::new();
let (idx, val) = f.u_rows.line(p);
for (&j, &u) in idx.iter().zip(val) {
w[j as usize] = u;
flag[j as usize] = true;
heap.push(Reverse(f.col_step[j as usize]));
}
let mut sigma = f.spike[p];
let mut zeta = sigma.abs();
let mut eta: Vec<(u32, T)> = Vec::new();
while let Some(Reverse(k)) = heap.pop() {
debug_assert!(k > tau);
let j = f.step_col[k] as usize;
let w_j = w[j];
w[j] = T::ZERO;
flag[j] = false;
if is_zero(w_j, drop_tol) {
continue;
}
let p_k = f.step_row[k] as usize;
let mu = w_j / f.diag[p_k];
let (idx, val) = f.u_rows.line(p_k);
for (&l, &u) in idx.iter().zip(val) {
let l = l as usize;
if !flag[l] {
flag[l] = true;
heap.push(Reverse(f.col_step[l]));
}
w[l] -= mu * u;
}
sigma -= mu * f.spike[p_k];
zeta += (mu * f.spike[p_k]).abs();
eta.push((idx32(p_k), mu));
}
(eta, sigma, zeta)
}