petektools 0.0.2

//! Briggs minimum-curvature gridding via biharmonic (∇⁴z = 0) SOR relaxation.
//!
//! Ported from petekio 0.2.0 `src/core/gridding.rs` (`grid_min_curvature` +
//! `biharmonic_target` / `harmonic_target`, the author's own prior art) with the
//! `GridGeometry` parameter swapped 1:1 for [`Lattice`]. The algorithm, defaults
//! (ω = 1.5, TOL = 1e-6, MAX_ITERS = 5000) and tolerances are preserved.
//!
//! This is *one-shot* gridding. The **convergent / warm-start** variant
//! (re-grid only the region affected by a newly-added control point) is a
//! planned extension — see `dev-docs/designs/roadmap.md`.

use crate::foundation::Lattice;
use ndarray::Array2;

use super::idw::grid_idw;

/// Briggs minimum-curvature gridding via biharmonic (∇⁴z = 0) SOR relaxation,
/// anchored at the grid nodes nearest each data sample.
///
/// **Scope:** a straightforward, convergent implementation for small/moderate
/// grids — interior nodes use the 13-point biharmonic stencil; near-edge nodes
/// fall back to the 5-point harmonic (Laplacian) update. Data are honoured by
/// snapping each sample to its nearest node and holding those fixed. Linear
/// trends (the exact biharmonic solution) are reproduced to tolerance.
///
/// When `seed` is `Some` and matches the lattice shape `(ncol, nrow)`, the SOR
/// relaxation starts from it (a **warm start**) instead of the cold IDW seed —
/// for an incremental re-grid this converges in far fewer iterations while
/// reaching the same field (warm == cold to tolerance). `None` or a wrong-shape
/// seed falls back to the cold IDW seed. The whole field is re-solved either
/// way; the seed only changes the starting point.
pub(crate) fn grid_min_curvature(
    coords: &[[f64; 3]],
    lattice: &Lattice,
    seed: Option<&Array2<f64>>,
) -> Array2<f64> {
    let (nc, nr) = (lattice.ncol, lattice.nrow);
    // Warm-start from `seed` when shape-matched, else cold IDW seed (a smooth,
    // in-range field). Parity with petekio's `grid_min_curvature(.., seed)`.
    let mut z = match seed {
        Some(s) if s.dim() == (nc, nr) => s.clone(),
        _ => grid_idw(coords, lattice),
    };

    // Anchor nodes: snap each sample to the nearest node, averaging collisions.
    let mut fixed = Array2::from_elem((nc, nr), false);
    let mut acc: std::collections::HashMap<(usize, usize), (f64, usize)> =
        std::collections::HashMap::new();
    for c in coords {
        if let Some((fi, fj)) = lattice.xy_to_ij(c[0], c[1]) {
            let i = fi.round();
            let j = fj.round();
            if i < 0.0 || j < 0.0 {
                continue;
            }
            let (i, j) = (i as usize, j as usize);
            if i < nc && j < nr {
                let e = acc.entry((i, j)).or_insert((0.0, 0));
                e.0 += c[2];
                e.1 += 1;
            }
        }
    }
    for ((i, j), (sum, n)) in acc {
        z[[i, j]] = sum / n as f64;
        fixed[[i, j]] = true;
    }

    if nc < 2 || nr < 2 {
        return z;
    }

    const MAX_ITERS: usize = 5000;
    const TOL: f64 = 1e-6;
    const OMEGA: f64 = 1.5; // SOR over-relaxation

    let range = {
        let (mut lo, mut hi) = (f64::INFINITY, f64::NEG_INFINITY);
        for c in coords {
            lo = lo.min(c[2]);
            hi = hi.max(c[2]);
        }
        (hi - lo).abs().max(1.0)
    };

    for _ in 0..MAX_ITERS {
        let mut max_delta = 0.0_f64;
        for j in 0..nr {
            for i in 0..nc {
                if fixed[[i, j]] {
                    continue;
                }
                let target = if i >= 2 && i + 2 < nc && j >= 2 && j + 2 < nr {
                    biharmonic_target(&z, i, j)
                } else {
                    harmonic_target(&z, i, j, nc, nr)
                };
                let old = z[[i, j]];
                let new = old + OMEGA * (target - old);
                z[[i, j]] = new;
                max_delta = max_delta.max((new - old).abs());
            }
        }
        if max_delta < TOL * range {
            break;
        }
    }
    z
}

/// 13-point biharmonic stencil solved for the centre node (∇⁴z = 0):
/// `z = [8·(N+S+E+W) − 2·(NE+NW+SE+SW) − (NN+SS+EE+WW)] / 20`.
fn biharmonic_target(z: &Array2<f64>, i: usize, j: usize) -> f64 {
    let n = z[[i, j + 1]];
    let s = z[[i, j - 1]];
    let e = z[[i + 1, j]];
    let w = z[[i - 1, j]];
    let ne = z[[i + 1, j + 1]];
    let nw = z[[i - 1, j + 1]];
    let se = z[[i + 1, j - 1]];
    let sw = z[[i - 1, j - 1]];
    let nn = z[[i, j + 2]];
    let ss = z[[i, j - 2]];
    let ee = z[[i + 2, j]];
    let ww = z[[i - 2, j]];
    (8.0 * (n + s + e + w) - 2.0 * (ne + nw + se + sw) - (nn + ss + ee + ww)) / 20.0
}

/// 5-point harmonic (Laplacian) update, averaging available orthogonal
/// neighbours — used for near-edge nodes where the biharmonic stencil overruns.
fn harmonic_target(z: &Array2<f64>, i: usize, j: usize, nc: usize, nr: usize) -> f64 {
    let mut sum = 0.0;
    let mut n = 0.0;
    if i > 0 {
        sum += z[[i - 1, j]];
        n += 1.0;
    }
    if i + 1 < nc {
        sum += z[[i + 1, j]];
        n += 1.0;
    }
    if j > 0 {
        sum += z[[i, j - 1]];
        n += 1.0;
    }
    if j + 1 < nr {
        sum += z[[i, j + 1]];
        n += 1.0;
    }
    if n > 0.0 {
        sum / n
    } else {
        z[[i, j]]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A planar linear trend `z = a·x + b·y + c` is the exact biharmonic
    /// solution (∇⁴z = 0), so minimum curvature must reproduce it to tolerance
    /// when samples define the plane.
    #[test]
    fn linear_trend_reproduced_exactly() {
        let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 7, 7);
        let (a, b, c) = (2.0, -3.0, 5.0);

        // Sample every node so the plane is fully constrained; the relaxation
        // then has to hold the exact plane everywhere.
        let mut coords = Vec::new();
        for j in 0..7 {
            for i in 0..7 {
                let (x, y) = lattice.node_xy(i, j);
                coords.push([x, y, a * x + b * y + c]);
            }
        }

        let out = grid_min_curvature(&coords, &lattice, None);
        for j in 0..7 {
            for i in 0..7 {
                let (x, y) = lattice.node_xy(i, j);
                let expected = a * x + b * y + c;
                assert!(
                    (out[[i, j]] - expected).abs() < 1e-6,
                    "node ({i},{j}): got {}, want {expected}",
                    out[[i, j]]
                );
            }
        }
    }

    /// With only sparse corner+edge samples on a plane, the interior is *solved*
    /// (not sampled) — the biharmonic relaxation should still recover the plane
    /// to a looser tolerance, since a plane is the minimum-curvature surface.
    #[test]
    fn sparse_linear_samples_recover_plane() {
        let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 9, 9);
        let (a, b, c) = (1.0, 0.5, -2.0);
        // Sample only the boundary ring.
        let mut coords = Vec::new();
        for j in 0..9 {
            for i in 0..9 {
                if i == 0 || i == 8 || j == 0 || j == 8 {
                    let (x, y) = lattice.node_xy(i, j);
                    coords.push([x, y, a * x + b * y + c]);
                }
            }
        }
        let out = grid_min_curvature(&coords, &lattice, None);
        // Interior node, fully solved by relaxation.
        let (x, y) = lattice.node_xy(4, 4);
        let expected = a * x + b * y + c;
        assert!(
            (out[[4, 4]] - expected).abs() < 1e-3,
            "interior node: got {}, want {expected}",
            out[[4, 4]]
        );
    }

    /// Samples are honoured: a node snapped to a sample holds that sample's Z.
    #[test]
    fn anchors_are_honoured() {
        let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 5, 5);
        let coords = [[0.0, 0.0, 0.0], [4.0, 4.0, 100.0], [2.0, 2.0, 50.0]];
        let out = grid_min_curvature(&coords, &lattice, None);
        assert!((out[[0, 0]] - 0.0).abs() < 1e-9);
        assert!((out[[4, 4]] - 100.0).abs() < 1e-9);
        assert!((out[[2, 2]] - 50.0).abs() < 1e-9);
    }

    /// Degenerate 1-D lattices short-circuit before relaxation (the `nc < 2 ||
    /// nr < 2` guard) and just return the snapped/IDW seed without panicking.
    #[test]
    fn degenerate_single_row_does_not_panic() {
        let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 5, 1);
        let coords = [[0.0, 0.0, 10.0], [4.0, 0.0, 20.0]];
        let out = grid_min_curvature(&coords, &lattice, None);
        assert_eq!(out.dim(), (5, 1));
        assert!(out.iter().all(|v| v.is_finite()));
    }

    /// Warm == cold to tolerance: seeding the relaxation with the cold solution
    /// must reach the same field (the seed only changes the starting point, not
    /// the fixed point the SOR converges to).
    #[test]
    fn warm_start_matches_cold_to_tolerance() {
        let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 12, 10);
        let coords = [
            [1.0, 1.0, 10.0],
            [9.0, 2.0, 25.0],
            [3.0, 8.0, 5.0],
            [10.0, 8.0, 40.0],
            [5.0, 5.0, 18.0],
        ];
        let cold = grid_min_curvature(&coords, &lattice, None);
        let warm = grid_min_curvature(&coords, &lattice, Some(&cold));
        // Both solves stop at `max_delta < TOL·range` (TOL = 1e-6), so they
        // agree to the solver's own tolerance — not to 1e-6 absolute. 1e-3 is
        // comfortably above TOL·range here yet ~four orders below the field
        // magnitude: a real continuity guarantee.
        for (w, c) in warm.iter().zip(cold.iter()) {
            assert!((w - c).abs() < 1e-3, "warm {w} vs cold {c}");
        }
    }

    /// A wrong-shape seed is ignored — falls back to the cold IDW seed, so the
    /// result is identical to a `None` (cold) solve.
    #[test]
    fn wrong_shape_seed_falls_back_to_cold() {
        let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 6, 6);
        let coords = [[0.0, 0.0, 0.0], [5.0, 5.0, 50.0], [2.0, 3.0, 20.0]];
        let cold = grid_min_curvature(&coords, &lattice, None);
        let bogus = Array2::from_elem((3, 3), 999.0); // wrong shape
        let out = grid_min_curvature(&coords, &lattice, Some(&bogus));
        for (o, c) in out.iter().zip(cold.iter()) {
            assert_eq!(o, c);
        }
    }
}