use crate::compute::ComputeBackend;
use crate::ldu_matrix::parallel::vecop_backend_for;
use crate::ldu_matrix::LduMatrix;
#[cfg(feature = "parallel")]
use crate::ldu_matrix::parallel::CELL_BLOCK;
#[cfg(feature = "parallel")]
use rayon::prelude::*;
const ILU_PIVOT_FLOOR: f64 = 1.0e-300;
pub struct JacobiPreconditioner {
inv_diag: Vec<f64>,
}
impl JacobiPreconditioner {
pub fn new(a: &LduMatrix) -> Self {
let inv_diag = a
.diag
.iter()
.map(|&d| {
if d != 0.0 && d.is_finite() {
1.0 / d
} else {
1.0
}
})
.collect();
Self { inv_diag }
}
pub fn apply_on(&self, r: &[f64], z: &mut [f64], backend: ComputeBackend) {
debug_assert_eq!(r.len(), self.inv_diag.len());
debug_assert_eq!(z.len(), self.inv_diag.len());
match vecop_backend_for(backend, z.len()) {
#[cfg(feature = "parallel")]
ComputeBackend::CpuMulti => {
z.par_chunks_mut(CELL_BLOCK)
.zip(r.par_chunks(CELL_BLOCK))
.zip(self.inv_diag.par_chunks(CELL_BLOCK))
.for_each(|((zc, rc), dc)| {
for ((zi, ri), di) in zc.iter_mut().zip(rc.iter()).zip(dc.iter()) {
*zi = ri * di;
}
});
}
_ => {
for i in 0..z.len() {
z[i] = r[i] * self.inv_diag[i];
}
}
}
}
}
pub struct Ilu0Preconditioner {
n: usize,
row_ptr: Vec<usize>,
col_idx: Vec<usize>,
lu: Vec<f64>,
diag_ptr: Vec<usize>,
}
impl Ilu0Preconditioner {
pub fn new(a: &LduMatrix) -> Self {
let n = a.n_cells;
let mut row_count = vec![1usize; n]; for f in 0..a.n_internal_faces {
let o = a.owner[f];
let nb = a.neighbour[f];
row_count[o] += 1;
row_count[nb] += 1;
}
let mut row_ptr = vec![0usize; n + 1];
for i in 0..n {
row_ptr[i + 1] = row_ptr[i] + row_count[i];
}
let nnz = row_ptr[n];
let mut rows: Vec<Vec<(usize, f64)>> = vec![Vec::new(); n];
for i in 0..n {
rows[i].push((i, a.diag[i])); }
for f in 0..a.n_internal_faces {
let o = a.owner[f];
let nb = a.neighbour[f];
rows[o].push((nb, a.upper[f]));
rows[nb].push((o, a.lower[f]));
}
let mut col_idx = Vec::with_capacity(nnz);
let mut lu = Vec::with_capacity(nnz);
let mut diag_ptr = vec![0usize; n];
for i in 0..n {
rows[i].sort_by_key(|&(c, _)| c);
for &(c, v) in &rows[i] {
if c == i {
diag_ptr[i] = col_idx.len();
}
col_idx.push(c);
lu.push(v);
}
}
let mut me = Self {
n,
row_ptr,
col_idx,
lu,
diag_ptr,
};
me.factor();
me
}
fn find(&self, row: usize, col: usize) -> Option<usize> {
let lo = self.row_ptr[row];
let hi = self.row_ptr[row + 1];
let slice = &self.col_idx[lo..hi];
match slice.binary_search(&col) {
Ok(rel) => Some(lo + rel),
Err(_) => None,
}
}
fn factor(&mut self) {
for i in 0..self.n {
let ri = self.row_ptr[i];
let ri_end = self.row_ptr[i + 1];
for p in ri..ri_end {
let k = self.col_idx[p];
if k >= i {
break; }
let mut akk = self.lu[self.diag_ptr[k]];
if akk.abs() < ILU_PIVOT_FLOOR {
akk = if akk < 0.0 {
-ILU_PIVOT_FLOOR
} else {
ILU_PIVOT_FLOOR
};
}
let aik = self.lu[p] / akk;
self.lu[p] = aik;
for q in (p + 1)..ri_end {
let j = self.col_idx[q];
if let Some(kj) = self.find(k, j) {
let akj = self.lu[kj];
self.lu[q] -= aik * akj;
}
}
}
let dptr = self.diag_ptr[i];
if self.lu[dptr].abs() < ILU_PIVOT_FLOOR {
self.lu[dptr] = if self.lu[dptr] < 0.0 {
-ILU_PIVOT_FLOOR
} else {
ILU_PIVOT_FLOOR
};
}
}
}
pub fn apply_on(&self, r: &[f64], z: &mut [f64], _backend: ComputeBackend) {
self.apply(r, z);
}
fn apply(&self, r: &[f64], z: &mut [f64]) {
debug_assert_eq!(r.len(), self.n);
debug_assert_eq!(z.len(), self.n);
for i in 0..self.n {
let mut sum = r[i];
let ri = self.row_ptr[i];
let ri_end = self.row_ptr[i + 1];
for p in ri..ri_end {
let j = self.col_idx[p];
if j >= i {
break; }
sum -= self.lu[p] * z[j];
}
z[i] = sum;
}
for i in (0..self.n).rev() {
let mut sum = z[i];
let ri_end = self.row_ptr[i + 1];
let dptr = self.diag_ptr[i];
for p in (dptr + 1)..ri_end {
let j = self.col_idx[p];
sum -= self.lu[p] * z[j];
}
z[i] = sum / self.lu[dptr];
}
}
}
pub(crate) fn jacobi_apply(
p: &JacobiPreconditioner,
r: &[f64],
z: &mut [f64],
backend: ComputeBackend,
) {
p.apply_on(r, z, backend);
}
pub(crate) fn ilu0_apply(
p: &Ilu0Preconditioner,
r: &[f64],
z: &mut [f64],
backend: ComputeBackend,
) {
p.apply_on(r, z, backend);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ldu_matrix::LduMatrix;
fn tridiag(n: usize, d: f64, e: f64) -> LduMatrix {
let mut owner = Vec::new();
let mut neigh = Vec::new();
for i in 0..n - 1 {
owner.push(i);
neigh.push(i + 1);
}
let mut a = LduMatrix::new(n, owner, neigh);
for i in 0..n {
a.diag[i] = d;
}
for f in 0..a.n_internal_faces {
a.lower[f] = e;
a.upper[f] = e;
}
a
}
#[test]
fn jacobi_reciprocal_diag() {
let a = tridiag(4, 2.0, -1.0);
let p = JacobiPreconditioner::new(&a);
let r = [2.0, 4.0, 6.0, 8.0];
let mut z = [0.0; 4];
p.apply_on(&r, &mut z, ComputeBackend::Serial);
assert_eq!(z, [1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn jacobi_zero_diag_guard() {
let mut a = tridiag(2, 0.0, -1.0);
a.diag[0] = 0.0;
let p = JacobiPreconditioner::new(&a);
let r = [5.0, 3.0];
let mut z = [0.0; 2];
p.apply_on(&r, &mut z, ComputeBackend::Serial);
assert_eq!(z[0], 5.0); }
#[test]
fn ilu0_exact_for_tridiagonal() {
let n = 5;
let a = tridiag(n, 4.0, -1.0);
let ilu = Ilu0Preconditioner::new(&a);
let x_true: Vec<f64> = (0..n).map(|i| (i as f64 + 1.0) * 0.5).collect();
let b = a.multiply(&x_true);
let mut z = vec![0.0; n];
ilu.apply(&b, &mut z);
for i in 0..n {
assert!(
(z[i] - x_true[i]).abs() < 1e-10,
"row {i}: {} vs {}",
z[i],
x_true[i]
);
}
}
#[test]
fn ilu0_nonsymmetric_finite() {
let mut a = tridiag(4, 5.0, 0.0);
for f in 0..a.n_internal_faces {
a.lower[f] = -1.0;
a.upper[f] = -2.0;
}
let ilu = Ilu0Preconditioner::new(&a);
let r = [1.0, 2.0, 3.0, 4.0];
let mut z = [0.0; 4];
ilu.apply(&r, &mut z);
assert!(z.iter().all(|v| v.is_finite()));
}
}