rust/vectorhd-core/src/energy.rs

//! Composite + data energy — the `render` composite loop and `data_energy` from Python.
//!
//! `E_data = (1/l0)·Σ_p ‖Render(P)[p] − I[p]‖²`.
//!
//! **Summation.** Python computes `np.sum(r*r, axis=2)` (an exact 3-term per-pixel sum) and then
//! `.sum()` over the H×W plane, where numpy uses *pairwise* summation. Reproducing numpy's exact
//! pairwise blocking would be brittle, so the per-pixel grouping is mirrored exactly and the outer
//! reduction uses **Neumaier compensated summation** instead. Both are accurate to ~1 ulp of the
//! true sum, so they agree to ~1e-16 relative — five orders of magnitude inside the 1e-9 gate,
//! and *closer to the true value* than a naive sequential sum would be.

use crate::model::Vectors;

/// Neumaier (improved Kahan) compensated summation.
#[derive(Default, Clone, Copy)]
pub struct Neumaier {
    sum: f64,
    c: f64,
}

impl Neumaier {
    #[inline]
    pub fn add(&mut self, x: f64) {
        let t = self.sum + x;
        if self.sum.abs() >= x.abs() {
            self.c += (self.sum - t) + x;
        } else {
            self.c += (x - t) + self.sum;
        }
        self.sum = t;
    }

    #[inline]
    pub fn total(&self) -> f64 {
        self.sum + self.c
    }
}

pub fn sum_compensated(xs: impl Iterator<Item = f64>) -> f64 {
    let mut s = Neumaier::default();
    for x in xs {
        s.add(x);
    }
    s.total()
}

/// The composite loop from `render`: `img = bg + Σ_R cov_R·(c_R − bg)`, then clipped to [0,1].
/// Regions are visited in face order, matching Python's dict iteration.
pub fn compose(v: &Vectors, cov: &[(i64, Vec<f64>)]) -> Vec<f64> {
    let bg = v.background;
    let n = v.width * v.height;
    let mut img = vec![bg; n * 3];
    let colors = v.colors01();
    for (label, cvr) in cov {
        let c = colors
            .get(label)
            .unwrap_or_else(|| panic!("region {label} has no exported color"));
        let (d0, d1, d2) = (c[0] - bg, c[1] - bg, c[2] - bg);
        for p in 0..n {
            let a = cvr[p];
            img[p * 3] += a * d0;
            img[p * 3 + 1] += a * d1;
            img[p * 3 + 2] += a * d2;
        }
    }
    for x in img.iter_mut() {
        *x = x.clamp(0.0, 1.0);
    }
    img
}

/// `E_data`, normalized by `l0`.
pub fn e_data(img: &[f64], target: &[f64], l0: f64) -> f64 {
    let n = img.len() / 3;
    let per_pixel = (0..n).map(|p| {
        // exact 3-term group, mirroring np.sum(r*r, axis=2)
        let d0 = img[p * 3] - target[p * 3];
        let d1 = img[p * 3 + 1] - target[p * 3 + 1];
        let d2 = img[p * 3 + 2] - target[p * 3 + 2];
        d0 * d0 + d1 * d1 + d2 * d2
    });
    sum_compensated(per_pixel) / l0
}

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

    #[test]
    fn neumaier_beats_naive_on_a_hostile_sequence() {
        // 1.0 followed by many tiny values: naive f64 loses all of them.
        let mut xs = vec![1.0_f64];
        xs.extend(std::iter::repeat(1e-17).take(1000));
        let naive = xs.iter().fold(0.0_f64, |a, b| a + b);
        let comp = sum_compensated(xs.iter().copied());
        assert_eq!(naive, 1.0, "naive summation should lose the tail");
        assert!(comp > 1.0, "compensated summation should keep it, got {comp}");
    }

    #[test]
    fn energy_is_zero_for_a_perfect_match() {
        let img = vec![0.25, 0.5, 0.75, 0.1, 0.2, 0.3];
        let e = e_data(&img, &img, 2.0);
        assert_eq!(e, 0.0);
    }

    #[test]
    fn energy_scales_by_inverse_l0() {
        let img = vec![1.0, 1.0, 1.0];
        let tgt = vec![0.0, 0.0, 0.0];
        assert!((e_data(&img, &tgt, 1.0) - 3.0).abs() < 1e-15);
        assert!((e_data(&img, &tgt, 3.0) - 1.0).abs() < 1e-15);
    }
}