use std::hash::{Hash, Hasher};
pub const BH_THRESHOLD: usize = 2_000;
pub const BH_THETA: f32 = 0.7;
const MAX_DEPTH: usize = 48;
const MIN_DIST: f32 = 1e-3;
const MAX_CH: usize = 8;
#[derive(Clone)]
struct Cell<const D: usize> {
center: [f32; D],
half: f32,
com: [f32; D],
mass: f32,
body: i32,
children: [i32; MAX_CH],
leaf: bool,
}
impl<const D: usize> Cell<D> {
fn new(center: [f32; D], half: f32) -> Self {
Self { center, half, com: [0.0; D], mass: 0.0, body: -1, children: [-1; MAX_CH], leaf: true }
}
}
#[inline]
fn child_index<const D: usize>(center: [f32; D], pb: [f32; D]) -> usize {
let mut idx = 0usize;
for d in 0..D {
if pb[d] >= center[d] {
idx |= 1 << d;
}
}
idx
}
fn arena_new_child<const D: usize>(cells: &mut Vec<Cell<D>>, pc: [f32; D], ph: f32, idx: usize) -> usize {
let ch = ph * 0.5;
let mut center = [0.0f32; D];
for d in 0..D {
let hi = (idx >> d) & 1 == 1;
center[d] = pc[d] + if hi { ch } else { -ch };
}
cells.push(Cell::new(center, ch));
cells.len() - 1
}
fn arena_insert<const D: usize>(
cells: &mut Vec<Cell<D>>,
all_pos: &[[f32; D]],
ci: usize,
b: usize,
pb: [f32; D],
depth: usize,
) {
let m = cells[ci].mass;
for d in 0..D {
cells[ci].com[d] = (cells[ci].com[d] * m + pb[d]) / (m + 1.0);
}
cells[ci].mass = m + 1.0;
if cells[ci].mass == 1.0 {
cells[ci].body = b as i32;
cells[ci].leaf = true;
return;
}
if depth >= MAX_DEPTH || cells[ci].half <= MIN_DIST {
cells[ci].body = -1;
cells[ci].leaf = true;
return;
}
if cells[ci].leaf && cells[ci].body >= 0 {
let ob = cells[ci].body as usize;
let obp = all_pos[ob];
cells[ci].body = -1;
cells[ci].leaf = false;
arena_route(cells, all_pos, ci, ob, obp, depth);
}
arena_route(cells, all_pos, ci, b, pb, depth);
}
fn arena_route<const D: usize>(
cells: &mut Vec<Cell<D>>,
all_pos: &[[f32; D]],
ci: usize,
b: usize,
pb: [f32; D],
depth: usize,
) {
let idx = child_index(cells[ci].center, pb);
let existing = cells[ci].children[idx];
let child = if existing < 0 {
let (pc, ph) = (cells[ci].center, cells[ci].half);
let nc = arena_new_child(cells, pc, ph, idx);
cells[ci].children[idx] = nc as i32;
nc
} else {
existing as usize
};
arena_insert(cells, all_pos, child, b, pb, depth + 1);
}
pub struct NBodyTree<const D: usize> {
arenas: Vec<Vec<Cell<D>>>,
}
impl<const D: usize> NBodyTree<D> {
#[inline]
fn nch() -> usize {
1usize << D
}
fn root_bounds(pos: &[[f32; D]]) -> ([f32; D], f32) {
let mut mn = [f32::MAX; D];
let mut mx = [f32::MIN; D];
for p in pos {
for d in 0..D {
mn[d] = mn[d].min(p[d]);
mx[d] = mx[d].max(p[d]);
}
}
let mut span = 0.0f32;
let mut center = [0.0f32; D];
for d in 0..D {
center[d] = 0.5 * (mn[d] + mx[d]);
span = span.max(mx[d] - mn[d]);
}
(center, (span * 0.5).max(MIN_DIST) * 1.0001)
}
#[must_use]
pub fn build(pos: &[[f32; D]]) -> Self {
if pos.is_empty() {
return Self { arenas: Vec::new() };
}
let (center, half) = Self::root_bounds(pos);
let mut cells = vec![Cell::new(center, half)];
for (i, &p) in pos.iter().enumerate() {
arena_insert(&mut cells, pos, 0, i, p, 0);
}
Self { arenas: vec![cells] }
}
#[must_use]
pub fn build_parallel(pos: &[[f32; D]]) -> Self {
let n = pos.len();
if n < 4_096 {
return Self::build(pos);
}
let (root_center, root_half) = Self::root_bounds(pos);
let nch = 1usize << D;
let want = znippy_zoomies::gatling_forkjoin::default_workers().saturating_mul(4).max(nch);
let mut pd = 1usize;
while nch.pow((pd + 1) as u32) <= want && pd < 4 {
pd += 1;
}
let buckets = nch.pow(pd as u32);
let mut bucket_of = vec![0u32; n];
for (i, &p) in pos.iter().enumerate() {
let mut c = root_center;
let mut h = root_half;
let mut idx = 0usize;
for _ in 0..pd {
h *= 0.5;
let ci = child_index(c, p);
for d in 0..D {
c[d] += if (ci >> d) & 1 == 1 { h } else { -h };
}
idx = idx * nch + ci;
}
bucket_of[i] = idx as u32;
}
let mut groups: Vec<Vec<u32>> = vec![Vec::new(); buckets];
for (i, &b) in bucket_of.iter().enumerate() {
groups[b as usize].push(i as u32);
}
let bucket_cell = |b: usize| -> ([f32; D], f32) {
let mut c = root_center;
let mut h = root_half;
for l in 0..pd {
h *= 0.5;
let shift = (pd - 1 - l) * D;
let ci = (b >> shift) & (nch - 1);
for d in 0..D {
c[d] += if (ci >> d) & 1 == 1 { h } else { -h };
}
}
(c, h)
};
let arenas: Vec<Vec<Cell<D>>> = znippy_zoomies::gatling_forkjoin::gatling_for_each(buckets, 0, |b| {
let bodies = &groups[b];
if bodies.is_empty() {
return Vec::new();
}
let (center, half) = bucket_cell(b);
let mut cells = vec![Cell::new(center, half)];
for &gi in bodies {
let gi = gi as usize;
arena_insert(&mut cells, pos, 0, gi, pos[gi], pd);
}
cells
});
Self { arenas }
}
#[must_use]
pub fn force(&self, p: [f32; D], self_idx: usize, k: f32, theta: f32) -> [f32; D] {
let mut acc = [0.0f32; D];
let k2 = k * k;
let theta2 = theta * theta;
const CAP: usize = MAX_DEPTH * MAX_CH + 16;
let mut stack: [usize; CAP] = [0; CAP];
for cells in &self.arenas {
if cells.is_empty() {
continue;
}
stack[0] = 0; let mut top = 1usize;
while top > 0 {
top -= 1;
let ci = stack[top];
let c = &cells[ci];
if c.mass == 0.0 {
continue;
}
if c.leaf && c.body == self_idx as i32 {
continue;
}
let mut dv = [0.0f32; D];
let mut d2 = 0.0f32;
for d in 0..D {
let dd = p[d] - c.com[d];
dv[d] = dd;
d2 += dd * dd;
}
let width = c.half * 2.0;
if c.leaf || width * width < theta2 * d2 {
let dist = d2.sqrt().max(MIN_DIST);
let f = k2 / (dist * dist) * c.mass; for d in 0..D {
acc[d] += dv[d] * f;
}
} else {
for ch in 0..Self::nch() {
let cc = c.children[ch];
if cc >= 0 {
debug_assert!(top < CAP, "Barnes–Hut traversal stack overflow");
stack[top] = cc as usize;
top += 1;
}
}
}
}
}
acc
}
}
#[must_use]
pub fn repulsion_forces<const D: usize>(pos: &[[f32; D]], k: f32, theta: f32) -> Vec<[f32; D]> {
let n = pos.len();
if n == 0 {
return Vec::new();
}
let tree = NBodyTree::<D>::build_parallel(pos);
znippy_zoomies::gatling_forkjoin::gatling_for_each(n, 0, |i| tree.force(pos[i], i, k, theta))
}
#[must_use]
pub fn repulsion_forces_serial<const D: usize>(pos: &[[f32; D]], k: f32, theta: f32) -> Vec<[f32; D]> {
let n = pos.len();
if n == 0 {
return Vec::new();
}
let tree = NBodyTree::<D>::build(pos);
(0..n).map(|i| tree.force(pos[i], i, k, theta)).collect()
}
#[must_use]
pub fn fr_step<const D: usize>(
pos: &[[f32; D]],
adj: &[Vec<u32>],
k: f32,
theta: f32,
temp: f32,
) -> Vec<[f32; D]> {
let n = pos.len();
if n == 0 {
return Vec::new();
}
let tree = NBodyTree::<D>::build_parallel(pos);
znippy_zoomies::gatling_forkjoin::gatling_for_each(n, 0, |i| {
let pi = pos[i];
let mut disp = tree.force(pi, i, k, theta);
for &jn in &adj[i] {
let j = jn as usize;
if j >= n {
continue;
}
let pj = pos[j];
let mut d2 = 0.0f32;
let mut dv = [0.0f32; D];
for d in 0..D {
let dd = pi[d] - pj[d];
dv[d] = dd;
d2 += dd * dd;
}
let dist = d2.sqrt().max(MIN_DIST);
let f = dist / k;
for d in 0..D {
disp[d] -= dv[d] * f;
}
}
let mut dl = 0.0f32;
for d in 0..D {
dl += disp[d] * disp[d];
}
let dl = dl.sqrt().max(MIN_DIST);
let cap = dl.min(temp) / dl;
let mut np = pi;
for d in 0..D {
np[d] += disp[d] * cap;
}
np
})
}
#[must_use]
pub fn fr_step_scatter<const D: usize>(
pos: &[[f32; D]],
edges: &[(usize, usize)],
k: f32,
theta: f32,
temp: f32,
) -> Vec<[f32; D]> {
let n = pos.len();
if n == 0 {
return Vec::new();
}
let rep = repulsion_forces::<D>(pos, k, theta);
let attr = znippy_zoomies::gatling_forkjoin::gatling_reduce(
edges.len(),
0,
|| vec![[0.0f32; D]; n],
|buf, e| {
let (a, b) = edges[e];
if a >= n || b >= n || a == b {
return;
}
let mut d2 = 0.0f32;
let mut dv = [0.0f32; D];
for d in 0..D {
let dd = pos[a][d] - pos[b][d];
dv[d] = dd;
d2 += dd * dd;
}
let dist = d2.sqrt().max(MIN_DIST);
let f = dist / k; for d in 0..D {
let push = dv[d] * f;
buf[a][d] -= push;
buf[b][d] += push;
}
},
|acc, part| {
for i in 0..n {
for d in 0..D {
acc[i][d] += part[i][d];
}
}
},
);
znippy_zoomies::gatling_forkjoin::gatling_for_each(n, 0, |i| {
let mut disp = rep[i];
for d in 0..D {
disp[d] += attr[i][d];
}
let mut dl = 0.0f32;
for d in 0..D {
dl += disp[d] * disp[d];
}
let dl = dl.sqrt().max(MIN_DIST);
let cap = dl.min(temp) / dl;
let mut np = pos[i];
for d in 0..D {
np[d] += disp[d] * cap;
}
np
})
}
#[must_use]
pub fn adjacency(n: usize, edges: &[(usize, usize)]) -> Vec<Vec<u32>> {
let mut adj = vec![Vec::new(); n];
for &(a, b) in edges {
if a >= n || b >= n || a == b {
continue;
}
adj[a].push(b as u32);
adj[b].push(a as u32);
}
adj
}
#[must_use]
pub fn fr_step_serial<const D: usize>(
pos: &[[f32; D]],
adj: &[Vec<u32>],
k: f32,
theta: f32,
temp: f32,
) -> Vec<[f32; D]> {
let n = pos.len();
if n == 0 {
return Vec::new();
}
let mut disp = if n >= BH_THRESHOLD {
repulsion_forces_serial::<D>(pos, k, theta)
} else {
let mut d = vec![[0.0f32; D]; n];
for i in 0..n {
for j in (i + 1)..n {
let mut dv = [0.0f32; D];
let mut d2 = 0.0f32;
for c in 0..D {
let dd = pos[i][c] - pos[j][c];
dv[c] = dd;
d2 += dd * dd;
}
let dist = d2.sqrt().max(MIN_DIST);
let f = k * k / (dist * dist);
for c in 0..D {
let push = dv[c] * f;
d[i][c] += push;
d[j][c] -= push;
}
}
}
d
};
for i in 0..n {
let pi = pos[i];
for &jn in &adj[i] {
let j = jn as usize;
if j >= n {
continue;
}
let pj = pos[j];
let mut d2 = 0.0f32;
let mut dv = [0.0f32; D];
for c in 0..D {
let dd = pi[c] - pj[c];
dv[c] = dd;
d2 += dd * dd;
}
let dist = d2.sqrt().max(MIN_DIST);
let f = dist / k;
for c in 0..D {
disp[i][c] -= dv[c] * f;
}
}
}
let mut out = pos.to_vec();
for i in 0..n {
let mut dl = 0.0f32;
for c in 0..D {
dl += disp[i][c] * disp[i][c];
}
let dl = dl.sqrt().max(MIN_DIST);
let cap = dl.min(temp) / dl;
for c in 0..D {
out[i][c] += disp[i][c] * cap;
}
}
out
}
pub fn normalize_box(pos: &mut [[f32; 2]]) {
let (mut x0, mut y0, mut x1, mut y1) = (f32::MAX, f32::MAX, f32::MIN, f32::MIN);
for p in pos.iter() {
x0 = x0.min(p[0]);
y0 = y0.min(p[1]);
x1 = x1.max(p[0]);
y1 = y1.max(p[1]);
}
let span = (x1 - x0).max(y1 - y0).max(1e-3);
let (cx, cy) = ((x0 + x1) * 0.5, (y0 + y1) * 0.5);
for p in pos.iter_mut() {
p[0] = (p[0] - cx) / span * 1.8;
p[1] = (p[1] - cy) / span * 1.8;
}
}
pub fn relax_new<const D: usize>(
pos: &mut [[f32; D]],
edges: &[(usize, usize)],
new_lo: usize,
k: f32,
theta: f32,
iters: usize,
step: f32,
) {
let n = pos.len();
if n == 0 || new_lo >= n {
return;
}
for _ in 0..iters {
let rep = repulsion_forces::<D>(pos, k, theta);
let mut disp = vec![[0.0f32; D]; n];
for i in new_lo..n {
disp[i] = rep[i];
}
for &(a, b) in edges {
if a >= n || b >= n || a == b {
continue;
}
let a_new = a >= new_lo;
let b_new = b >= new_lo;
if !a_new && !b_new {
continue;
}
let mut dv = [0.0f32; D];
let mut d2 = 0.0f32;
for d in 0..D {
let dd = pos[a][d] - pos[b][d];
dv[d] = dd;
d2 += dd * dd;
}
let dist = d2.sqrt().max(MIN_DIST);
let f = dist / k; for d in 0..D {
let push = dv[d] * f;
if a_new {
disp[a][d] -= push;
}
if b_new {
disp[b][d] += push;
}
}
}
for i in new_lo..n {
let mut dl = 0.0f32;
for d in 0..D {
dl += disp[i][d] * disp[i][d];
}
let dl = dl.sqrt().max(MIN_DIST);
let cap = dl.min(step) / dl;
for d in 0..D {
pos[i][d] += disp[i][d] * cap;
}
}
}
}
#[must_use]
pub fn structure_sig(layout_key: &str, n: usize, edges: &[(usize, usize)]) -> u64 {
let mut h = std::collections::hash_map::DefaultHasher::new();
layout_key.hash(&mut h);
n.hash(&mut h);
edges.len().hash(&mut h);
for e in edges {
e.hash(&mut h);
}
h.finish()
}
#[cfg(test)]
mod tests {
use super::*;
fn cloud2(n: usize) -> Vec<[f32; 2]> {
(0..n)
.map(|i| {
let a = std::f32::consts::TAU * i as f32 / n as f32;
let r = 1.0 + (i as f32 * 0.37).sin() * 0.5;
[r * a.cos(), r * a.sin()]
})
.collect()
}
fn naive2(pos: &[[f32; 2]], k: f32) -> Vec<[f32; 2]> {
let n = pos.len();
let k2 = k * k;
let mut out = vec![[0.0f32; 2]; n];
for i in 0..n {
for j in 0..n {
if i == j {
continue;
}
let dx = pos[i][0] - pos[j][0];
let dy = pos[i][1] - pos[j][1];
let dist = (dx * dx + dy * dy).sqrt().max(MIN_DIST);
let f = k2 / (dist * dist);
out[i][0] += dx * f;
out[i][1] += dy * f;
}
}
out
}
#[test]
fn bh_theta_zero_matches_naive_2d() {
let pos = cloud2(400);
let k = (4.0 / pos.len() as f32).sqrt();
let bh = repulsion_forces::<2>(&pos, k, 0.0);
let naive = naive2(&pos, k);
let mut max_err = 0.0f32;
for i in 0..pos.len() {
for d in 0..2 {
max_err = max_err.max((bh[i][d] - naive[i][d]).abs());
}
}
assert!(max_err < 1e-3, "theta=0 Barnes–Hut must equal naive all-pairs; max_err={max_err}");
}
#[test]
fn bh_theta_moderate_close_to_naive_2d() {
let pos = cloud2(600);
let k = (4.0 / pos.len() as f32).sqrt();
let bh = repulsion_forces::<2>(&pos, k, 0.7);
let naive = naive2(&pos, k);
let mut sum_rel = 0.0f32;
let mut cnt = 0.0f32;
for i in 0..pos.len() {
let bm = (bh[i][0] * bh[i][0] + bh[i][1] * bh[i][1]).sqrt();
let nm = (naive[i][0] * naive[i][0] + naive[i][1] * naive[i][1]).sqrt();
if nm > 1e-4 {
sum_rel += (bm - nm).abs() / nm;
cnt += 1.0;
}
}
let mean_rel = sum_rel / cnt.max(1.0);
assert!(mean_rel < 0.15, "theta=0.7 stays within ~15% of exact; mean_rel={mean_rel}");
}
#[test]
fn repulsion_is_deterministic() {
let pos = cloud2(1500);
let k = (4.0 / pos.len() as f32).sqrt();
let a = repulsion_forces::<2>(&pos, k, 0.7);
let b = repulsion_forces::<2>(&pos, k, 0.7);
assert_eq!(a, b, "parallel Barnes–Hut repulsion is deterministic");
}
#[test]
fn parallel_build_matches_serial_build() {
let pos = cloud2(8_000);
let k = (4.0 / pos.len() as f32).sqrt();
let serial = NBodyTree::<2>::build(&pos);
let parallel = NBodyTree::<2>::build_parallel(&pos);
let mut max_err = 0.0f32;
for i in 0..pos.len() {
let fs = serial.force(pos[i], i, k, 0.0);
let fp = parallel.force(pos[i], i, k, 0.0);
for d in 0..2 {
max_err = max_err.max((fs[d] - fp[d]).abs());
}
}
assert!(max_err < 1e-3, "forest build matches serial build at theta=0; max_err={max_err}");
let mass: f32 = parallel.arenas.iter().filter(|a| !a.is_empty()).map(|a| a[0].mass).sum();
assert!((mass - pos.len() as f32).abs() < 0.5, "forest holds every body; mass={mass}");
}
#[test]
fn scatter_reduce_step_matches_gather_step() {
let n = 6_000usize;
let mut edges = Vec::new();
for t in 0..n / 3 {
let (a, b, c) = (t * 3, t * 3 + 1, t * 3 + 2);
edges.push((a, b));
edges.push((b, c));
edges.push((c, a));
edges.push((a, (a + 3) % n));
}
let pos: Vec<[f32; 2]> = (0..n)
.map(|i| {
let ang = std::f32::consts::TAU * i as f32 / n as f32;
[ang.cos(), ang.sin()]
})
.collect();
let k = (4.0 / n as f32).sqrt();
let adj = adjacency(n, &edges);
let gather = fr_step::<2>(&pos, &adj, k, BH_THETA, 0.05);
let scatter = fr_step_scatter::<2>(&pos, &edges, k, BH_THETA, 0.05);
let mut max_err = 0.0f32;
for (g, s) in gather.iter().zip(&scatter) {
max_err = max_err.max((g[0] - s[0]).abs()).max((g[1] - s[1]).abs());
}
assert!(max_err < 1e-4, "scatter-reduce == gather attraction; max_err={max_err}");
}
#[test]
fn parallel_build_is_deterministic() {
let pos = cloud2(6_000);
let k = (4.0 / pos.len() as f32).sqrt();
let a = repulsion_forces::<2>(&pos, k, 0.7);
let t = NBodyTree::<2>::build_parallel(&pos);
let b: Vec<[f32; 2]> = (0..pos.len()).map(|i| t.force(pos[i], i, k, 0.7)).collect();
assert_eq!(a, b, "forest build + query is deterministic");
}
#[test]
fn bh_theta_zero_matches_naive_3d() {
let n = 300;
let golden = std::f32::consts::PI * (3.0 - 5.0_f32.sqrt());
let pos: Vec<[f32; 3]> = (0..n)
.map(|i| {
let y = 1.0 - (i as f32 / (n - 1) as f32) * 2.0;
let r = (1.0 - y * y).max(0.0).sqrt();
let th = golden * i as f32;
[th.cos() * r, y, th.sin() * r]
})
.collect();
let k = 0.3;
let k2 = k * k;
let bh = repulsion_forces::<3>(&pos, k, 0.0);
let mut max_err = 0.0f32;
for i in 0..n {
let mut naive = [0.0f32; 3];
for j in 0..n {
if i == j {
continue;
}
let mut dv = [0.0f32; 3];
let mut d2 = 0.0f32;
for d in 0..3 {
dv[d] = pos[i][d] - pos[j][d];
d2 += dv[d] * dv[d];
}
let dist = d2.sqrt().max(MIN_DIST);
let f = k2 / (dist * dist);
for d in 0..3 {
naive[d] += dv[d] * f;
}
}
for d in 0..3 {
max_err = max_err.max((bh[i][d] - naive[d]).abs());
}
}
assert!(max_err < 1e-3, "3-D theta=0 octree must equal naive; max_err={max_err}");
}
#[test]
fn relax_new_freezes_existing_nodes() {
let mut pos = cloud2(50);
let frozen = pos.clone();
let old_n = pos.len();
for i in 0..5 {
pos.push([0.01 * i as f32, 0.02 * i as f32]);
}
let edges: Vec<(usize, usize)> = (0..5).map(|i| (old_n + i, i)).collect();
let k = (4.0 / pos.len() as f32).sqrt();
relax_new::<2>(&mut pos, &edges, old_n, k, 0.7, 10, 0.05);
for i in 0..old_n {
assert_eq!(pos[i], frozen[i], "existing node {i} stays frozen during absorb");
}
for i in old_n..pos.len() {
assert!(pos[i][0].is_finite() && pos[i][1].is_finite(), "new node {i} is finite");
}
}
#[test]
fn structure_sig_tracks_changes() {
let e1 = vec![(0usize, 1usize), (1, 2)];
let e2 = vec![(0usize, 1usize), (1, 2), (2, 3)];
assert_eq!(structure_sig("force", 3, &e1), structure_sig("force", 3, &e1));
assert_ne!(structure_sig("force", 3, &e1), structure_sig("force", 4, &e2));
assert_ne!(structure_sig("force", 3, &e1), structure_sig("circular", 3, &e1));
}
}