graph_explorer_layout/

quadtree.rs

//! Barnes-Hut quadtree for many-body repulsion in `O(n log n)`.
//!
//! Each cell accumulates the mass and position-sum of the bodies beneath it.
//! When a cell is far enough away relative to its size (`width / distance <
//! theta`) it is treated as one body at its centroid instead of being descended
//! into — which is the entire reason this is not `O(n²)`.
//!
//! The force law is d3-force's many-body law: the force on `a` from `b` is
//! `strength * (b - a) / d²`, whose *magnitude* therefore falls off as `1/d`.
//! Negative `strength` repels. Both properties are load-bearing — see the
//! comment on [`QuadTree::repulsion`].

/// Sentinel for "no cell" / "no body". `usize::MAX` cannot collide with a real
/// index because the arena would have to hold `usize::MAX` cells first.
const NONE: usize = usize::MAX;

/// Subdivision depth cap. Without it, coincident or near-coincident bodies
/// subdivide until the cell size underflows and the build never terminates.
const MAX_DEPTH: u32 = 24;

/// Below this squared distance two bodies are treated as exactly coincident and
/// separated by a deterministic unit jitter rather than dividing by ~zero.
const EPS2: f32 = 1e-6;

/// Squared minimum interaction distance (d3's `distanceMin`, default 1).
///
/// Without it the `1/d` law is unbounded as `d -> 0`: two bodies 0.01 apart
/// would exchange a force 100x larger than the strength, which in one tick
/// launches them off-screen. Clamping bounds every pairwise force by
/// `|strength|`.
const MIN_DIST2: f32 = 1.0;

#[derive(Clone)]
struct Cell {
    center: [f32; 2],
    half: f32,
    /// Number of bodies beneath this cell; every node weighs 1.
    mass: f32,
    /// Sum of member positions. Centroid is `sum / mass`.
    sum: [f32; 2],
    /// Child cells (NW, NE, SW, SE), `NONE` when absent.
    kids: [usize; 4],
    /// Head of this leaf's body chain, or `NONE` for an internal cell.
    body: usize,
}

pub struct QuadTree {
    cells: Vec<Cell>,
    /// Intrusive linked list: `next[b]` is the next body sharing `b`'s leaf.
    /// Lets a leaf hold several coincident bodies without allocating per leaf.
    next: Vec<usize>,
}

impl QuadTree {
    pub fn build(pos: &[[f32; 2]]) -> Self {
        let mut t = QuadTree { cells: Vec::new(), next: vec![NONE; pos.len()] };
        if pos.is_empty() {
            return t;
        }
        // `f32::MIN` is the most-negative *finite* f32, so it is a valid
        // identity for a running max over finite inputs (unlike, say,
        // `f32::MIN_POSITIVE`). Non-finite inputs would poison it, but they
        // would poison every downstream force anyway.
        let (mut lo, mut hi) = ([f32::MAX; 2], [f32::MIN; 2]);
        for p in pos {
            lo[0] = lo[0].min(p[0]); lo[1] = lo[1].min(p[1]);
            hi[0] = hi[0].max(p[0]); hi[1] = hi[1].max(p[1]);
        }
        // A square root cell keeps quadrants square, so `half` is a valid
        // proxy for cell width in the theta test.
        let half = ((hi[0] - lo[0]).max(hi[1] - lo[1]) * 0.5).max(1.0);
        let center = [(lo[0] + hi[0]) * 0.5, (lo[1] + hi[1]) * 0.5];
        t.cells.push(Cell { center, half, mass: 0.0, sum: [0.0, 0.0], kids: [NONE; 4], body: NONE });
        for i in 0..pos.len() {
            t.insert(0, i, pos, 0);
        }
        t
    }

    fn quadrant(center: [f32; 2], p: [f32; 2]) -> usize {
        // 0 NW, 1 NE, 2 SW, 3 SE
        (if p[0] >= center[0] { 1 } else { 0 }) | (if p[1] < center[1] { 2 } else { 0 })
    }

    fn child_cell(parent: &Cell, q: usize) -> Cell {
        let h = parent.half * 0.5;
        let dx = if q & 1 == 1 { h } else { -h };
        let dy = if q & 2 == 2 { -h } else { h };
        Cell {
            center: [parent.center[0] + dx, parent.center[1] + dy],
            half: h,
            mass: 0.0,
            sum: [0.0, 0.0],
            kids: [NONE; 4],
            body: NONE,
        }
    }

    fn insert(&mut self, cell: usize, b: usize, pos: &[[f32; 2]], depth: u32) {
        self.cells[cell].mass += 1.0;
        self.cells[cell].sum[0] += pos[b][0];
        self.cells[cell].sum[1] += pos[b][1];

        let is_internal = self.cells[cell].kids.iter().any(|&k| k != NONE);
        if is_internal {
            let q = Self::quadrant(self.cells[cell].center, pos[b]);
            let kid = self.ensure_kid(cell, q);
            self.insert(kid, b, pos, depth + 1);
            return;
        }

        let existing = self.cells[cell].body;
        if existing == NONE {
            self.cells[cell].body = b;
            return;
        }
        if depth >= MAX_DEPTH {
            // Coincident (or effectively so). Chain rather than subdivide —
            // `repulsion` walks the chain and skips self, so a multi-body leaf
            // stays exact instead of self-interacting.
            self.next[b] = existing;
            self.cells[cell].body = b;
            return;
        }

        // Split: push the sitting body down, then the new one.
        self.cells[cell].body = NONE;
        for who in [existing, b] {
            let q = Self::quadrant(self.cells[cell].center, pos[who]);
            let kid = self.ensure_kid(cell, q);
            // The parent's mass/sum already counted both, so recurse into the
            // child directly rather than back through `insert` on the parent.
            self.insert(kid, who, pos, depth + 1);
        }
    }

    fn ensure_kid(&mut self, cell: usize, q: usize) -> usize {
        let existing = self.cells[cell].kids[q];
        if existing != NONE {
            return existing;
        }
        let c = Self::child_cell(&self.cells[cell], q);
        self.cells.push(c);
        let idx = self.cells.len() - 1;
        self.cells[cell].kids[q] = idx;
        idx
    }

    /// Force on body `i`. Negative `strength` repels (d3's convention).
    ///
    /// The contribution of body `b` is `strength * (b - i) / d²`, so a negative
    /// strength points the force *away* from `b`. The magnitude falls off as
    /// `1/d`, not `1/d²`: that is d3's law, and it is what makes the Barnes-Hut
    /// monopole approximation accurate enough to run at `theta = 0.9`.
    pub fn repulsion(&self, i: usize, pos: &[[f32; 2]], theta: f32, strength: f32) -> [f32; 2] {
        let mut f = [0.0f32, 0.0];
        if self.cells.is_empty() || i >= pos.len() {
            return f;
        }
        self.walk(0, i, pos, theta, strength, &mut f);
        f
    }

    fn walk(&self, cell: usize, i: usize, pos: &[[f32; 2]], theta: f32, strength: f32, f: &mut [f32; 2]) {
        let c = &self.cells[cell];
        if c.mass == 0.0 {
            return;
        }
        let is_internal = c.kids.iter().any(|&k| k != NONE);

        if !is_internal {
            // Leaf: apply each chained body individually, skipping self.
            let mut b = c.body;
            while b != NONE {
                if b != i {
                    Self::pair(pos[i], pos[b], i, b, strength, f);
                }
                b = self.next[b];
            }
            return;
        }

        let centroid = [c.sum[0] / c.mass, c.sum[1] / c.mass];
        let dx = centroid[0] - pos[i][0];
        let dy = centroid[1] - pos[i][1];
        let d2 = dx * dx + dy * dy;
        // `half * 2` is the cell width, since `child_cell` halves `half` at
        // every level and a cell spans `center ± half`. So this is d3's
        // `width / distance < theta`, compared in squared form.
        if d2 > 0.0 && (c.half * 2.0) * (c.half * 2.0) < theta * theta * d2 {
            let d2 = Self::clamp_near(d2);
            let mag = strength * c.mass / d2;
            f[0] += dx * mag;
            f[1] += dy * mag;
            return;
        }
        for &k in &c.kids {
            if k != NONE {
                self.walk(k, i, pos, theta, strength, f);
            }
        }
    }

    /// d3's `distanceMin` softening: pull sub-`MIN_DIST2` separations back up
    /// toward the floor so the `1/d` law cannot produce an unbounded force.
    fn clamp_near(d2: f32) -> f32 {
        if d2 < MIN_DIST2 { (MIN_DIST2 * d2).sqrt() } else { d2 }
    }

    fn pair(a: [f32; 2], b: [f32; 2], ai: usize, bi: usize, strength: f32, f: &mut [f32; 2]) {
        let mut dx = b[0] - a[0];
        let mut dy = b[1] - a[1];
        let mut d2 = dx * dx + dy * dy;
        if d2 < EPS2 {
            // Exactly (or all but) coincident: there is no meaningful direction,
            // so pick one deterministically from the index pair — a rebuild of
            // the same graph then produces the same layout, the seeded-layout
            // property the now-deleted one-shot layout also guaranteed.
            //
            // The jitter is a *unit* vector applied at the minimum interaction
            // distance, so the resulting force is exactly `|strength|`: the same
            // nudge two bodies one unit apart would feel. Using the raw jitter
            // length here instead would divide by ~1e-6 and yield a force ~1e5x
            // the strength, which is a catapult, not a nudge.
            let s = (ai.wrapping_mul(31).wrapping_add(bi)) as f32;
            let (jx, jy) = ((s * 0.7).sin(), (s * 1.3).cos());
            let n = (jx * jx + jy * jy).sqrt().max(1e-6);
            dx = jx / n;
            dy = jy / n;
            d2 = MIN_DIST2;
        }
        let d2 = Self::clamp_near(d2);
        let mag = strength / d2;
        f[0] += dx * mag;
        f[1] += dy * mag;
    }
}

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

    /// Reference implementation: every pair, no approximation.
    ///
    /// Mirrors d3-force's many-body law exactly — `strength * (p - i) / d²`,
    /// magnitude `|strength| / d`. The direction is carried by the unnormalised
    /// delta, so there is no separate unit-vector step; normalising *and*
    /// dividing by `d²` would be a `1/d²` law, which is a different force.
    fn brute_force(pos: &[[f32; 2]], i: usize, strength: f32) -> [f32; 2] {
        let mut f = [0.0f32, 0.0];
        for (j, p) in pos.iter().enumerate() {
            if j == i { continue; }
            let dx = p[0] - pos[i][0];
            let dy = p[1] - pos[i][1];
            let d2 = (dx * dx + dy * dy).max(1e-6);
            let mag = strength / d2;
            f[0] += dx * mag;
            f[1] += dy * mag;
        }
        f
    }

    /// Deterministic scatter — no rand dependency, and reproducible failures.
    fn scatter(n: usize) -> Vec<[f32; 2]> {
        let mut s = 12345u64;
        (0..n)
            .map(|_| {
                let mut nxt = || {
                    s ^= s >> 12; s ^= s << 25; s ^= s >> 27;
                    ((s.wrapping_mul(0x2545F4914F6CDD1D) >> 33) as f32) / (1u64 << 31) as f32
                };
                [(nxt() - 0.5) * 1000.0, (nxt() - 0.5) * 1000.0]
            })
            .collect()
    }

    #[test]
    fn theta_zero_is_exact_because_no_cell_can_be_approximated() {
        // theta = 0 forces full recursion, so the tree must reproduce
        // brute force to floating-point noise. This is the test that proves
        // the traversal visits every body exactly once.
        let pos = scatter(60);
        let t = QuadTree::build(&pos);
        for i in 0..pos.len() {
            let got = t.repulsion(i, &pos, 0.0, -100.0);
            let want = brute_force(&pos, i, -100.0);
            assert!(
                (got[0] - want[0]).abs() < 0.05 && (got[1] - want[1]).abs() < 0.05,
                "body {i}: tree {got:?} vs brute {want:?}"
            );
        }
    }

    #[test]
    fn default_theta_approximates_brute_force_within_tolerance() {
        let pos = scatter(200);
        let t = QuadTree::build(&pos);
        let mut worst = 0.0f32;
        for i in 0..pos.len() {
            let got = t.repulsion(i, &pos, 0.9, -100.0);
            let want = brute_force(&pos, i, -100.0);
            let mag = (want[0] * want[0] + want[1] * want[1]).sqrt().max(1e-3);
            let err = ((got[0] - want[0]).powi(2) + (got[1] - want[1]).powi(2)).sqrt() / mag;
            worst = worst.max(err);
        }
        assert!(worst < 0.30, "worst relative error {worst} exceeds 30% at theta=0.9");
    }

    /// The monopole approximation should degrade smoothly with theta, and
    /// vanish at theta = 0. A tree whose mass or centroid accumulation is wrong
    /// can still pass the theta = 0 test (which never reads either) but will
    /// break this one, because every approximated cell then carries the error.
    #[test]
    fn accuracy_degrades_monotonically_with_theta() {
        let pos = scatter(200);
        let t = QuadTree::build(&pos);
        let err_at = |theta: f32| {
            let mut worst = 0.0f32;
            for i in 0..pos.len() {
                let got = t.repulsion(i, &pos, theta, -100.0);
                let want = brute_force(&pos, i, -100.0);
                let mag = (want[0] * want[0] + want[1] * want[1]).sqrt().max(1e-3);
                let e = ((got[0] - want[0]).powi(2) + (got[1] - want[1]).powi(2)).sqrt() / mag;
                worst = worst.max(e);
            }
            worst
        };
        let (e0, e3, e5, e9) = (err_at(0.0), err_at(0.3), err_at(0.5), err_at(0.9));
        assert!(e0 < 1e-4, "theta=0 must be exact, got {e0}");
        assert!(e3 < e5 && e5 < e9, "error must grow with theta: {e3} {e5} {e9}");
    }

    /// Mass and centroid are only exercised by *approximated* cells, so the
    /// theta = 0 test cannot catch a double-count. Check the accumulation
    /// directly: every internal cell must equal the sum of its children, and
    /// the root must hold every body exactly once.
    #[test]
    fn every_cell_holds_each_body_beneath_it_exactly_once() {
        let pos = scatter(200);
        let t = QuadTree::build(&pos);
        let n = pos.len() as f32;
        assert_eq!(t.cells[0].mass, n, "root mass must be the body count");
        let mean = [
            pos.iter().map(|p| p[0]).sum::<f32>() / n,
            pos.iter().map(|p| p[1]).sum::<f32>() / n,
        ];
        assert!((t.cells[0].sum[0] / n - mean[0]).abs() < 0.01);
        assert!((t.cells[0].sum[1] / n - mean[1]).abs() < 0.01);

        let mut reachable = 0usize;
        for c in &t.cells {
            if c.kids.iter().any(|&k| k != NONE) {
                let km: f32 = c.kids.iter().filter(|&&k| k != NONE).map(|&k| t.cells[k].mass).sum();
                assert!((km - c.mass).abs() < 1e-3, "cell mass {} vs kids {km}", c.mass);
            } else {
                // Leaf: its mass must equal the length of its body chain.
                let mut b = c.body;
                let mut chain = 0.0f32;
                while b != NONE {
                    reachable += 1;
                    chain += 1.0;
                    b = t.next[b];
                }
                assert!((chain - c.mass).abs() < 1e-3, "leaf mass {} vs chain {chain}", c.mass);
            }
        }
        assert_eq!(reachable, pos.len(), "every body must sit in exactly one leaf");
    }

    #[test]
    fn a_body_never_repels_itself() {
        // One body: the only possible force is a self-interaction, so any
        // non-zero result means the traversal failed to exclude it.
        let pos = vec![[10.0, 10.0]];
        let t = QuadTree::build(&pos);
        assert_eq!(t.repulsion(0, &pos, 0.9, -100.0), [0.0, 0.0]);
    }

    #[test]
    fn coincident_bodies_do_not_produce_nan() {
        // Two nodes at the identical point is not exotic — it is what a fresh
        // seed or a degenerate layout produces, and a 1/0 here poisons every
        // position downstream for the rest of the session.
        let pos = vec![[5.0, 5.0], [5.0, 5.0], [5.0, 5.0]];
        let t = QuadTree::build(&pos);
        for i in 0..3 {
            let f = t.repulsion(i, &pos, 0.9, -100.0);
            assert!(f[0].is_finite() && f[1].is_finite(), "body {i} produced {f:?}");
        }
    }

    /// A degenerate cluster must produce a *nudge*, not a catapult. Without a
    /// minimum interaction distance the coincident case divides by ~1e-6 and
    /// returns a force ~1e5x the strength, which ejects the cluster to infinity
    /// on the first tick while still being perfectly finite.
    #[test]
    fn coincident_bodies_are_nudged_not_launched() {
        let pos = vec![[5.0, 5.0], [5.0, 5.0], [5.0, 5.0]];
        let t = QuadTree::build(&pos);
        for i in 0..3 {
            let f = t.repulsion(i, &pos, 0.9, -100.0);
            let mag = (f[0] * f[0] + f[1] * f[1]).sqrt();
            assert!(mag <= 200.0, "body {i} force {mag} exceeds two bodies' worth of strength");
        }
        // And the nudge must be deterministic across rebuilds.
        let t2 = QuadTree::build(&pos);
        assert_eq!(t.repulsion(0, &pos, 0.9, -100.0), t2.repulsion(0, &pos, 0.9, -100.0));
    }

    /// The `distanceMin` softening in [`QuadTree::clamp_near`] is what makes a
    /// near-coincident pair a nudge rather than a catapult, and *deleting it
    /// passes every other test in this file* — the `1/d` law stays perfectly
    /// finite as it runs away. At `d = 0.01` with the shipping charge strength
    /// the softened force is 800; unsoftened it is 80,000, which ejects the
    /// pair off-screen on the very first tick.
    #[test]
    fn a_near_coincident_pair_is_softened_to_about_the_strength() {
        let pos = vec![[0.0, 0.0], [0.01, 0.0]];
        let t = QuadTree::build(&pos);
        let f = t.repulsion(0, &pos, 0.9, -800.0);
        let mag = (f[0] * f[0] + f[1] * f[1]).sqrt();
        assert!(
            (mag - 800.0).abs() < 1.0,
            "expected the softened |strength| = 800, got {mag}; unsoftened this is 80000"
        );
        assert!(f[0] < 0.0, "and it must still point away from the other body, got {f:?}");
    }

    /// The coincident-body jitter is keyed on the *index pair*, so each member
    /// of a degenerate cluster is pushed a different way and the cluster comes
    /// apart. Replacing the key with a constant also passes every other test
    /// here: the forces stay finite and bounded, but become identical, so the
    /// cluster translates rigidly and never separates — which looks, on screen,
    /// exactly like nodes that have fused.
    #[test]
    fn coincident_bodies_are_pushed_in_different_directions() {
        let pos = vec![[5.0, 5.0], [5.0, 5.0]];
        let t = QuadTree::build(&pos);
        let f0 = t.repulsion(0, &pos, 0.9, -100.0);
        let f1 = t.repulsion(1, &pos, 0.9, -100.0);
        let spread = ((f0[0] - f1[0]).powi(2) + (f0[1] - f1[1]).powi(2)).sqrt();
        assert!(
            spread > 1.0,
            "bodies got near-identical pushes {f0:?} and {f1:?}; a rigid cluster never separates"
        );
    }

    #[test]
    fn deeply_clustered_bodies_terminate() {
        // All-but-identical points would subdivide forever without a depth cap.
        let pos: Vec<[f32; 2]> = (0..50).map(|i| [1.0 + i as f32 * 1e-7, 1.0]).collect();
        let t = QuadTree::build(&pos);
        let f = t.repulsion(0, &pos, 0.9, -100.0);
        assert!(f[0].is_finite() && f[1].is_finite());
    }

    #[test]
    fn an_empty_set_builds_and_answers_nothing() {
        let pos: Vec<[f32; 2]> = vec![];
        let t = QuadTree::build(&pos);
        assert_eq!(t.repulsion(0, &pos, 0.9, -100.0), [0.0, 0.0]);
    }

    #[test]
    fn repulsion_pushes_apart_and_attraction_pulls_together() {
        // Sign convention: negative strength repels, which is d3's.
        let pos = vec![[0.0, 0.0], [10.0, 0.0]];
        let t = QuadTree::build(&pos);
        let f = t.repulsion(0, &pos, 0.9, -100.0);
        assert!(f[0] < 0.0, "body 0 should be pushed away from body 1 (leftward), got {f:?}");
        let a = t.repulsion(0, &pos, 0.9, 100.0);
        assert!(a[0] > 0.0, "positive strength should attract, got {a:?}");
    }

    /// The force must fall off as `1/d`, not `1/d²` — doubling the separation
    /// must halve the force. This is the property the `theta = 0.9` accuracy
    /// budget depends on, and it is invisible to every other test here because
    /// they all compare the tree against a reference that shares its law.
    #[test]
    fn force_magnitude_falls_off_as_inverse_distance() {
        let near = vec![[0.0, 0.0], [10.0, 0.0]];
        let far = vec![[0.0, 0.0], [20.0, 0.0]];
        let fn_ = QuadTree::build(&near).repulsion(0, &near, 0.0, -100.0)[0].abs();
        let ff = QuadTree::build(&far).repulsion(0, &far, 0.0, -100.0)[0].abs();
        assert!((fn_ - 10.0).abs() < 1e-3, "expected |strength|/d = 10, got {fn_}");
        assert!((ff - 5.0).abs() < 1e-3, "expected |strength|/d = 5, got {ff}");
    }
}