Browse Source

Initial commit. 3d stuff should be done. 2d stuff unfinished. More documentation required

master
Stephen 4 months ago
commit
1e9bb0f2bd
11 changed files with 646 additions and 0 deletions
  1. +3
    -0
      .gitignore
  2. +17
    -0
      Cargo.toml
  3. +47
    -0
      benches/3d_benchmark.rs
  4. +2
    -0
      rustfmt.toml
  5. +0
    -0
      src/barnes_hut_2d.rs
  6. +254
    -0
      src/barnes_hut_3d.rs
  7. +120
    -0
      src/lib.rs
  8. +0
    -0
      src/particle_2d.rs
  9. +12
    -0
      src/particle_3d.rs
  10. +90
    -0
      src/vector_2d.rs
  11. +101
    -0
      src/vector_3d.rs

+ 3
- 0
.gitignore View File

@ -0,0 +1,3 @@
/target
Cargo.lock
.idea

+ 17
- 0
Cargo.toml View File

@ -0,0 +1,17 @@
[package]
name = "nbody_barnes_hut"
version = "0.1.0"
authors = ["Stephen <webmaster@scd31.com>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
[dev-dependencies]
criterion = "0.3"
rand = "0.7"
[[bench]]
name = "3d_benchmark"
harness = false

+ 47
- 0
benches/3d_benchmark.rs View File

@ -0,0 +1,47 @@
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use nbody_barnes_hut::barnes_hut_3d::OctTree;
use nbody_barnes_hut::particle_3d::Particle3D;
use nbody_barnes_hut::vector_3d::Vector3D;
use rand::prelude::*;
fn criterion_benchmark(c: &mut Criterion) {
// Create 10 000 random points
let mut rng = rand::thread_rng();
let points: Vec<Particle3D> = (0..10_000)
.map(|_| {
let pos = Vector3D::new(
rng.gen_range(-1000.0, 1000.0),
rng.gen_range(-1000.0, 1000.0),
rng.gen_range(-1000.0, 1000.0),
);
Particle3D::new(pos, 30.0)
})
.collect();
// This is pretty hacky
let points_ref = &points.iter().collect::<Vec<&Particle3D>>()[..];
c.bench_function("insert 10k points", |b| {
b.iter(|| OctTree::new(black_box(points_ref), black_box(0.5)))
});
let tree = OctTree::new(points_ref, 0.5);
c.bench_function("calc gravity for 10k points", |b| {
b.iter(|| {
for p in &points {
black_box(tree.calc_forces_on_particle(
p.position,
(),
|d_squared, mass, dis_vec, _| {
black_box(d_squared);
black_box(mass);
black_box(dis_vec);
black_box(Vector3D::ZERO)
},
));
}
})
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);

+ 2
- 0
rustfmt.toml View File

@ -0,0 +1,2 @@
tab_spaces = 4
hard_tabs = true

+ 0
- 0
src/barnes_hut_2d.rs View File


+ 254
- 0
src/barnes_hut_3d.rs View File

@ -0,0 +1,254 @@
use crate::particle_3d::Particle3D;
use crate::vector_3d::Vector3D;
use std::ops::AddAssign;
pub struct OctNode {
children_base: usize,
children_mask: u8,
quad_size_squared: f64,
mass: f64,
}
impl OctNode {
fn new(min: Vector3D, max: Vector3D) -> Self {
let s = (min - max).dist_squared();
OctNode {
children_base: 0,
children_mask: 0,
mass: 0.0,
quad_size_squared: s,
}
}
}
pub struct OctTree {
nodes: Vec<OctNode>,
min_coords: Vec<Vector3D>,
max_coords: Vec<Vector3D>,
avg_weighted_positions: Vec<Vector3D>,
center_of_masses: Vec<Vector3D>,
theta_squared: f64,
}
impl OctTree {
pub fn new(particles: &[&Particle3D], theta: f64) -> Self {
// TODO fix this - we shouldn't need to construct a vec here
let (min, max) = find_min_max(&particles);
let mut tree = OctTree::construct_empty_tree(min, max, theta, particles.len());
tree.add_particles_to_specific_node(particles, 0);
//calculate center of masses
for (i, n) in &mut tree.nodes.iter().enumerate() {
tree.center_of_masses
.push(tree.avg_weighted_positions[i] / n.mass);
}
tree
}
pub fn update_theta(&mut self, theta: f64) {
assert!(theta >= 0.0);
assert!(theta <= 1.0);
self.theta_squared = theta * theta;
}
fn construct_empty_tree(min: Vector3D, max: Vector3D, theta: f64, n: usize) -> Self {
let mut o = OctTree {
nodes: Vec::with_capacity(n),
min_coords: Vec::with_capacity(n),
max_coords: Vec::with_capacity(n),
avg_weighted_positions: Vec::with_capacity(n),
center_of_masses: Vec::with_capacity(n),
theta_squared: theta * theta,
};
o.nodes.push(OctNode::new(min, max)); // Create our root node
o.min_coords.push(min);
o.max_coords.push(max);
o.avg_weighted_positions.push(Vector3D::ZERO);
o
}
pub fn calc_forces_on_particle<T: Copy, F: Fn(f64, f64, Vector3D, T) -> Vector3D + Copy>(
&self,
particle_position: Vector3D,
other_attributes: T,
calc: F,
) -> Vector3D {
self.barnes_hut_specific_node(particle_position, 0, other_attributes, calc)
}
fn barnes_hut_specific_node<T: Copy, F: Fn(f64, f64, Vector3D, T) -> Vector3D + Copy>(
&self,
particle_position: Vector3D,
cur_node: usize,
other_attributes: T,
calc: F,
) -> Vector3D {
let c_node = &self.nodes[cur_node];
let s_squared = c_node.quad_size_squared;
let dis_vec = self.center_of_masses[cur_node] - particle_position;
let d_squared = dis_vec.dist_squared();
let theta_squared = self.theta_squared;
if c_node.children_mask == 0 || s_squared < theta_squared * d_squared {
if d_squared.abs() < 0.00001 {
return Vector3D::ZERO;
}
calc(d_squared, c_node.mass, dis_vec, other_attributes)
} else {
let mut c = c_node.children_base;
//Divide and conquer
let mut sum: Vector3D = Vector3D::ZERO;
for i in 0..8 {
if c_node.children_mask & (1 << i) != 0 {
sum +=
self.barnes_hut_specific_node(particle_position, c, other_attributes, calc);
c += 1;
}
}
sum
}
}
fn add_particles_to_specific_node(&mut self, particles: &[&Particle3D], cur_node: usize) {
if particles.is_empty() {
return;
}
if particles.len() == 1 {
self.nodes[cur_node].mass += particles[0].mass;
self.avg_weighted_positions[cur_node] += particles[0].position * particles[0].mass;
return;
}
let center_point = (self.min_coords[cur_node] + self.max_coords[cur_node]) / 2.0;
let mut particle_octs: [Vec<&Particle3D>; 8] = [
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
Vec::with_capacity(particles.len() / 4),
];
for &particle in particles {
// Update ourself
self.nodes[cur_node].mass += particle.mass;
self.avg_weighted_positions[cur_node] += particle.position * particle.mass;
let octant_index = OctTree::get_id_from_center(center_point, particle.position);
particle_octs[octant_index].push(particle);
}
//Recurse
self.nodes[cur_node].children_base = self.nodes.len();
for (i, particle_oct) in particle_octs.iter().enumerate() {
if !particle_oct.is_empty() {
let (min, max) = OctTree::get_octant_bounding_box_from_id(
self.min_coords[cur_node],
self.max_coords[cur_node],
center_point,
i,
);
self.nodes.push(OctNode::new(min, max));
self.min_coords.push(min);
self.max_coords.push(max);
self.avg_weighted_positions.push(Vector3D::ZERO);
self.nodes[cur_node].children_mask |= 1 << i;
}
}
let mut c = self.nodes[cur_node].children_base;
for particle_oct in particle_octs.iter() {
if !particle_oct.is_empty() {
self.add_particles_to_specific_node(particle_oct, c);
c += 1;
}
}
}
pub fn get_id_from_center(center: Vector3D, point: Vector3D) -> usize {
let offset = point - center;
//We can look at the sign of the components of the offset to figure out which octant it is in.
let x_offset = if offset.x > 0.0 { 0 } else { 1 };
let y_offset = if offset.y > 0.0 { 0 } else { 1 };
let z_offset = if offset.z > 0.0 { 0 } else { 1 };
x_offset + y_offset * 2 + z_offset * 4 // basic binary stuff here
}
pub fn get_octant_bounding_box_from_id(
min: Vector3D,
max: Vector3D,
center: Vector3D,
idx: usize,
) -> (Vector3D, Vector3D) {
let mut min_coord: Vector3D = Vector3D::ZERO;
let mut max_coord: Vector3D = Vector3D::ZERO;
if (idx & 1) != 0 {
min_coord.x = min.x;
max_coord.x = center.x;
} else {
min_coord.x = center.x;
max_coord.x = max.x;
}
if (idx & 2) != 0 {
min_coord.y = min.y;
max_coord.y = center.y;
} else {
min_coord.y = center.y;
max_coord.y = max.y;
}
if (idx & 4) != 0 {
min_coord.z = min.z;
max_coord.z = center.z;
} else {
min_coord.z = center.z;
max_coord.z = max.z;
}
(min_coord, max_coord)
}
}
fn find_min_max(particles: &[&Particle3D]) -> (Vector3D, Vector3D) {
let mut min: Vector3D = Vector3D::ZERO;
let mut max: Vector3D = Vector3D::ZERO;
for particle in particles {
let particle = particle.position;
if particle.x < min.x {
min.x = particle.x;
}
if particle.y < min.y {
min.y = particle.y;
}
if particle.z < min.z {
min.z = particle.z;
}
if particle.x > max.x {
max.x = particle.x;
}
if particle.y > max.y {
max.y = particle.y;
}
if particle.z > max.z {
max.z = particle.z;
}
}
(min, max)
}

+ 120
- 0
src/lib.rs View File

@ -0,0 +1,120 @@
//! # Nbody barnes hut
//!
//! `nbody_barnes_hut` is designed to facilitate the simulation of N-body systems in `O(nlogn)` time.
//! This is useful for many applications: common ones are gravitational simulations and electrostatic simulations.
//! Simulations in 2D and 3D are both supported.
//!
//! This crate is not multithreaded. Rather, call `calc_forces_on_particle` in a multithreaded loop (for example, with `rayon`).
//! The time to create the tree is negligible in comparison to the time used calculating forces.
//!
//! # Example
//!
//! Here is a basic 3D gravitational simulator:
//! ```
//! // Create 10 000 random points
//! let mut rng = rand::thread_rng();
//! let points: Vec<Particle3D> = (0..10_000)
//! .map(|_| {
//! let pos = Vector3D::new(
//! rng.gen_range(-1000.0, 1000.0),
//! rng.gen_range(-1000.0, 1000.0),
//! rng.gen_range(-1000.0, 1000.0),
//! );
//! Particle3D::new(pos, 30.0)
//! })
//! .collect();
//!
//! // This is pretty hacky
//! let points_ref = &points.iter().collect::<Vec<&Particle3D>>()[..];
//!
//! let tree = OctTree::new(points_ref, 0.5);
//!
//! for p in &points {
//! // Do something with this value
//! let acceleration_on_particle = tree.calc_forces_on_particle(
//! p.position,
//! (),
//! |d_squared, mass, dis_vec, _| {
//! // dis_vec is not normalized, so we have to normalize it here
//! G * mass * dis_vec / (d_squared * d_squared.sqrt())
//! },
//! );
//! }
//! ```
pub mod barnes_hut_2d;
pub mod barnes_hut_3d;
pub mod particle_2d;
pub mod particle_3d;
pub mod vector_2d;
pub mod vector_3d;
#[cfg(test)]
mod tests {
use crate::barnes_hut_3d::OctTree;
use crate::particle_3d::Particle3D;
use crate::vector_3d::Vector3D;
#[test]
fn gravity_tests_3d() {
let particle_1 = Particle3D::new(Vector3D::ZERO, 10.0);
let p2_position = Vector3D::new(5.0, 10.0, 15.0);
let particle_2 = Particle3D::new(p2_position, 100.0);
let dist_squared = 5.0 * 5.0 + 10.0 * 10.0 + 15.0 * 15.0;
let force_expected = -10.0 * p2_position.normalize() / dist_squared;
let tree = OctTree::new(&[&particle_1, &particle_2][..], 0.5);
let force_actual = tree.calc_forces_on_particle(
particle_2.position,
(),
|dist_squared, mass, dis_vec, _| mass * dis_vec / (dist_squared * dist_squared.sqrt()),
);
assert_eq!(force_expected, force_actual);
let tree = OctTree::new(&[&particle_1][..], 0.5);
assert_vector3d_eq(
Vector3D::ZERO,
tree.calc_forces_on_particle(
particle_1.position,
(),
|dist_squared, mass, dis_vec, _| {
mass * dis_vec / (dist_squared * dist_squared.sqrt())
},
),
)
}
#[test]
fn electrostatics_3d() {
let charge_one = 100.0;
// NOTE we are using the charge as mass
let particle_1 = Particle3D::new(Vector3D::ZERO, charge_one);
let charge_two = -50.0;
let particle_2_pos = Vector3D::new(15.0, 20.0, -1.1);
let particle_2 = Particle3D::new(particle_2_pos, charge_two);
let dis_squared = 15.0 * 15.0 + 20.0 * 20.0 + 1.1 * 1.1;
// F = Kq1q2/r^2
let accel_expected = -charge_one * charge_two * particle_2_pos.normalize() / dis_squared;
let tree = OctTree::new(&[&particle_1, &particle_2][..], 0.5);
let accel_actual = tree.calc_forces_on_particle(
particle_2_pos,
charge_two,
|dist_squared, charge_other, dis_vec, charge_self| {
charge_other * charge_self * dis_vec / (dist_squared * dist_squared.sqrt())
},
);
assert_vector3d_eq(accel_expected, accel_actual);
}
fn assert_vector3d_eq(a: Vector3D, b: Vector3D) {
let d = (a - b).dist();
assert!(d < 0.00001);
}
}

+ 0
- 0
src/particle_2d.rs View File


+ 12
- 0
src/particle_3d.rs View File

@ -0,0 +1,12 @@
use crate::vector_3d::Vector3D;
pub struct Particle3D {
pub position: Vector3D,
pub mass: f64,
}
impl Particle3D {
pub fn new(position: Vector3D, mass: f64) -> Self {
Self { position, mass }
}
}

+ 90
- 0
src/vector_2d.rs View File

@ -0,0 +1,90 @@
use std::ops::{Add, AddAssign, Div, Mul, Sub};
// Fast vectors
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Vector2D {
pub x: f64,
pub y: f64,
}
impl Vector2D {
pub const ZERO: Vector2D = Vector2D { x: 0.0, y: 0.0 };
pub fn new(x: f64, y: f64) -> Self {
Self { x, y }
}
pub fn dist_squared(self) -> f64 {
self.x * self.x + self.y * self.y
}
pub fn dist(self) -> f64 {
self.dist_squared().sqrt()
}
pub fn normalize(self) -> Self {
self / self.dist()
}
}
impl Mul<f64> for Vector2D {
type Output = Self;
fn mul(self, rhs: f64) -> Self {
Self {
x: self.x * rhs,
y: self.y * rhs,
}
}
}
impl Mul<Vector2D> for f64 {
type Output = Vector2D;
fn mul(self, rhs: Vector2D) -> Vector2D {
Vector2D {
x: self * rhs.x,
y: self * rhs.y,
}
}
}
impl Div<f64> for Vector2D {
type Output = Self;
fn div(self, rhs: f64) -> Self {
Self {
x: self.x / rhs,
y: self.y / rhs,
}
}
}
impl Sub for Vector2D {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self {
x: self.x - rhs.x,
y: self.y - rhs.y,
}
}
}
impl Add for Vector2D {
type Output = Self;
fn add(self, other: Self) -> Self {
Self {
x: self.x + other.x,
y: self.y + other.y,
}
}
}
impl AddAssign for Vector2D {
fn add_assign(&mut self, other: Self) {
self.x += other.x;
self.y += other.y;
}
}

+ 101
- 0
src/vector_3d.rs View File

@ -0,0 +1,101 @@
use std::ops::{Add, AddAssign, Div, Mul, Sub};
// Fast vectors
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Vector3D {
pub x: f64,
pub y: f64,
pub z: f64,
}
impl Vector3D {
pub const ZERO: Vector3D = Vector3D {
x: 0.0,
y: 0.0,
z: 0.0,
};
pub fn new(x: f64, y: f64, z: f64) -> Self {
Self { x, y, z }
}
pub fn dist_squared(self) -> f64 {
self.x * self.x + self.y * self.y + self.z * self.z
}
pub fn dist(self) -> f64 {
self.dist_squared().sqrt()
}
pub fn normalize(self) -> Self {
self / self.dist()
}
}
impl Mul<f64> for Vector3D {
type Output = Self;
fn mul(self, rhs: f64) -> Self {
Self {
x: self.x * rhs,
y: self.y * rhs,
z: self.z * rhs,
}
}
}
impl Mul<Vector3D> for f64 {
type Output = Vector3D;
fn mul(self, rhs: Vector3D) -> Vector3D {
Vector3D {
x: self * rhs.x,
y: self * rhs.y,
z: self * rhs.z,
}
}
}
impl Div<f64> for Vector3D {
type Output = Self;
fn div(self, rhs: f64) -> Self {
Self {
x: self.x / rhs,
y: self.y / rhs,
z: self.z / rhs,
}
}
}
impl Sub for Vector3D {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self {
x: self.x - rhs.x,
y: self.y - rhs.y,
z: self.z - rhs.z,
}
}
}
impl Add for Vector3D {
type Output = Self;
fn add(self, other: Self) -> Self {
Self {
x: self.x + other.x,
y: self.y + other.y,
z: self.z + other.z,
}
}
}
impl AddAssign for Vector3D {
fn add_assign(&mut self, other: Self) {
self.x += other.x;
self.y += other.y;
self.z += other.z;
}
}

Loading…
Cancel
Save