use crate::foundation::{AlgoError, Lattice, Result};
use ndarray::Array2;
use std::collections::HashMap;
use super::band_lu::BandLu;
use super::min_curvature::{classify_sample, SampleKind, DATA_WEIGHT, TENSION};
use super::Conditioning;
fn stencil() -> [(isize, isize, f64); 13] {
let t = TENSION;
let denom = 20.0 * (1.0 - t) + 4.0 * t;
let orth = -(8.0 * (1.0 - t) + t);
let diag = 2.0 * (1.0 - t);
let two = 1.0 - t;
[
(0, 0, denom),
(0, 1, orth),
(0, -1, orth),
(1, 0, orth),
(-1, 0, orth),
(1, 1, diag),
(-1, 1, diag),
(1, -1, diag),
(-1, -1, diag),
(0, 2, two),
(0, -2, two),
(2, 0, two),
(-2, 0, two),
]
}
fn accumulate<F: FnMut(usize, usize, f64)>(
nc: isize,
nr: isize,
i: isize,
j: isize,
coeff: f64,
out: &mut F,
) {
if i < 0 {
accumulate(nc, nr, i + 1, j, 2.0 * coeff, out);
accumulate(nc, nr, i + 2, j, -coeff, out);
} else if i >= nc {
accumulate(nc, nr, i - 1, j, 2.0 * coeff, out);
accumulate(nc, nr, i - 2, j, -coeff, out);
} else if j < 0 {
accumulate(nc, nr, i, j + 1, 2.0 * coeff, out);
accumulate(nc, nr, i, j + 2, -coeff, out);
} else if j >= nr {
accumulate(nc, nr, i, j - 1, 2.0 * coeff, out);
accumulate(nc, nr, i, j - 2, -coeff, out);
} else {
out(i as usize, j as usize, coeff);
}
}
struct OffSample {
sample: u32,
nodes: [(usize, usize); 4],
w: [f64; 4],
}
struct Geometry {
nc: usize,
anchor_of: Vec<i32>,
anchor_node: Vec<(usize, usize)>,
anchor_samples: Vec<Vec<u32>>,
off: Vec<OffSample>,
free_of: Vec<i32>,
free_node: Vec<(usize, usize)>,
}
fn classify_geometry(
lattice: &Lattice,
sample_xy: &[[f64; 2]],
conditioning: Conditioning,
) -> Geometry {
let (nc, nr) = (lattice.ncol, lattice.nrow);
let n_nodes = nc * nr;
let mut anchor_of = vec![-1i32; n_nodes];
let mut anchor_node: Vec<(usize, usize)> = Vec::new();
let mut anchor_samples: Vec<Vec<u32>> = Vec::new();
let mut off: Vec<OffSample> = Vec::new();
for (s, xy) in sample_xy.iter().enumerate() {
match classify_sample(lattice, xy[0], xy[1], conditioning) {
SampleKind::OnNode(i, j) => {
let f = i + j * nc;
let a = if anchor_of[f] < 0 {
anchor_of[f] = anchor_node.len() as i32;
anchor_node.push((i, j));
anchor_samples.push(Vec::new());
anchor_node.len() - 1
} else {
anchor_of[f] as usize
};
anchor_samples[a].push(s as u32);
}
SampleKind::Off { nodes, w } => off.push(OffSample {
sample: s as u32,
nodes,
w,
}),
SampleKind::Skip => {}
}
}
let mut free_of = vec![-1i32; n_nodes];
let mut free_node: Vec<(usize, usize)> = Vec::new();
let push = |i: usize, j: usize, free_of: &mut [i32], fnv: &mut Vec<(usize, usize)>| {
let f = i + j * nc;
if anchor_of[f] < 0 {
free_of[f] = fnv.len() as i32;
fnv.push((i, j));
}
};
if nc <= nr {
for j in 0..nr {
for i in 0..nc {
push(i, j, &mut free_of, &mut free_node);
}
}
} else {
for i in 0..nc {
for j in 0..nr {
push(i, j, &mut free_of, &mut free_node);
}
}
}
Geometry {
nc,
anchor_of,
anchor_node,
anchor_samples,
off,
free_of,
free_node,
}
}
#[allow(clippy::type_complexity)]
fn assemble(
geom: &Geometry,
nc: isize,
nr: isize,
) -> (
HashMap<(u32, u32), f64>,
Vec<Vec<(u32, f64)>>,
Vec<Vec<(u32, f64)>>,
) {
let n_free = geom.free_node.len();
let mut full: HashMap<(u32, u32), f64> = HashMap::new();
let mut rhs_anchor: Vec<Vec<(u32, f64)>> = vec![Vec::new(); n_free];
let mut rhs_data: Vec<Vec<(u32, f64)>> = vec![Vec::new(); n_free];
let ncu = geom.nc;
let stencil = stencil();
for (r, &(i, j)) in geom.free_node.iter().enumerate() {
let r = r as u32;
for &(di, dj, coeff) in &stencil {
accumulate(
nc,
nr,
i as isize + di,
j as isize + dj,
coeff,
&mut |ni, nj, c| {
let f = ni + nj * ncu;
let a = geom.anchor_of[f];
if a >= 0 {
rhs_anchor[r as usize].push((a as u32, c));
} else {
let q = geom.free_of[f] as u32;
*full.entry((r, q)).or_insert(0.0) += c;
}
},
);
}
}
for s in &geom.off {
for k in 0..4 {
if s.w[k] == 0.0 {
continue;
}
let (ik, jk) = s.nodes[k];
let fk = ik + jk * ncu;
if geom.anchor_of[fk] >= 0 {
continue; }
let rk = geom.free_of[fk] as usize;
rhs_data[rk].push((s.sample, DATA_WEIGHT * s.w[k]));
for m in 0..4 {
if s.w[m] == 0.0 {
continue;
}
let (im, jm) = s.nodes[m];
let fm = im + jm * ncu;
let coeff = DATA_WEIGHT * s.w[k] * s.w[m];
if geom.anchor_of[fm] >= 0 {
rhs_anchor[rk].push((geom.anchor_of[fm] as u32, coeff));
} else {
let rm = geom.free_of[fm] as u32;
*full.entry((rk as u32, rm)).or_insert(0.0) += coeff;
}
}
}
}
(full, rhs_anchor, rhs_data)
}
const MIN_CONTROLS: usize = 4;
const RIDGE_REL: f64 = 1.0e-8;
fn build_and_factor(
full: &HashMap<(u32, u32), f64>,
n_free: usize,
bw: usize,
ridge: f64,
) -> Option<BandLu> {
let mut lu = BandLu::zeros(n_free, bw);
for (&(r, c), &v) in full {
lu.add(r as usize, c as usize, v);
}
if ridge != 0.0 {
for d in 0..n_free {
lu.add(d, d, ridge);
}
}
lu.factor().then_some(lu)
}
pub struct MinCurvatureOperator {
lattice: Lattice,
nc: usize,
nr: usize,
n_samples: usize,
lu: Option<BandLu>,
n_free: usize,
free_node: Vec<(usize, usize)>,
anchor_node: Vec<(usize, usize)>,
anchor_samples: Vec<Vec<u32>>,
rhs_anchor: Vec<Vec<(u32, f64)>>,
rhs_data: Vec<Vec<(u32, f64)>>,
solve_scratch: std::cell::RefCell<Vec<f64>>,
}
impl MinCurvatureOperator {
pub fn factor(
lattice: &Lattice,
sample_xy: &[[f64; 2]],
conditioning: Conditioning,
) -> Result<Self> {
let (nc, nr) = (lattice.ncol, lattice.nrow);
if nc < 2 || nr < 2 {
return Err(AlgoError::InvalidGeometry(
"MinCurvatureOperator::factor: lattice must be at least 2x2",
));
}
let geom = classify_geometry(lattice, sample_xy, conditioning);
let n_free = geom.free_node.len();
if n_free == 0 {
return Ok(MinCurvatureOperator {
lattice: lattice.clone(),
nc,
nr,
n_samples: sample_xy.len(),
lu: None,
n_free: 0,
free_node: geom.free_node,
anchor_node: geom.anchor_node,
anchor_samples: geom.anchor_samples,
rhs_anchor: Vec::new(),
rhs_data: Vec::new(),
solve_scratch: std::cell::RefCell::new(Vec::new()),
});
}
let (full, rhs_anchor, rhs_data) = assemble(&geom, nc as isize, nr as isize);
let mut bw = 0usize;
for &(r, c) in full.keys() {
bw = bw.max((r as isize - c as isize).unsigned_abs());
}
let lu = build_and_factor(&full, n_free, bw, 0.0).or_else(|| {
let n_controls = geom.anchor_node.len() + geom.off.len();
if n_controls < MIN_CONTROLS {
return None;
}
let scale = full
.iter()
.filter(|(&(r, c), _)| r == c)
.map(|(_, &v)| v.abs())
.fold(0.0_f64, f64::max);
build_and_factor(&full, n_free, bw, RIDGE_REL * scale)
});
let Some(lu) = lu else {
return Err(AlgoError::InvalidArgument(
"MinCurvatureOperator: conditioning system is singular \
(too few / collinear hard anchors to pin the surface)"
.to_string(),
));
};
Ok(MinCurvatureOperator {
lattice: lattice.clone(),
nc,
nr,
n_samples: sample_xy.len(),
lu: Some(lu),
n_free,
free_node: geom.free_node,
anchor_node: geom.anchor_node,
anchor_samples: geom.anchor_samples,
rhs_anchor,
rhs_data,
solve_scratch: std::cell::RefCell::new(Vec::new()),
})
}
pub fn lattice(&self) -> &Lattice {
&self.lattice
}
pub fn sample_count(&self) -> usize {
self.n_samples
}
pub fn solve(&self, z: &[f64]) -> Result<Array2<f64>> {
if z.len() != self.n_samples {
return Err(AlgoError::InvalidArgument(format!(
"MinCurvatureOperator::solve: expected {} z-values, got {}",
self.n_samples,
z.len()
)));
}
let anchorval: Vec<f64> = self
.anchor_samples
.iter()
.map(|ss| ss.iter().map(|&s| z[s as usize]).sum::<f64>() / ss.len() as f64)
.collect();
let mut field = Array2::from_elem((self.nc, self.nr), f64::NAN);
for (a, &(i, j)) in self.anchor_node.iter().enumerate() {
field[[i, j]] = anchorval[a];
}
if let Some(lu) = &self.lu {
let mut b = vec![0.0f64; self.n_free];
for (r, slot) in b.iter_mut().enumerate() {
let mut acc = 0.0;
for &(sid, coeff) in &self.rhs_data[r] {
acc += coeff * z[sid as usize];
}
for &(aid, coeff) in &self.rhs_anchor[r] {
acc -= coeff * anchorval[aid as usize];
}
*slot = acc;
}
let mut out = self.solve_scratch.borrow_mut();
lu.solve_into(&b, &mut out);
for (r, &(i, j)) in self.free_node.iter().enumerate() {
field[[i, j]] = out[r];
}
}
Ok(field)
}
#[cfg(test)]
pub(crate) fn anchor_count(&self) -> usize {
self.anchor_node.len()
}
#[cfg(test)]
pub(crate) fn free_count(&self) -> usize {
self.n_free
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn direct_matches_converged_sor_nonplanar() {
use super::super::min_curvature::grid_min_curvature_sor;
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 11, 9);
let mut coords = Vec::new();
for &(i, j) in &[
(0usize, 0usize),
(10, 0),
(0, 8),
(10, 8),
(5, 4),
(2, 6),
(8, 2),
(3, 3),
(7, 6),
] {
let (x, y) = (i as f64, j as f64);
let z = 100.0 + 3.0 * (x * 0.4).sin() - 2.0 * (y * 0.3).cos() + 0.05 * x * y;
coords.push([x, y, z]);
}
let sor = grid_min_curvature_sor(&coords, &lattice, None, Conditioning::NearestNode);
let sample_xy: Vec<[f64; 2]> = coords.iter().map(|c| [c[0], c[1]]).collect();
let z: Vec<f64> = coords.iter().map(|c| c[2]).collect();
let op =
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::NearestNode).unwrap();
let direct = op.solve(&z).unwrap();
let mut maxd = 0.0f64;
for (a, b) in direct.iter().zip(sor.iter()) {
maxd = maxd.max((a - b).abs());
}
assert!(maxd < 1e-3, "direct vs converged SOR max diff = {maxd}");
}
#[test]
fn anchorless_bilinear_cloud_factors_and_solves() {
let (xori, yori, cell, nc, nr) = (431_000.0, 6_521_000.0, 100.0, 41usize, 37usize);
let lattice = Lattice::regular(xori, yori, cell, cell, nc, nr);
let span_x = cell * (nc - 1) as f64;
let span_y = cell * (nr - 1) as f64;
let mut s: u64 = 0x9E37_79B9_7F4A_7C15;
let mut next = || {
s = s
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((s >> 33) as f64) / (u32::MAX as f64)
};
let mut sample_xy: Vec<[f64; 2]> = Vec::new();
let mut z: Vec<f64> = Vec::new();
for _ in 0..1500 {
let fx = next();
let fy = next();
let x = xori + (0.2 + 0.6 * fx) * span_x + 0.37;
let y = yori + (0.2 + 0.6 * fy) * span_y + 0.29;
sample_xy.push([x, y]);
z.push(2000.0 + 30.0 * fx + 20.0 * fy); }
assert_eq!(
sample_xy
.iter()
.filter(|xy| matches!(
classify_sample(&lattice, xy[0], xy[1], Conditioning::Bilinear),
SampleKind::OnNode(..)
))
.count(),
0,
"fixture must be anchorless"
);
let op = MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::Bilinear)
.expect("anchorless bilinear system must factor (residual null mode pinned)");
assert_eq!(op.anchor_count(), 0);
let field = op.solve(&z).unwrap();
assert!(
field.iter().all(|v| v.is_finite()),
"solved field must be all-finite (not the old all-NaN singular result)"
);
let mut max_misfit = 0.0_f64;
let mut sse = 0.0_f64;
for (xy, &d) in sample_xy.iter().zip(z.iter()) {
if let SampleKind::Off { nodes, w } =
classify_sample(&lattice, xy[0], xy[1], Conditioning::Bilinear)
{
let interp: f64 = (0..4).map(|k| w[k] * field[[nodes[k].0, nodes[k].1]]).sum();
let e = (interp - d).abs();
max_misfit = max_misfit.max(e);
sse += e * e;
}
}
let rms = (sse / sample_xy.len() as f64).sqrt();
assert!(
max_misfit < 0.5 && rms < 0.1,
"off-node data not honoured: max {max_misfit} m, rms {rms} m"
);
}
#[test]
fn under_constrained_is_not_pd() {
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 6, 6);
let sample_xy = [[2.0, 3.0]];
let r = MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::NearestNode);
assert!(
matches!(r, Err(AlgoError::InvalidArgument(_))),
"single anchor must be non-PD"
);
}
#[test]
fn solve_length_mismatch_errors() {
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 5, 5);
let sample_xy = [[0.0, 0.0], [4.0, 0.0], [0.0, 4.0], [4.0, 4.0]];
let op =
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::NearestNode).unwrap();
assert!(matches!(
op.solve(&[1.0, 2.0]),
Err(AlgoError::InvalidArgument(_))
));
}
#[test]
fn degenerate_lattice_errors() {
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 5, 1);
let sample_xy = [[0.0, 0.0], [4.0, 0.0]];
assert!(matches!(
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::NearestNode),
Err(AlgoError::InvalidGeometry(_))
));
}
#[test]
fn all_anchored_returns_anchor_values() {
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 3, 3);
let mut sample_xy = Vec::new();
let mut z = Vec::new();
for j in 0..3 {
for i in 0..3 {
sample_xy.push([i as f64, j as f64]);
z.push((i + 10 * j) as f64);
}
}
let op =
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::NearestNode).unwrap();
assert_eq!(op.free_count(), 0);
assert_eq!(op.anchor_count(), 9);
let field = op.solve(&z).unwrap();
for j in 0..3 {
for i in 0..3 {
assert_eq!(field[[i, j]], (i + 10 * j) as f64);
}
}
}
#[test]
#[ignore]
fn hotspot_timing() {
use super::super::min_curvature::grid_min_curvature_sor;
use std::time::Instant;
let (nc, nr) = (122usize, 116usize);
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, nc, nr);
let step = 0.6f64;
let mut coords: Vec<[f64; 3]> = Vec::new();
let mut y = 0.3;
while y < (nr - 1) as f64 {
let mut x = 0.2;
while x < (nc - 1) as f64 {
let z = 2000.0 + 0.02 * x - 0.015 * y + 6.0 * ((x * 0.11).sin() + (y * 0.09).cos());
coords.push([x, y, z]);
x += step;
}
y += step;
}
let sample_xy: Vec<[f64; 2]> = coords.iter().map(|c| [c[0], c[1]]).collect();
let z: Vec<f64> = coords.iter().map(|c| c[2]).collect();
println!(
"\nhotspot: {}x{} lattice ({} nodes), {} off-node samples",
nc,
nr,
nc * nr,
coords.len()
);
let t = Instant::now();
let op =
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::Bilinear).unwrap();
let t_factor = t.elapsed();
let t = Instant::now();
let _field = op.solve(&z).unwrap();
let t_solve = t.elapsed();
println!(
" DIRECT factor = {:?}, solve = {:?}, one-shot total = {:?}",
t_factor,
t_solve,
t_factor + t_solve
);
let n_horizons = 20;
let t = Instant::now();
for h in 0..n_horizons {
let zz: Vec<f64> = z.iter().map(|v| v + h as f64).collect();
let _ = op.solve(&zz).unwrap();
}
let per = t.elapsed() / n_horizons;
println!(" REUSE per-horizon solve (x{n_horizons}) = {per:?}");
let t = Instant::now();
let _ = grid_min_curvature_sor(&coords, &lattice, None, Conditioning::Bilinear);
println!(" SOR (old cold one-shot) = {:?}\n", t.elapsed());
}
#[test]
#[ignore]
fn anchorless_bench() {
use std::time::Instant;
let (xori, yori, cell, nc, nr) = (431_000.0, 6_521_000.0, 100.0, 41usize, 37usize);
let lattice = Lattice::regular(xori, yori, cell, cell, nc, nr);
let span_x = cell * (nc - 1) as f64;
let span_y = cell * (nr - 1) as f64;
let mut s: u64 = 0x9E37_79B9_7F4A_7C15;
let mut next = || {
s = s
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((s >> 33) as f64) / (u32::MAX as f64)
};
let mut sample_xy: Vec<[f64; 2]> = Vec::new();
let mut z: Vec<f64> = Vec::new();
for _ in 0..1500 {
let fx = next();
let fy = next();
sample_xy.push([
xori + (0.2 + 0.6 * fx) * span_x + 0.37,
yori + (0.2 + 0.6 * fy) * span_y + 0.29,
]);
z.push(2000.0 + 30.0 * fx + 20.0 * fy);
}
let t = Instant::now();
let op =
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::Bilinear).unwrap();
let t_factor = t.elapsed();
let t = Instant::now();
let field = op.solve(&z).unwrap();
let t_solve = t.elapsed();
let (mut mx, mut sse) = (0.0f64, 0.0f64);
for (xy, &d) in sample_xy.iter().zip(z.iter()) {
if let SampleKind::Off { nodes, w } =
classify_sample(&lattice, xy[0], xy[1], Conditioning::Bilinear)
{
let v: f64 = (0..4).map(|k| w[k] * field[[nodes[k].0, nodes[k].1]]).sum();
mx = mx.max((v - d).abs());
sse += (v - d) * (v - d);
}
}
println!(
"\nanchorless {nc}x{nr} ({} nodes), {} off-node samples, 0 anchors\n \
stabilized factor = {t_factor:?}, solve = {t_solve:?}, one-shot = {:?}\n \
data misfit: max = {mx:e} m, rms = {:e} m\n",
nc * nr,
sample_xy.len(),
t_factor + t_solve,
(sse / sample_xy.len() as f64).sqrt()
);
}
#[test]
fn factor_once_solve_many_horizons() {
let lattice = Lattice::regular(0.0, 0.0, 1.0, 1.0, 7, 7);
let sample_xy = [[0.0, 0.0], [6.0, 0.0], [0.0, 6.0], [6.0, 6.0], [3.0, 3.0]];
let op =
MinCurvatureOperator::factor(&lattice, &sample_xy, Conditioning::NearestNode).unwrap();
let plane = |a: f64, b: f64, c: f64| {
sample_xy
.iter()
.map(|p| a * p[0] + b * p[1] + c)
.collect::<Vec<_>>()
};
for (a, b, c) in [(2.0, -3.0, 5.0), (-1.0, 4.0, 1.0)] {
let field = op.solve(&plane(a, b, c)).unwrap();
for j in 0..7 {
for i in 0..7 {
let want = a * i as f64 + b * j as f64 + c;
assert!(
(field[[i, j]] - want).abs() < 1e-6,
"({i},{j}): {} vs {want}",
field[[i, j]]
);
}
}
}
}
}