use crate::particle::Particle;
pub const DEFAULT_THETA: f32 = 0.6;
pub const BRUTE_FORCE_THRESHOLD: usize = 500;
pub const MIN_DIST2: f32 = 1.0;
pub const MIN_SPLIT_DIST2: f32 = 1e-8;
pub const MIN_CELL_SIZE: f32 = 1e-3;
const CELL_ACCEPTANCE_MIN_DIST: f32 = 0.001;
pub fn apply_repulsion_brute_force(particles: &[Particle], strength: f32, min_dist2: f32, masses: Option<&[f32]>, out: &mut [(f32, f32)]) {
let n = particles.len();
let mass_of = |i: usize| masses.and_then(|m| m.get(i)).copied().unwrap_or(1.0);
for i in 0..n {
let (xi, yi) = (particles[i].x, particles[i].y);
let mass_i = mass_of(i);
for j in (i + 1)..n {
let dx = xi - particles[j].x;
let dy = yi - particles[j].y;
let d2 = (dx * dx + dy * dy).max(min_dist2);
let d = d2.sqrt();
let mass_j = mass_of(j);
let f_on_i = strength * mass_j / d2;
let f_on_j = strength * mass_i / d2;
out[i].0 += dx / d * f_on_i;
out[i].1 += dy / d * f_on_i;
out[j].0 -= dx / d * f_on_j;
out[j].1 -= dy / d * f_on_j;
}
}
}
#[derive(Clone, Copy, Debug)]
struct QuadBounds {
min_x: f32,
min_y: f32,
size: f32,
}
impl QuadBounds {
fn quadrant(&self, x: f32, y: f32) -> usize {
let mid_x = self.min_x + self.size * 0.5;
let mid_y = self.min_y + self.size * 0.5;
match (x >= mid_x, y >= mid_y) {
(false, false) => 0,
(true, false) => 1,
(false, true) => 2,
(true, true) => 3,
}
}
fn child_bounds(&self, quadrant: usize) -> QuadBounds {
let half = self.size * 0.5;
let (ox, oy) = match quadrant {
0 => (0.0, 0.0),
1 => (half, 0.0),
2 => (0.0, half),
_ => (half, half),
};
QuadBounds { min_x: self.min_x + ox, min_y: self.min_y + oy, size: half }
}
fn intersects_circle(&self, cx: f32, cy: f32, r: f32) -> bool {
let max_x = self.min_x + self.size;
let max_y = self.min_y + self.size;
let nearest_x = cx.clamp(self.min_x, max_x);
let nearest_y = cy.clamp(self.min_y, max_y);
let dx = cx - nearest_x;
let dy = cy - nearest_y;
dx * dx + dy * dy <= r * r
}
}
enum NodeContent {
Empty,
Leaf { x: f32, y: f32, mass: f32, entries: Vec<(u32, f32)> },
Internal { children: Box<[QuadNode; 4]> },
}
struct QuadNode {
bounds: QuadBounds,
mass: f32,
com_x: f32,
com_y: f32,
content: NodeContent,
}
impl QuadNode {
fn new_empty(bounds: QuadBounds) -> Self {
Self { bounds, mass: 0.0, com_x: 0.0, com_y: 0.0, content: NodeContent::Empty }
}
fn insert(&mut self, x: f32, y: f32, mass: f32, index: u32, min_split_dist2: f32, min_cell_size: f32) {
let new_mass = self.mass + mass;
self.com_x = (self.com_x * self.mass + x * mass) / new_mass;
self.com_y = (self.com_y * self.mass + y * mass) / new_mass;
self.mass = new_mass;
match &mut self.content {
NodeContent::Empty => {
self.content = NodeContent::Leaf { x, y, mass, entries: vec![(index, mass)] };
}
NodeContent::Leaf { x: lx, y: ly, mass: lmass, entries } => {
let (dx, dy) = (x - *lx, y - *ly);
if dx * dx + dy * dy <= min_split_dist2 || self.bounds.size <= min_cell_size {
*lmass += mass;
entries.push((index, mass));
return;
}
let (lx, ly) = (*lx, *ly);
let prior_entries = std::mem::take(entries);
let mut children = [
QuadNode::new_empty(self.bounds.child_bounds(0)),
QuadNode::new_empty(self.bounds.child_bounds(1)),
QuadNode::new_empty(self.bounds.child_bounds(2)),
QuadNode::new_empty(self.bounds.child_bounds(3)),
];
let prior_quadrant = self.bounds.quadrant(lx, ly);
for (prior_index, prior_mass) in prior_entries {
children[prior_quadrant].insert(lx, ly, prior_mass, prior_index, min_split_dist2, min_cell_size);
}
children[self.bounds.quadrant(x, y)].insert(x, y, mass, index, min_split_dist2, min_cell_size);
self.content = NodeContent::Internal { children: Box::new(children) };
}
NodeContent::Internal { children } => {
children[self.bounds.quadrant(x, y)].insert(x, y, mass, index, min_split_dist2, min_cell_size);
}
}
}
fn accumulate(&self, x: f32, y: f32, theta: f32, strength: f32, min_dist2: f32, out: &mut (f32, f32)) {
match &self.content {
NodeContent::Empty => {}
NodeContent::Leaf { x: lx, y: ly, mass, .. } => {
apply_point(x, y, *lx, *ly, *mass, strength, min_dist2, out);
}
NodeContent::Internal { children } => {
let dx = x - self.com_x;
let dy = y - self.com_y;
let d = (dx * dx + dy * dy).sqrt().max(CELL_ACCEPTANCE_MIN_DIST);
if self.bounds.size / d < theta {
apply_point(x, y, self.com_x, self.com_y, self.mass, strength, min_dist2, out);
} else {
for child in children.iter() {
child.accumulate(x, y, theta, strength, min_dist2, out);
}
}
}
}
}
fn collect_within(&self, qx: f32, qy: f32, r: f32, out: &mut Vec<u32>) {
match &self.content {
NodeContent::Empty => {}
NodeContent::Leaf { entries, .. } => out.extend(entries.iter().map(|(idx, _)| *idx)),
NodeContent::Internal { children } => {
for child in children.iter() {
if child.bounds.intersects_circle(qx, qy, r) {
child.collect_within(qx, qy, r, out);
}
}
}
}
}
}
fn apply_point(x: f32, y: f32, px: f32, py: f32, mass: f32, strength: f32, min_dist2: f32, out: &mut (f32, f32)) {
let dx = x - px;
let dy = y - py;
let d2 = (dx * dx + dy * dy).max(min_dist2);
let d = d2.sqrt();
let f = strength * mass / d2;
out.0 += dx / d * f;
out.1 += dy / d * f;
}
pub struct Quadtree {
root: Option<QuadNode>,
}
impl Quadtree {
pub fn build(particles: &[Particle], min_split_dist2: f32, min_cell_size: f32) -> Self {
Self::build_weighted(particles, None, min_split_dist2, min_cell_size)
}
pub fn build_weighted(particles: &[Particle], masses: Option<&[f32]>, min_split_dist2: f32, min_cell_size: f32) -> Self {
if particles.is_empty() {
return Self { root: None };
}
let (mut min_x, mut min_y, mut max_x, mut max_y) = (f32::MAX, f32::MAX, f32::MIN, f32::MIN);
for p in particles {
min_x = min_x.min(p.x);
max_x = max_x.max(p.x);
min_y = min_y.min(p.y);
max_y = max_y.max(p.y);
}
let size = (max_x - min_x).max(max_y - min_y).max(1.0) * 1.01;
let bounds = QuadBounds { min_x, min_y, size };
let mut root = QuadNode::new_empty(bounds);
for (i, p) in particles.iter().enumerate() {
let mass = masses.and_then(|m| m.get(i)).copied().unwrap_or(1.0);
root.insert(p.x, p.y, mass, i as u32, min_split_dist2, min_cell_size);
}
Self { root: Some(root) }
}
pub fn accumulate_forces(&self, particles: &[Particle], theta: f32, strength: f32, min_dist2: f32, out: &mut [(f32, f32)]) {
let Some(root) = &self.root else { return };
for (i, p) in particles.iter().enumerate() {
let mut acc = (0.0, 0.0);
root.accumulate(p.x, p.y, theta, strength, min_dist2, &mut acc);
out[i].0 += acc.0;
out[i].1 += acc.1;
}
}
pub fn apply_collision(&self, particles: &[Particle], radii: &[f32], strength: f32, out: &mut [(f32, f32)]) {
let Some(root) = &self.root else { return };
let n = particles.len();
if n < 2 {
return;
}
let max_radius = radii.iter().copied().fold(1.0f32, f32::max);
let mut candidates: Vec<u32> = Vec::new();
for i in 0..n {
let query_radius = radii.get(i).copied().unwrap_or(1.0) + max_radius;
candidates.clear();
root.collect_within(particles[i].x, particles[i].y, query_radius, &mut candidates);
for &j_u32 in &candidates {
let j = j_u32 as usize;
if j <= i {
continue;
}
collision_pair_force(i, j, particles, radii, strength, out);
}
}
}
}
pub(crate) fn collision_pair_force(i: usize, j: usize, particles: &[Particle], radii: &[f32], strength: f32, out: &mut [(f32, f32)]) {
let dx = particles[j].x - particles[i].x;
let dy = particles[j].y - particles[i].y;
let dist2 = dx * dx + dy * dy;
let min_dist = radii.get(i).copied().unwrap_or(1.0) + radii.get(j).copied().unwrap_or(1.0);
if dist2 <= 1e-6 {
out[i].0 -= 0.5;
out[j].0 += 0.5;
return;
}
if dist2 < min_dist * min_dist {
let dist = dist2.sqrt();
let overlap = (min_dist - dist) * strength;
let nx = dx / dist;
let ny = dy / dist;
out[i].0 -= nx * overlap * 0.5;
out[i].1 -= ny * overlap * 0.5;
out[j].0 += nx * overlap * 0.5;
out[j].1 += ny * overlap * 0.5;
}
}
#[cfg(test)]
mod tests {
use super::*;
fn deterministic_particles(n: usize) -> Vec<Particle> {
let mut particles = Vec::with_capacity(n);
let mut state: u64 = 0x9E3779B97F4A7C15;
for _ in 0..n {
state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let rx = ((state >> 33) as u32 % 2000) as f32 - 1000.0;
state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let ry = ((state >> 33) as u32 % 2000) as f32 - 1000.0;
particles.push(Particle::at(rx, ry));
}
particles
}
#[test]
fn barnes_hut_matches_brute_force_within_tolerance() {
let particles = deterministic_particles(96);
let strength = 400.0;
let mut brute = vec![(0f32, 0f32); particles.len()];
apply_repulsion_brute_force(&particles, strength, MIN_DIST2, None, &mut brute);
let qt = Quadtree::build(&particles, MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let mut approx = vec![(0f32, 0f32); particles.len()];
qt.accumulate_forces(&particles, 0.6, strength, MIN_DIST2, &mut approx);
let mut max_rel_err = 0f32;
for (b, a) in brute.iter().zip(approx.iter()) {
let bmag = (b.0 * b.0 + b.1 * b.1).sqrt();
let diff = ((a.0 - b.0).powi(2) + (a.1 - b.1).powi(2)).sqrt();
if bmag > 1e-3 {
max_rel_err = max_rel_err.max(diff / bmag);
}
}
assert!(max_rel_err < 0.35, "Barnes-Hut relative error too high: {max_rel_err}");
}
#[test]
fn barnes_hut_matches_brute_force_within_tolerance_at_the_shipped_theta() {
let particles = deterministic_particles(96);
let strength = 400.0;
let mut brute = vec![(0f32, 0f32); particles.len()];
apply_repulsion_brute_force(&particles, strength, MIN_DIST2, None, &mut brute);
let qt = Quadtree::build(&particles, MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let mut approx = vec![(0f32, 0f32); particles.len()];
qt.accumulate_forces(&particles, DEFAULT_THETA, strength, MIN_DIST2, &mut approx);
let mut max_rel_err = 0f32;
for (b, a) in brute.iter().zip(approx.iter()) {
let bmag = (b.0 * b.0 + b.1 * b.1).sqrt();
let diff = ((a.0 - b.0).powi(2) + (a.1 - b.1).powi(2)).sqrt();
if bmag > 1e-3 {
max_rel_err = max_rel_err.max(diff / bmag);
}
}
assert!(max_rel_err < 0.35, "Barnes-Hut relative error at the SHIPPED theta={DEFAULT_THETA} too high: {max_rel_err}");
}
#[test]
fn empty_quadtree_produces_no_force() {
let particles: Vec<Particle> = Vec::new();
let qt = Quadtree::build(&particles, MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let mut out: Vec<(f32, f32)> = Vec::new();
qt.accumulate_forces(&particles, DEFAULT_THETA, 100.0, MIN_DIST2, &mut out);
assert!(out.is_empty());
}
#[test]
fn single_particle_produces_no_self_force() {
let particles = vec![Particle::at(3.0, 4.0)];
let qt = Quadtree::build(&particles, MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let mut out = vec![(0f32, 0f32)];
qt.accumulate_forces(&particles, DEFAULT_THETA, 100.0, MIN_DIST2, &mut out);
assert_eq!(out[0], (0.0, 0.0));
}
#[test]
fn coincident_particle_stack_builds_and_acts_as_one_point_mass() {
let mut particles = Vec::new();
for _ in 0..8 {
particles.push(Particle::at(-50.0, 0.0));
}
for _ in 0..8 {
particles.push(Particle::at(40.0, 30.0));
}
particles.push(Particle::at(0.0, 0.0));
let strength = 100.0;
let qt = Quadtree::build(&particles, MIN_SPLIT_DIST2, MIN_CELL_SIZE); let mut out = vec![(0f32, 0f32); particles.len()];
qt.accumulate_forces(&particles, DEFAULT_THETA, strength, MIN_DIST2, &mut out);
let probe = out[16];
assert!(probe.0.is_finite() && probe.1.is_finite());
let mut expected = (0.0f32, 0.0f32);
apply_point(0.0, 0.0, -50.0, 0.0, 8.0, strength, MIN_DIST2, &mut expected);
apply_point(0.0, 0.0, 40.0, 30.0, 8.0, strength, MIN_DIST2, &mut expected);
let diff = ((probe.0 - expected.0).powi(2) + (probe.1 - expected.1).powi(2)).sqrt();
let mag = (expected.0 * expected.0 + expected.1 * expected.1).sqrt();
assert!(
diff <= mag * 0.05,
"coincident stack should act as one 8-mass point: got {probe:?}, expected {expected:?}"
);
}
#[test]
fn build_weighted_with_none_masses_matches_build_exactly() {
let particles = deterministic_particles(40);
let plain = Quadtree::build(&particles, MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let weighted = Quadtree::build_weighted(&particles, None, MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let mut out_plain = vec![(0f32, 0f32); particles.len()];
let mut out_weighted = vec![(0f32, 0f32); particles.len()];
plain.accumulate_forces(&particles, DEFAULT_THETA, 400.0, MIN_DIST2, &mut out_plain);
weighted.accumulate_forces(&particles, DEFAULT_THETA, 400.0, MIN_DIST2, &mut out_weighted);
assert_eq!(out_plain, out_weighted, "None masses must reproduce the uniform-1.0 default exactly");
}
#[test]
fn build_weighted_preserves_each_entrys_own_mass_through_a_later_split() {
let particles = vec![Particle::at(0.0, 0.0), Particle::at(0.0, 0.0), Particle::at(500.0, 500.0)];
let masses = vec![2.0f32, 5.0, 1.0];
let strength = 100.0;
let qt = Quadtree::build_weighted(&particles, Some(&masses), MIN_SPLIT_DIST2, MIN_CELL_SIZE);
let mut out = vec![(0f32, 0f32); particles.len()];
qt.accumulate_forces(&particles, 1e-9, strength, MIN_DIST2, &mut out);
let mut expected = (0.0f32, 0.0f32);
apply_point(500.0, 500.0, 0.0, 0.0, 7.0, strength, MIN_DIST2, &mut expected);
let diff = ((out[2].0 - expected.0).powi(2) + (out[2].1 - expected.1).powi(2)).sqrt();
let mag = (expected.0 * expected.0 + expected.1 * expected.1).sqrt();
assert!(
diff <= mag * 0.02,
"the far particle must feel the merged (0,0) pair as one mass-7.0 point: got {:?}, expected {expected:?}",
out[2]
);
let mut buggy = (0.0f32, 0.0f32);
apply_point(500.0, 500.0, 0.0, 0.0, 2.0, strength, MIN_DIST2, &mut buggy);
let buggy_diff = ((out[2].0 - buggy.0).powi(2) + (out[2].1 - buggy.1).powi(2)).sqrt();
assert!(buggy_diff > mag * 0.3, "the fixed and buggy answers must be clearly distinguishable, not coincidentally close");
}
}