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}