use glam::{Vec2, Vec3};
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use crate::force_3d::{OctNode, accumulate_repulsion_3d, build_octree};
const MIN_CELL_SIZE: f32 = 1.0;
#[derive(Debug, Clone, Copy)]
pub struct SimulationConfig {
pub repulsion_strength: f32,
pub link_strength: f32,
pub link_distance: f32,
pub center_strength: f32,
pub velocity_decay: f32,
pub theta: f32,
pub seed: u64,
pub dimensions: u8,
}
impl Default for SimulationConfig {
fn default() -> Self {
Self {
repulsion_strength: -30.0,
link_strength: 0.1,
link_distance: 30.0,
center_strength: 0.05,
velocity_decay: 0.4,
theta: 0.9,
seed: 0,
dimensions: 2,
}
}
}
#[derive(Debug, Clone, Copy)]
struct Edge {
src: u32,
dst: u32,
weight: f32,
}
pub struct Simulation {
config: SimulationConfig,
positions: Vec<Vec2>,
velocities: Vec<Vec2>,
forces: Vec<Vec2>,
edges: Vec<Edge>,
tree: Vec<QuadNode>,
positions_3d: Vec<Vec3>,
velocities_3d: Vec<Vec3>,
forces_3d: Vec<Vec3>,
oct_tree: Vec<OctNode>,
}
impl Simulation {
pub fn new(n_nodes: u32, config: SimulationConfig) -> Self {
let n = n_nodes as usize;
let mut rng = ChaCha8Rng::seed_from_u64(config.seed);
let radius = (n_nodes as f32).max(1.0).sqrt() * 10.0;
let dims = if config.dimensions == 3 { 3u8 } else { 2u8 };
let mut cfg = config;
cfg.dimensions = dims;
let mut positions = Vec::with_capacity(n);
let mut positions_3d: Vec<Vec3> = if dims == 3 {
Vec::with_capacity(n)
} else {
Vec::new()
};
for _ in 0..n {
if dims == 3 {
loop {
let x: f32 = rng.gen_range(-1.0..=1.0);
let y: f32 = rng.gen_range(-1.0..=1.0);
let z: f32 = rng.gen_range(-1.0..=1.0);
if x * x + y * y + z * z <= 1.0 {
let p = Vec3::new(x * radius, y * radius, z * radius);
positions_3d.push(p);
positions.push(Vec2::new(p.x, p.y));
break;
}
}
} else {
loop {
let x: f32 = rng.gen_range(-1.0..=1.0);
let y: f32 = rng.gen_range(-1.0..=1.0);
if x * x + y * y <= 1.0 {
positions.push(Vec2::new(x * radius, y * radius));
break;
}
}
}
}
let (velocities_3d, forces_3d) = if dims == 3 {
(vec![Vec3::ZERO; n], vec![Vec3::ZERO; n])
} else {
(Vec::new(), Vec::new())
};
Self {
config: cfg,
positions,
velocities: vec![Vec2::ZERO; n],
forces: vec![Vec2::ZERO; n],
edges: Vec::new(),
tree: Vec::new(),
positions_3d,
velocities_3d,
forces_3d,
oct_tree: Vec::new(),
}
}
pub fn set_edges(&mut self, edges: &[(u32, u32, f32)]) {
let n = self.positions.len() as u32;
self.edges.clear();
self.edges.reserve(edges.len());
for &(src, dst, weight) in edges {
if src < n && dst < n && src != dst {
self.edges.push(Edge { src, dst, weight });
}
}
}
#[inline]
pub fn node_count(&self) -> usize {
self.positions.len()
}
#[inline]
pub fn positions(&self) -> &[Vec2] {
&self.positions
}
#[inline]
pub fn positions_3d(&self) -> &[Vec3] {
&self.positions_3d
}
pub fn tick(&mut self, dt: f32) {
if self.config.dimensions == 3 {
self.tick_3d(dt);
return;
}
self.tick_2d(dt);
}
fn tick_2d(&mut self, dt: f32) {
let n = self.positions.len();
if n == 0 {
return;
}
for f in self.forces.iter_mut() {
*f = Vec2::ZERO;
}
self.build_tree();
let theta2 = self.config.theta * self.config.theta;
let repulsion = self.config.repulsion_strength;
if !self.tree.is_empty() {
for i in 0..n {
let pos = self.positions[i];
let f = accumulate_repulsion(&self.tree, 0, pos, i as u32, theta2, repulsion);
self.forces[i] += f;
}
}
let link_k = self.config.link_strength;
let link_d = self.config.link_distance;
for e in &self.edges {
let a = self.positions[e.src as usize];
let b = self.positions[e.dst as usize];
let delta = b - a;
let dist = delta.length().max(1e-6);
let mag = (dist - link_d) * link_k * e.weight;
let dir = delta / dist;
self.forces[e.src as usize] += dir * mag;
self.forces[e.dst as usize] -= dir * mag;
}
let c = self.config.center_strength;
if c != 0.0 {
for i in 0..n {
self.forces[i] -= self.positions[i] * c;
}
}
let decay = self.config.velocity_decay;
for i in 0..n {
self.velocities[i] *= decay;
self.velocities[i] += self.forces[i] * dt;
self.positions[i] += self.velocities[i] * dt;
}
}
fn tick_3d(&mut self, dt: f32) {
let n = self.positions_3d.len();
if n == 0 {
return;
}
for f in self.forces_3d.iter_mut() {
*f = Vec3::ZERO;
}
build_octree(&mut self.oct_tree, &self.positions_3d, MIN_CELL_SIZE);
let theta2 = self.config.theta * self.config.theta;
let repulsion = self.config.repulsion_strength;
if !self.oct_tree.is_empty() {
for i in 0..n {
let pos = self.positions_3d[i];
let f = accumulate_repulsion_3d(
&self.oct_tree,
0,
pos,
i as u32,
theta2,
repulsion,
);
self.forces_3d[i] += f;
}
}
let link_k = self.config.link_strength;
let link_d = self.config.link_distance;
for e in &self.edges {
let a = self.positions_3d[e.src as usize];
let b = self.positions_3d[e.dst as usize];
let delta = b - a;
let dist = delta.length().max(1e-6);
let mag = (dist - link_d) * link_k * e.weight;
let dir = delta / dist;
self.forces_3d[e.src as usize] += dir * mag;
self.forces_3d[e.dst as usize] -= dir * mag;
}
let c = self.config.center_strength;
if c != 0.0 {
for i in 0..n {
self.forces_3d[i] -= self.positions_3d[i] * c;
}
}
let decay = self.config.velocity_decay;
for i in 0..n {
self.velocities_3d[i] *= decay;
self.velocities_3d[i] += self.forces_3d[i] * dt;
self.positions_3d[i] += self.velocities_3d[i] * dt;
self.positions[i] = Vec2::new(self.positions_3d[i].x, self.positions_3d[i].y);
}
}
fn build_tree(&mut self) {
self.tree.clear();
let n = self.positions.len();
if n == 0 {
return;
}
let mut min = self.positions[0];
let mut max = self.positions[0];
for p in &self.positions[1..] {
min = min.min(*p);
max = max.max(*p);
}
let extent = (max - min).max_element().max(MIN_CELL_SIZE);
let center = (min + max) * 0.5;
let half = Vec2::splat(extent * 0.5);
let min = center - half;
let max = center + half;
self.tree.push(QuadNode::empty(min, max));
for i in 0..n {
let p = self.positions[i];
insert(&mut self.tree, 0, i as u32, p);
}
finalise(&mut self.tree, 0);
}
}
#[derive(Debug, Clone, Copy)]
struct QuadNode {
min: Vec2,
max: Vec2,
mass: f32,
com: Vec2,
body: Option<u32>,
children: [u32; 4], has_children: bool,
}
impl QuadNode {
fn empty(min: Vec2, max: Vec2) -> Self {
Self {
min,
max,
mass: 0.0,
com: Vec2::ZERO,
body: None,
children: [u32::MAX; 4],
has_children: false,
}
}
#[inline]
fn size(&self) -> f32 {
(self.max - self.min).max_element()
}
#[inline]
fn center(&self) -> Vec2 {
(self.min + self.max) * 0.5
}
#[inline]
fn quadrant_of(&self, p: Vec2) -> usize {
let c = self.center();
let east = (p.x >= c.x) as usize;
let north = (p.y >= c.y) as usize;
(north << 1) | east
}
#[inline]
fn child_bounds(&self, q: usize) -> (Vec2, Vec2) {
let c = self.center();
let (x_min, x_max) = if q & 1 == 1 {
(c.x, self.max.x)
} else {
(self.min.x, c.x)
};
let (y_min, y_max) = if q & 2 == 2 {
(c.y, self.max.y)
} else {
(self.min.y, c.y)
};
(Vec2::new(x_min, y_min), Vec2::new(x_max, y_max))
}
}
fn insert(tree: &mut Vec<QuadNode>, idx: usize, body: u32, pos: Vec2) {
if !tree[idx].has_children && tree[idx].body.is_none() && tree[idx].mass == 0.0 {
tree[idx].body = Some(body);
tree[idx].com = pos;
tree[idx].mass = 1.0;
return;
}
if !tree[idx].has_children {
if tree[idx].size() <= MIN_CELL_SIZE {
let new_mass = tree[idx].mass + 1.0;
tree[idx].com = (tree[idx].com * tree[idx].mass + pos) / new_mass;
tree[idx].mass = new_mass;
return;
}
let existing = tree[idx].body.take().expect("leaf without a body");
let existing_pos = tree[idx].com;
tree[idx].mass = 0.0;
tree[idx].com = Vec2::ZERO;
tree[idx].has_children = true;
let q_existing = tree[idx].quadrant_of(existing_pos);
let c_existing = create_or_get_child(tree, idx, q_existing);
insert(tree, c_existing, existing, existing_pos);
let q_new = tree[idx].quadrant_of(pos);
let c_new = create_or_get_child(tree, idx, q_new);
insert(tree, c_new, body, pos);
return;
}
let q = tree[idx].quadrant_of(pos);
let c = create_or_get_child(tree, idx, q);
insert(tree, c, body, pos);
}
fn create_or_get_child(tree: &mut Vec<QuadNode>, parent: usize, quadrant: usize) -> usize {
let existing = tree[parent].children[quadrant];
if existing != u32::MAX {
return existing as usize;
}
let (cmin, cmax) = tree[parent].child_bounds(quadrant);
let idx = tree.len();
tree.push(QuadNode::empty(cmin, cmax));
tree[parent].children[quadrant] = idx as u32;
idx
}
fn finalise(tree: &mut [QuadNode], idx: usize) {
if !tree[idx].has_children {
return;
}
let children = tree[idx].children;
let mut mass = 0.0;
let mut com = Vec2::ZERO;
for &c in &children {
if c == u32::MAX {
continue;
}
finalise(tree, c as usize);
let child = tree[c as usize];
mass += child.mass;
com += child.com * child.mass;
}
if mass > 0.0 {
com /= mass;
}
tree[idx].mass = mass;
tree[idx].com = com;
}
fn accumulate_repulsion(
tree: &[QuadNode],
idx: usize,
target_pos: Vec2,
target_idx: u32,
theta2: f32,
repulsion: f32,
) -> Vec2 {
let node = &tree[idx];
if node.mass <= 0.0 {
return Vec2::ZERO;
}
if !node.has_children {
if let Some(body) = node.body {
if body == target_idx {
let residual_mass = node.mass - 1.0;
if residual_mass <= 0.0 {
return Vec2::ZERO;
}
return pair_force(target_pos, node.com, residual_mass, repulsion);
}
}
return pair_force(target_pos, node.com, node.mass, repulsion);
}
let delta = node.com - target_pos;
let dist2 = delta.length_squared();
let size = node.size();
if size * size < theta2 * dist2 {
return pair_force(target_pos, node.com, node.mass, repulsion);
}
let mut total = Vec2::ZERO;
for &c in &node.children {
if c == u32::MAX {
continue;
}
total += accumulate_repulsion(tree, c as usize, target_pos, target_idx, theta2, repulsion);
}
total
}
#[inline]
fn pair_force(from: Vec2, to: Vec2, mass: f32, repulsion: f32) -> Vec2 {
let delta = to - from;
let dist2 = delta.length_squared() + 0.01;
let dist = dist2.sqrt();
let dir = delta / dist;
let magnitude = repulsion * mass / dist2;
dir * magnitude
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn two_linked_nodes_converge() {
let mut sim = Simulation::new(
2,
SimulationConfig {
seed: 42,
link_strength: 0.8,
link_distance: 30.0,
repulsion_strength: -30.0,
center_strength: 0.0,
velocity_decay: 0.9,
..Default::default()
},
);
sim.set_edges(&[(0, 1, 1.0)]);
for _ in 0..2000 {
sim.tick(0.02);
}
let d = (sim.positions()[0] - sim.positions()[1]).length();
let target = 30.0_f32;
assert!(
d > target * 0.8 && d < target * 1.2,
"expected ~30 got {d}"
);
}
#[test]
fn star_topology_stabilizes() {
let n = 9u32;
let mut sim = Simulation::new(
n,
SimulationConfig {
seed: 1,
link_strength: 0.2,
link_distance: 30.0,
repulsion_strength: -30.0,
center_strength: 0.1,
velocity_decay: 0.5,
..Default::default()
},
);
let edges: Vec<_> = (1..n).map(|i| (0u32, i, 1.0)).collect();
sim.set_edges(&edges);
for _ in 0..1000 {
sim.tick(0.02);
}
let max_v = sim
.velocities
.iter()
.map(|v| v.length())
.fold(0.0_f32, f32::max);
assert!(max_v < 0.5, "max_v={max_v}");
}
#[test]
fn deterministic_tick_sequence() {
let config = SimulationConfig {
seed: 12345,
..Default::default()
};
let mut a = Simulation::new(50, config);
let mut b = Simulation::new(50, config);
let edges: Vec<_> = (0..49u32).map(|i| (i, i + 1, 1.0)).collect();
a.set_edges(&edges);
b.set_edges(&edges);
for _ in 0..100 {
a.tick(0.02);
b.tick(0.02);
}
for (pa, pb) in a.positions().iter().zip(b.positions().iter()) {
assert_eq!(pa, pb, "deterministic run diverged");
}
}
#[test]
fn isolated_nodes_repel() {
let n = 10u32;
let mut sim = Simulation::new(
n,
SimulationConfig {
seed: 7,
repulsion_strength: -80.0,
center_strength: 0.0,
velocity_decay: 0.6,
..Default::default()
},
);
let cluster_rng = &mut rand_chacha::ChaCha8Rng::seed_from_u64(99);
for p in sim.positions.iter_mut() {
loop {
let x: f32 = cluster_rng.gen_range(-1.0..=1.0);
let y: f32 = cluster_rng.gen_range(-1.0..=1.0);
if x * x + y * y <= 1.0 {
*p = Vec2::new(x, y);
break;
}
}
}
for v in sim.velocities.iter_mut() {
*v = Vec2::ZERO;
}
for _ in 0..500 {
sim.tick(0.02);
}
let mut min_d = f32::INFINITY;
for i in 0..sim.positions.len() {
for j in (i + 1)..sim.positions.len() {
let d = (sim.positions[i] - sim.positions[j]).length();
if d < min_d {
min_d = d;
}
}
}
assert!(min_d > 5.0, "min pairwise distance {min_d} ≤ 5");
}
}