Skip to main content

barnes_hut_tree/
lib.rs

1#![cfg_attr(not(test), no_std)]
2#![doc = include_str!("../README.md")]
3
4extern crate alloc;
5
6mod arena;
7mod morton;
8
9pub use arena::BarnesHutTree;
10pub use morton::{Dim, Morton, MortonWord};
11
12/// Per-cell force contribution, the pluggable half of the Barnes-Hut traversal.
13///
14/// [`BarnesHutTree::accumulate_forces`] and [`BarnesHutTree::accumulate_all`] call
15/// [`ForceKernel::accumulate`] once per summarised cell (a leaf, or an internal cell
16/// passing the theta test), handing it the cell's summary. The kernel turns that summary
17/// into a force, so the force law lives entirely in the kernel and never in the tree.
18///
19/// The blanket implementation covers any closure of the matching shape, `#[inline]`d so
20/// a monomorphised call compiles down to inline arithmetic.
21///
22/// # Example
23///
24/// ```
25/// use barnes_hut_tree::BarnesHutTree;
26///
27/// // Inverse-cube repulsion: push the query away from a summarised cell with magnitude
28/// // `mass / distance^3` on the displacement, one axis at a time.
29/// let kernel = |displacement: &[f32; 2],
30///               distance_squared: f32,
31///               mass: f32,
32///               _is_leaf: bool,
33///               force: &mut [f32; 2]| {
34///     let scale = mass / (distance_squared * distance_squared.sqrt());
35///     force[0] += scale * displacement[0];
36///     force[1] += scale * displacement[1];
37/// };
38///
39/// let points = [0.0f32, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0];
40/// let tree = BarnesHutTree::<f32, u64, 2>::new_uniform(&points);
41/// let mut forces = [0.0f32; 8];
42/// tree.accumulate_all(&points, 0.5, &kernel, &mut forces);
43/// assert!(forces.iter().all(|value| value.is_finite()));
44/// ```
45pub trait ForceKernel<T, const D: usize> {
46    /// Adds the force from a cell of summed `mass` whose center of mass sits at
47    /// `displacement` from the query (the vector `query - center_of_mass`, squared length
48    /// `distance_squared`) into `force`. `distance_squared` is always strictly positive:
49    /// the traversal skips zero-distance cells before calling the kernel. `is_leaf` is
50    /// true for a leaf holding actual points, false for an internal summary. A kernel that
51    /// treats near contact differently from the far field (as FM3 does, doubling the leaf
52    /// charge) uses this flag.
53    fn accumulate(
54        &self,
55        displacement: &[T; D],
56        distance_squared: T,
57        mass: T,
58        is_leaf: bool,
59        force: &mut [T; D],
60    );
61}
62
63impl<T, const D: usize, F> ForceKernel<T, D> for F
64where
65    F: Fn(&[T; D], T, T, bool, &mut [T; D]),
66{
67    #[inline]
68    fn accumulate(
69        &self,
70        displacement: &[T; D],
71        distance_squared: T,
72        mass: T,
73        is_leaf: bool,
74        force: &mut [T; D],
75    ) {
76        self(displacement, distance_squared, mass, is_leaf, force);
77    }
78}