barnes_hut/
lib.rs

1//! This algorithm uses Tree Code to group source bodies, as an approximation. It leads to O(N(log N))
2//! computation time, where `N` is the number of bodies. Canonical use cases include gravity, and charged
3//! particle simulations.
4//!
5//! See the [readme](https://github.com/David-OConnor/barnes_hut/blob/main/README.md) for details,
6//! including an example.
7
8#![allow(non_ascii_idents)]
9#![allow(mixed_script_confusables)]
10
11// todo: Ideally make generic over f32 and f64, but we don't have a good way to do that with `Vec3`.
12
13use std::{fmt, fmt::Formatter};
14
15#[cfg(feature = "encode")]
16use bincode::{Decode, Encode};
17use lin_alg::f64::Vec3;
18use rayon::prelude::*;
19
20#[derive(Clone, Debug)]
21#[cfg_attr(feature = "bincode", derive(Encode, Decode))]
22pub struct BhConfig {
23    /// This determines how aggressively we group. It's no lower than 0. 0 means no grouping.
24    /// (Best accuracy; poorest performance; effectively a naive N-body). Higher values
25    /// decrease accuracy, and are more performant.
26    pub θ: f64,
27    pub max_bodies_per_node: usize,
28    /// This is a limit on tree division, preventing getting stuck in a loop, e.g. for particles with close.
29    /// (or identical) positions
30    pub max_tree_depth: usize,
31}
32
33impl Default for BhConfig {
34    fn default() -> Self {
35        Self {
36            θ: 0.5,
37            max_bodies_per_node: 1,
38            max_tree_depth: 15,
39        }
40    }
41}
42
43/// We use this to allow for arbitrary body (or particle etc) types in application code to
44/// use this library. Substitute `charge` for `mass` as required.
45pub trait BodyModel {
46    fn posit(&self) -> Vec3;
47    fn mass(&self) -> f64;
48}
49
50#[derive(Clone, Debug)]
51#[cfg_attr(feature = "bincode", derive(Encode, Decode))]
52/// A cubical bounding box. length=width=depth.
53pub struct Cube {
54    pub center: Vec3,
55    pub width: f64,
56}
57
58impl Cube {
59    /// Construct minimum limits that encompass all bodies. Run these each time the bodies change,
60    /// or perhaps use a pad and do it at a coarser interval.
61    ///
62    /// The pad allows us to keep the same cube for multiple timesteps, but taking into acacount
63    /// potential movement of bodies outside the cube between these updates.
64    ///
65    /// The z offset is intended for the case where the Z coordinate for all particles is 0.
66    /// This prevents the divisions straddling the points, doubling the number of nodes.
67    pub fn from_bodies<T: BodyModel>(bodies: &[T], pad: f64, z_offset: bool) -> Option<Self> {
68        if bodies.is_empty() {
69            return None;
70        }
71
72        let mut x_min = f64::MAX;
73        let mut x_max = f64::MIN;
74        let mut y_min = f64::MAX;
75        let mut y_max = f64::MIN;
76        let mut z_min = f64::MAX;
77        let mut z_max = f64::MIN;
78
79        for body in bodies {
80            let p = &body.posit();
81            x_min = x_min.min(p.x);
82            x_max = x_max.max(p.x);
83            y_min = y_min.min(p.y);
84            y_max = y_max.max(p.y);
85            z_min = z_min.min(p.z);
86            z_max = z_max.max(p.z);
87        }
88
89        x_min -= pad;
90        x_max += pad;
91        y_min -= pad;
92        y_max += pad;
93        z_min -= pad;
94        z_max += pad;
95
96        if z_offset {
97            z_max += 1e-5;
98        }
99
100        let x_size = x_max - x_min;
101        let y_size = y_max - y_min;
102        let z_size = z_max - z_min;
103
104        // Coerce to a cube.
105        let width = x_size.max(y_size).max(z_size);
106
107        let center = Vec3::new(
108            (x_max + x_min) / 2.,
109            (y_max + y_min) / 2.,
110            (z_max + z_min) / 2.,
111        );
112
113        Some(Self::new(center, width))
114    }
115
116    pub fn new(center: Vec3, width: f64) -> Self {
117        Self { center, width }
118    }
119
120    /// Divide this into equal-area octants.
121    pub(crate) fn divide_into_octants(&self) -> [Self; 8] {
122        let width = self.width / 2.;
123        let wd2 = self.width / 4.; // short for brevity below.
124
125        // Every combination of + and - for the center offset.
126        // The order matters, due to the binary index logic used when partitioning bodies into octants.
127        [
128            Self::new(self.center + Vec3::new(-wd2, -wd2, -wd2), width),
129            Self::new(self.center + Vec3::new(wd2, -wd2, -wd2), width),
130            Self::new(self.center + Vec3::new(-wd2, wd2, -wd2), width),
131            Self::new(self.center + Vec3::new(wd2, wd2, -wd2), width),
132            Self::new(self.center + Vec3::new(-wd2, -wd2, wd2), width),
133            Self::new(self.center + Vec3::new(wd2, -wd2, wd2), width),
134            Self::new(self.center + Vec3::new(-wd2, wd2, wd2), width),
135            Self::new(self.center + Vec3::new(wd2, wd2, wd2), width),
136        ]
137    }
138}
139
140#[derive(Debug)]
141pub struct Node {
142    /// We use `id` while building the tree, then sort by it, replacing with index.
143    /// Once complete, `id` == index in `Tree::nodes`.
144    /// Mass, center-of-mass, and body_ids include those from all sub-nodes.
145    pub id: usize,
146    pub bounding_box: Cube,
147    /// Node indices in the tree. We use this to guide the transversal process while finding
148    /// relevant nodes for a given target body.
149    pub children: Vec<usize>,
150    pub mass: f64,
151    pub center_of_mass: Vec3,
152    pub body_ids: Vec<usize>,
153}
154
155impl fmt::Display for Node {
156    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
157        write!(
158            f,
159            "Id: {}, Width: {:.3}, Ch: {:?}",
160            self.id, self.bounding_box.width, self.children
161        )
162    }
163}
164
165#[derive(Debug, Default)]
166/// A recursive tree. Each node can be subdivided  Terminates with `NodeType::NodeTerminal`.
167pub struct Tree {
168    /// Order matters; we index this by `Node::children`.
169    // Note: It doesn't appear that passing in a persistent, pre-allocated nodes Vec from the applicatoni
170    // has a significant impact on tree construction time.
171    pub nodes: Vec<Node>,
172}
173
174impl Tree {
175    /// Constructs a tree. Call this externaly using all bodies, once per time step.
176    /// It creates the entire tree, branching until each cell has `MAX_BODIES_PER_NODE` or fewer
177    /// bodies, or it reaches a maximum recursion depth.
178    ///
179    /// We partially transverse it as-required while calculating the force on a given target.
180    pub fn new<T: BodyModel>(bodies: &[T], bb: &Cube, config: &BhConfig) -> Self {
181        // Convert &[T] to &[&T].
182        let body_refs: Vec<&T> = bodies.iter().collect();
183
184        // todo: Refine this guess A/R.
185        // From an unrigorous benchmark, preallocating seems to be slightly faster, but not significantly so?
186        let mut nodes = Vec::with_capacity(bodies.len() * 7 / 4);
187
188        let mut current_node_i: usize = 0;
189
190        // Stack to simulate recursion: Each entry contains (bodies, bounding box, parent_id, child_index, depth).
191        let mut stack = Vec::new();
192
193        // body ids matches indexes with bodies.
194        let body_ids_init: Vec<usize> = body_refs.iter().enumerate().map(|(id, _)| id).collect();
195
196        stack.push((body_refs.to_vec(), body_ids_init, bb.clone(), None, 0));
197
198        while let Some((bodies_, body_ids, bb_, parent_id, depth)) = stack.pop() {
199            if depth > config.max_tree_depth {
200                break;
201            }
202            let (center_of_mass, mass) = center_of_mass(&bodies_);
203
204            let node_id = current_node_i;
205            nodes.push(Node {
206                id: node_id,
207                bounding_box: bb_.clone(),
208                mass,
209                center_of_mass,
210                children: Vec::new(),
211                body_ids: body_ids.clone(), // todo: The clone...
212            });
213
214            current_node_i += 1;
215
216            if let Some(pid) = parent_id {
217                // Rust is requesting an explicit type here.
218                let n: &mut Node = &mut nodes[pid];
219                n.children.push(node_id);
220            }
221
222            // If multiple (past our threshold) bodies are in this node, create an internal node and push its ID.
223            // Divide into octants and partition bodies. Otherwise, create a leaf node.
224            if bodies_.len() > config.max_bodies_per_node {
225                let octants = bb_.divide_into_octants();
226                let bodies_by_octant = partition(&bodies_, &body_ids, &bb_);
227
228                // Add each octant with bodies to the stack.
229                for (i, octant) in octants.into_iter().enumerate() {
230                    if !bodies_by_octant[i].is_empty() {
231                        let mut bto = Vec::with_capacity(bodies_by_octant[i].len());
232                        let mut ids_this_octant = Vec::with_capacity(bodies_by_octant[i].len());
233
234                        // todo: The clone etc?
235                        for (body, id) in &bodies_by_octant[i] {
236                            bto.push(*body);
237                            ids_this_octant.push(*id);
238                        }
239
240                        stack.push((bto, ids_this_octant, octant, Some(node_id), depth + 1));
241                    }
242                }
243            }
244        }
245
246        // Now that nodes are populated, rearrange so index == `id`. We will then index by `children`.
247        nodes.sort_by(|l, r| l.id.partial_cmp(&r.id).unwrap());
248
249        Self { nodes }
250    }
251
252    /// Get all leaves relevant to a given target. We use this to create a coarser
253    /// version of the tree, containing only the nodes we need to calculate acceleration
254    /// on a specific target.
255    pub fn leaves(&self, posit_target: Vec3, config: &BhConfig) -> Vec<&Node> {
256        let mut result = Vec::new();
257
258        if self.nodes.is_empty() {
259            return result;
260        }
261
262        let node_i = 0;
263
264        let mut stack = Vec::new();
265        stack.push(node_i);
266
267        while let Some(current_node_i) = stack.pop() {
268            let node = &self.nodes[current_node_i];
269
270            if node.children.len() <= config.max_bodies_per_node {
271                result.push(node);
272                continue;
273            }
274
275            let dist = (posit_target - node.center_of_mass).magnitude();
276
277            if node.bounding_box.width / dist < config.θ {
278                result.push(node);
279            } else {
280                // The source is near; add children to the stack to go deeper.
281                for &child_i in &node.children {
282                    stack.push(child_i);
283                }
284            }
285        }
286
287        result
288    }
289}
290
291/// Compute center of mass as a position, and mass value.
292fn center_of_mass<T: BodyModel>(bodies: &[&T]) -> (Vec3, f64) {
293    let mut mass = 0.;
294    let mut center_of_mass = Vec3::new_zero();
295
296    for body in bodies {
297        mass += body.mass();
298        center_of_mass += body.posit() * body.mass();
299    }
300
301    if mass.abs() > f64::EPSILON {
302        center_of_mass /= mass;
303    }
304
305    (center_of_mass, mass)
306}
307
308/// Partition bodies into each of the 8 octants.
309fn partition<'a, T: BodyModel>(
310    bodies: &[&'a T],
311    body_ids: &[usize],
312    bb: &Cube,
313) -> [Vec<(&'a T, usize)>; 8] {
314    let mut result: [Vec<(&'a T, usize)>; 8] = Default::default();
315
316    for (i, body) in bodies.iter().enumerate() {
317        let mut index = 0;
318        if body.posit().x > bb.center.x {
319            index |= 0b001;
320        }
321        if body.posit().y > bb.center.y {
322            index |= 0b010;
323        }
324        if body.posit().z > bb.center.z {
325            index |= 0b100;
326        }
327
328        result[index].push((body, body_ids[i]));
329    }
330
331    result
332}
333
334/// Calculate force using the Barnes Hut algorithm. The force function passed
335/// as a parameter has signature `(acc_dir: Vec3 (unit), mass_src: f64, distance: f64) -> Vec3`
336/// `id_target` is the index in the body array used to make the tree; it prevents self-interaction.
337/// Note that `mass` can be interchanged with `charge`, or similar.
338///
339/// When handling target mass or charge, reflect that in your `force_fn`; not here.
340pub fn run_bh<F>(
341    posit_target: Vec3,
342    id_target: usize,
343    tree: &Tree,
344    config: &BhConfig,
345    force_fn: &F,
346) -> Vec3
347where
348    F: Fn(Vec3, f64, f64) -> Vec3 + Send + Sync,
349{
350    tree.leaves(posit_target, config)
351        .par_iter()
352        .filter_map(|leaf| {
353            if leaf.body_ids.contains(&id_target) {
354                // Prevent self-interaction.
355                return None;
356            }
357
358            let acc_diff = leaf.center_of_mass - posit_target;
359            let dist = acc_diff.magnitude();
360
361            let acc_dir = acc_diff / dist; // Unit vec
362
363            Some(force_fn(acc_dir, leaf.mass, dist))
364        })
365        .reduce(Vec3::new_zero, |acc, elem| acc + elem)
366}
barnes_hut/lib.rs

barnes_hut/
lib.rs
