Adom Fields

//! A multi-layer grid conductance system solved by preconditioned conjugate gradients.
//! Unknowns: one value per cell per layer. Edges: right and down neighbours on a layer (gx, gy),
//! the cell below on the next layer (gz), and a leak from every cell to ground.

#[derive(Clone)]
pub struct System { pub nl: usize, pub h: usize, pub w: usize, pub gx: Vec<Vec<f64>>, pub gy: Vec<Vec<f64>>, pub gz: Vec<Vec<f64>>, pub leak: Vec<Vec<f64>>, pub extra_leak: Vec<f64>, pub edges: Vec<(usize, usize, f64)> }

impl System {
    pub fn new(nl: usize, h: usize, w: usize) -> System {
        System { nl, h, w, extra_leak: Vec::new(), edges: Vec::new(), gx: vec![vec![0.0; h * (w - 1)]; nl], gy: vec![vec![0.0; (h - 1) * w]; nl], gz: vec![vec![0.0; h * w]; nl.saturating_sub(1)], leak: vec![vec![0.0; h * w]; nl] }
    }
    pub fn n(&self) -> usize { self.nl * self.h * self.w + self.extra_leak.len() }
    pub fn add_node(&mut self, leak: f64) -> usize { let i = self.n(); self.extra_leak.push(leak); i }
    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        let (h, w) = (self.h, self.w); let n = h * w;
        for k in 0..self.nl {
            let (xo, yo) = (&x[k * n..(k + 1) * n], &mut y[k * n..(k + 1) * n]);
            for i in 0..n { yo[i] = self.leak[k][i] * xo[i]; }
            let gx = &self.gx[k];
            for r in 0..h { for c in 0..w - 1 { let g = gx[r * (w - 1) + c]; if g != 0.0 { let (a, b) = (r * w + c, r * w + c + 1); let d = g * (xo[a] - xo[b]); yo[a] += d; yo[b] -= d; } } }
            let gy = &self.gy[k];
            for r in 0..h - 1 { for c in 0..w { let g = gy[r * w + c]; if g != 0.0 { let (a, b) = (r * w + c, (r + 1) * w + c); let d = g * (xo[a] - xo[b]); yo[a] += d; yo[b] -= d; } } }
        }
        for k in 0..self.nl.saturating_sub(1) {
            let gz = &self.gz[k];
            for i in 0..n { let g = gz[i]; if g != 0.0 { let (a, b) = (k * n + i, (k + 1) * n + i); let d = g * (x[a] - x[b]); y[a] += d; y[b] -= d; } }
        }
        for (i, g) in self.extra_leak.iter().enumerate() { let a = self.nl * n + i; y[a] = g * x[a]; }
        for &(a,b,g) in &self.edges { let q = g * (x[a]-x[b]); y[a] += q; y[b] -= q; }
    }
    fn diagonal(&self) -> Vec<f64> {
        let (h, w) = (self.h, self.w); let n = h * w;
        let mut d = vec![0.0; self.n()];
        for k in 0..self.nl {
            for i in 0..n { d[k * n + i] = self.leak[k][i]; }
            for r in 0..h { for c in 0..w - 1 { let g = self.gx[k][r * (w - 1) + c]; d[k * n + r * w + c] += g; d[k * n + r * w + c + 1] += g; } }
            for r in 0..h - 1 { for c in 0..w { let g = self.gy[k][r * w + c]; d[k * n + r * w + c] += g; d[k * n + (r + 1) * w + c] += g; } }
        }
        for k in 0..self.nl.saturating_sub(1) { for i in 0..n { let g = self.gz[k][i]; d[k * n + i] += g; d[(k + 1) * n + i] += g; } }
        for (i,g) in self.extra_leak.iter().enumerate() { d[self.nl*n+i] += g; }
        for &(a,b,g) in &self.edges { d[a] += g; d[b] += g; }
        d
    }
    /// Solve A x = b to a TRUE relative residual of `tol`: CG, then iterative refinement (A dx = b - Ax) while the
    /// recomputed residual is over tol, because CG's own recursive residual drifts on badly conditioned nets (issue 11).
    /// Returns (x, iterations, true residual).
    pub fn solve(&self, b: &[f64], tol: f64, max_iter: usize) -> (Vec<f64>, usize, f64) {
        let (mut x, mut its, _) = self.cg(b, tol, max_iter);
        let bnorm = b.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-30);
        let mut ax = vec![0.0; x.len()];
        let mut true_res = 0.0;
        for _ in 0..6 {
            self.matvec(&x, &mut ax);
            let r: Vec<f64> = b.iter().zip(&ax).map(|(bi, ai)| bi - ai).collect();
            true_res = r.iter().map(|v| v * v).sum::<f64>().sqrt() / bnorm;
            if true_res < tol || its >= max_iter { break; }
            let (dx, it2, _) = self.cg(&r, (tol * bnorm / r.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-300)).min(0.5), max_iter - its);
            for (xi, d) in x.iter_mut().zip(&dx) { *xi += d; }
            its += it2;
        }
        (x, its, true_res)
    }
    /// The same solve started from a guess (the still-air temperatures for the moving-air pass): only the correction
    /// A·dx = b − A·x0 is solved, to the same true relative residual against b. Returns (x, iterations, true residual).
    pub fn solve_from(&self, b: &[f64], x0: &[f64], tol: f64, max_iter: usize) -> (Vec<f64>, usize, f64) {
        if x0.len() != b.len() { return self.solve(b, tol, max_iter); }
        let bnorm = b.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-30);
        let mut ax = vec![0.0; b.len()]; self.matvec(x0, &mut ax);
        let r: Vec<f64> = b.iter().zip(&ax).map(|(bi, ai)| bi - ai).collect();
        let rn = r.iter().map(|v| v * v).sum::<f64>().sqrt();
        if rn / bnorm < tol { return (x0.to_vec(), 0, rn / bnorm); }
        let (dx, its, _) = self.solve(&r, (tol * bnorm / rn).min(0.5), max_iter);
        let x: Vec<f64> = x0.iter().zip(&dx).map(|(a, d)| a + d).collect();
        self.matvec(&x, &mut ax);
        let res = b.iter().zip(&ax).map(|(bi, ai)| (bi - ai) * (bi - ai)).sum::<f64>().sqrt() / bnorm;
        (x, its, res)
    }
    fn cg(&self, b: &[f64], tol: f64, max_iter: usize) -> (Vec<f64>, usize, f64) {
        let n = self.n();
        let diag = self.diagonal();
        let minv: Vec<f64> = diag.iter().map(|d| if *d > 0.0 { 1.0 / d } else { 0.0 }).collect();
        let mut x = vec![0.0; n];
        let mut r = b.to_vec();
        let mut z: Vec<f64> = r.iter().zip(&minv).map(|(a, m)| a * m).collect();
        let mut p = z.clone();
        let mut ap = vec![0.0; n];
        let bnorm = b.iter().map(|v| v * v).sum::<f64>().sqrt().max(1e-30);
        let mut rz: f64 = r.iter().zip(&z).map(|(a, b)| a * b).sum();
        let mut it = 0; let mut res = if b.iter().all(|v| *v == 0.0) { 0.0 } else { 1.0 };
        while it < max_iter {
            self.matvec(&p, &mut ap);
            let pap: f64 = p.iter().zip(&ap).map(|(a, b)| a * b).sum();
            if pap.abs() < 1e-300 { break; }
            let alpha = rz / pap;
            for i in 0..n { x[i] += alpha * p[i]; r[i] -= alpha * ap[i]; }
            res = r.iter().map(|v| v * v).sum::<f64>().sqrt() / bnorm;
            it += 1;
            if res < tol { break; }
            for i in 0..n { z[i] = r[i] * minv[i]; }
            let rz2: f64 = r.iter().zip(&z).map(|(a, b)| a * b).sum();
            let beta = rz2 / rz; rz = rz2;
            for i in 0..n { p[i] = z[i] + beta * p[i]; }
        }
        (x, it, res)
    }
}

pub fn harm(a: f64, b: f64) -> f64 { if a + b > 0.0 { 2.0 * a * b / (a + b) } else { 0.0 } }