rust/vectorhd-core/src/gradient.rs

//! Boundary-integral gradient — a port of `geometry_gradient` / `_segment_tau_forces`, stopped
//! at **polyline vertices** (the Bernstein chain rule to control points stays in Python).
//!
//! For an edge separating regions L and R, displacing a boundary point along its unit normal
//! sweeps area from one region into the other:
//!
//! ```text
//! ∂E/∂v_i  =  (2/l0) · Σ_{segments touching v_i}  ∫ τ-ramp · (resid·(c_L − c_R)) ds · n
//! ```
//!
//! Two conventions are load-bearing and pinned by the Python tests:
//! * **Left normal**: `n = (−seg.y, seg.x)/‖seg‖`. Flipping it flips the whole gradient.
//! * **Truncation, not floor**: the pixel a sub-piece falls in is `int(a + t·d)`, which truncates
//!   toward zero in Python; Rust's `as i64` matches. They differ for negative coordinates, which
//!   do occur when geometry hangs off the raster.

use crate::model::{Pt, Vectors};
use std::collections::HashMap;

fn push_unique(v: &mut Vec<f64>, x: f64) {
    if !v.iter().any(|&y| y == x) {
        v.push(x);
    }
}

/// `(∫(1−τ)·field ds, ∫τ·field ds)` along `a → b`, exact for a per-pixel-constant field: split at
/// integer grid lines, and within each sub-piece τ is linear so `∫(1−τ)ds = len·(1−τ_mid)`.
pub fn segment_tau_forces(a: Pt, b: Pt, field: &[f64], width: usize, height: usize) -> (f64, f64) {
    let dx = b[0] - a[0];
    let dy = b[1] - a[1];
    let length = dx.hypot(dy);
    if length < 1e-12 {
        return (0.0, 0.0);
    }

    let mut cuts: Vec<f64> = Vec::with_capacity(8);
    push_unique(&mut cuts, 0.0);
    push_unique(&mut cuts, 1.0);

    if dx.abs() > 1e-12 {
        let k0 = a[0].min(b[0]).ceil() as i64;
        let k1 = a[0].max(b[0]).floor() as i64;
        for k in k0..=k1 {
            let t = (k as f64 - a[0]) / dx;
            if t > 0.0 && t < 1.0 {
                push_unique(&mut cuts, t);
            }
        }
    }
    if dy.abs() > 1e-12 {
        let k0 = a[1].min(b[1]).ceil() as i64;
        let k1 = a[1].max(b[1]).floor() as i64;
        for k in k0..=k1 {
            let t = (k as f64 - a[1]) / dy;
            if t > 0.0 && t < 1.0 {
                push_unique(&mut cuts, t);
            }
        }
    }
    cuts.sort_by(|p, q| p.partial_cmp(q).expect("no NaN in cut parameters"));

    let mut i0 = 0.0_f64;
    let mut i1 = 0.0_f64;
    for w in cuts.windows(2) {
        let (t0, t1) = (w[0], w[1]);
        let tm = 0.5 * (t0 + t1);
        let px = (a[0] + tm * dx) as i64; // Python int(): truncate toward zero
        let py = (a[1] + tm * dy) as i64;
        if px >= 0 && (px as usize) < width && py >= 0 && (py as usize) < height {
            let seg = field[py as usize * width + px as usize] * (t1 - t0) * length;
            i0 += seg * (1.0 - tm);
            i1 += seg * tm;
        }
    }
    (i0, i1)
}

/// `{edge_id: [per-cubic [per-vertex [gx, gy]]]}` for every interior edge.
pub fn vertex_gradients(
    v: &Vectors,
    img: &[f64],
    target: &[f64],
) -> HashMap<i64, Vec<Vec<[f64; 2]>>> {
    let (width, height) = (v.width, v.height);
    let n = width * height;
    let scale = 2.0 / v.l0;
    let colors = v.colors01();

    let mut out: HashMap<i64, Vec<Vec<[f64; 2]>>> = HashMap::new();
    let mut field = vec![0.0_f64; n];

    for e in &v.edges {
        if !e.interior {
            continue;
        }
        let cl = colors[&e.left_label];
        let cr = colors[&e.right_label];
        let cd = [cl[0] - cr[0], cl[1] - cr[1], cl[2] - cr[2]];

        // field = (img − target) · cdiff, per pixel — Python's `resid @ cdiff`.
        for p in 0..n {
            let r0 = img[p * 3] - target[p * 3];
            let r1 = img[p * 3 + 1] - target[p * 3 + 1];
            let r2 = img[p * 3 + 2] - target[p * 3 + 2];
            field[p] = r0 * cd[0] + r1 * cd[1] + r2 * cd[2];
        }

        let mut per_cubic: Vec<Vec<[f64; 2]>> = Vec::with_capacity(e.cubics.len());
        for pts in &e.cubics {
            let m = pts.len();
            let mut vgrad = vec![[0.0_f64; 2]; m];
            for i in 0..m - 1 {
                let sx = pts[i + 1][0] - pts[i][0];
                let sy = pts[i + 1][1] - pts[i][1];
                let length = sx.hypot(sy);
                if length < 1e-12 {
                    continue;
                }
                let n0 = -sy / length; // LEFT normal; sign pinned by the Python tests
                let n1 = sx / length;
                let (i0, i1) = segment_tau_forces(pts[i], pts[i + 1], &field, width, height);
                // Python evaluates `scale * i0` first, then scales the normal vector.
                let f0 = scale * i0;
                let f1 = scale * i1;
                vgrad[i][0] += f0 * n0;
                vgrad[i][1] += f0 * n1;
                vgrad[i + 1][0] += f1 * n0;
                vgrad[i + 1][1] += f1 * n1;
            }
            per_cubic.push(vgrad);
        }
        out.insert(e.id, per_cubic);
    }
    out
}

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

    /// A uniform field over a segment fully inside one pixel: the τ-ramp must split the total
    /// `field · length` into halves, since ∫(1−τ)dτ = ∫τ dτ = 1/2.
    #[test]
    fn tau_forces_split_evenly_for_a_uniform_field() {
        let field = vec![2.0; 4]; // 2x2 raster, constant
        let (i0, i1) = segment_tau_forces([0.2, 0.2], [0.8, 0.2], &field, 2, 2);
        let expected = 2.0 * 0.6 * 0.5;
        assert!((i0 - expected).abs() < 1e-12, "i0 {i0} != {expected}");
        assert!((i1 - expected).abs() < 1e-12, "i1 {i1} != {expected}");
    }

    #[test]
    fn tau_forces_vanish_on_a_degenerate_segment() {
        let field = vec![5.0; 4];
        let (i0, i1) = segment_tau_forces([0.5, 0.5], [0.5, 0.5], &field, 2, 2);
        assert_eq!((i0, i1), (0.0, 0.0));
    }

    #[test]
    fn tau_forces_ignore_geometry_off_the_raster() {
        let field = vec![7.0; 4];
        let (i0, i1) = segment_tau_forces([-5.0, -5.0], [-4.0, -5.0], &field, 2, 2);
        assert_eq!((i0, i1), (0.0, 0.0));
    }

    /// A segment crossing a pixel boundary must be split, weighting each pixel's field by the
    /// length inside it.
    #[test]
    fn tau_forces_split_at_pixel_boundaries() {
        // 2x1 raster, field 1.0 in column 0 and 3.0 in column 1.
        let field = vec![1.0, 3.0];
        let (i0, i1) = segment_tau_forces([0.0, 0.5], [2.0, 0.5], &field, 2, 1);
        // piece A: t∈[0,0.5], tm=0.25, px=0 → 1.0·0.5·2 = 1.0 ; i0 += 0.75, i1 += 0.25
        // piece B: t∈[0.5,1], tm=0.75, px=1 → 3.0·0.5·2 = 3.0 ; i0 += 0.75, i1 += 2.25
        assert!((i0 - 1.5).abs() < 1e-12, "i0 {i0}");
        assert!((i1 - 2.5).abs() < 1e-12, "i1 {i1}");
    }

    /// A clean two-region vertical split, used by both gradient tests below.
    ///
    /// Region A = [0, x] × [0, H] (color `CA`), region B = [x, W] × [0, H] (color `CB`). The only
    /// interior boundary is the vertical segment (x,0) → (x,H); the remaining edges lie on the
    /// raster frame and are horizontal, so they contribute nothing to a y-integral. Sliding `x`
    /// therefore perturbs exactly one edge, which is what makes the finite difference clean.
    ///
    /// **Orientation.** Traversed +y, the segment's left normal is `(−dy, dx)/‖d‖ = (−1, 0)`,
    /// pointing toward region A. The normal points at the *right* label, so `L = B`, `R = A` and
    /// the colour jump is `c_B − c_A`. Getting this backwards flips the gradient's sign.
    mod split {
        use super::*;

        pub const W: usize = 12;
        pub const H: usize = 12;
        pub const BG: f64 = 1.0;
        pub const CA: [f64; 3] = [0.2, 0.4, 0.6];
        pub const CB: [f64; 3] = [0.9, 0.8, 0.7];

        pub fn render(x: f64) -> Vec<f64> {
            use crate::coverage::polygon_coverage;
            let a = vec![[0.0, 0.0], [x, 0.0], [x, H as f64], [0.0, H as f64]];
            let b = vec![
                [x, 0.0],
                [W as f64, 0.0],
                [W as f64, H as f64],
                [x, H as f64],
            ];
            let ca = polygon_coverage(&[a], W, H);
            let cb = polygon_coverage(&[b], W, H);
            let mut img = vec![BG; W * H * 3];
            for p in 0..W * H {
                for ch in 0..3 {
                    img[p * 3 + ch] +=
                        ca[p].abs() * (CA[ch] - BG) + cb[p].abs() * (CB[ch] - BG);
                }
            }
            img
        }

        pub fn energy(x: f64, target: &[f64], l0: f64) -> f64 {
            crate::energy::e_data(&render(x), target, l0)
        }

        /// dE/dx from the kernel's own primitives, summing both endpoints' x-gradients (they
        /// translate together when the whole edge slides).
        pub fn analytic_dx(x: f64, target: &[f64], l0: f64) -> f64 {
            let img = render(x);
            // c_L − c_R = CB − CA, per the orientation note above.
            let cd = [CB[0] - CA[0], CB[1] - CA[1], CB[2] - CA[2]];
            let field: Vec<f64> = (0..W * H)
                .map(|p| {
                    (0..3)
                        .map(|ch| (img[p * 3 + ch] - target[p * 3 + ch]) * cd[ch])
                        .sum()
                })
                .collect();
            let (a, b) = ([x, 0.0], [x, H as f64]);
            let (sx, sy) = (b[0] - a[0], b[1] - a[1]);
            let length = sx.hypot(sy);
            let n0 = -sy / length;
            let (i0, i1) = segment_tau_forces(a, b, &field, W, H);
            let scale = 2.0 / l0;
            scale * i0 * n0 + scale * i1 * n0
        }
    }

    /// Internal finite-difference check: the analytic vertex gradient must match a central
    /// difference of `E_data`. Built standalone rather than from a fixture, so it does not depend
    /// on the exported vectors.
    #[test]
    fn gradient_matches_central_difference() {
        let l0 = 1.0_f64;
        let target = split::render(7.3); // minimum deliberately away from the probe point
        let x0 = 4.37_f64; // OFF the integer grid — see the test below for why that matters

        let analytic = split::analytic_dx(x0, &target, l0);
        let h = 1e-6;
        let fd = (split::energy(x0 + h, &target, l0) - split::energy(x0 - h, &target, l0))
            / (2.0 * h);

        let rel = (analytic - fd).abs() / fd.abs().max(1e-12);
        assert!(
            rel < 1e-5,
            "analytic {analytic:.9e} vs finite-difference {fd:.9e} (rel {rel:.2e})"
        );
    }

    /// Documents a real property of the Python kernel that a port must preserve: when a segment
    /// lies exactly on an integer grid line, `field` is sampled with `int()` (truncation toward
    /// zero), which reads the pixel to the **right**. `E_data` is genuinely kinked there, so the
    /// analytic value is a *one-sided* derivative and a symmetric finite difference — which
    /// straddles the kink — legitimately disagrees.
    ///
    /// This is not a defect to fix; it is a tie-break that must be replicated bit-for-bit, and it
    /// is not exotic: fitted control points come off the crack grid, so integer coordinates are
    /// the *initial* state of every optimization.
    #[test]
    fn on_grid_line_the_derivative_is_one_sided_by_design() {
        let l0 = 1.0_f64;
        let target = split::render(7.3);
        let h = 1e-6;

        // Off-grid: analytic and central difference agree.
        let off = 4.37_f64;
        let a_off = split::analytic_dx(off, &target, l0);
        let fd_off = (split::energy(off + h, &target, l0) - split::energy(off - h, &target, l0))
            / (2.0 * h);
        assert!(
            (a_off - fd_off).abs() / fd_off.abs().max(1e-12) < 1e-5,
            "off-grid should agree: {a_off} vs {fd_off}"
        );

        // On-grid: the analytic value equals the RIGHT-hand derivative, not the central one.
        let on = 5.0_f64;
        let a_on = split::analytic_dx(on, &target, l0);
        let right = (split::energy(on + h, &target, l0) - split::energy(on, &target, l0)) / h;
        let central =
            (split::energy(on + h, &target, l0) - split::energy(on - h, &target, l0)) / (2.0 * h);
        assert!(
            (a_on - right).abs() / right.abs().max(1e-12) < 1e-4,
            "on-grid analytic {a_on} should match the right derivative {right}"
        );
        assert!(
            (a_on - central).abs() / central.abs().max(1e-12) > 1e-3,
            "the kink should be visible: analytic {a_on} vs central {central}"
        );
    }
}