broccoli/queries/
nbody.rs

1//!
2//! Experimental nbody approximate solver
3//!
4//! The user can choose the distance at which to fallback on approximate solutions.
5//! The algorithm works similar to a Barnes–Hut simulation, but uses a kdtree instead of a quad tree.
6//!
7//! The user defines some geometric functions and their ideal accuracy.
8//!
9use super::*;
10
11type NodeWrapperVistr<'a, 'b, T, M> = VistrMut<'a, NodeWrapper<'b, T, M>, PreOrder>;
12
13///Helper enum indicating whether or not to gravitate a node as a whole, or as its individual parts.
14pub enum GravEnum<'a, T: Aabb, M> {
15    Mass(&'a mut M),
16    Bot(AabbPin<&'a mut [T]>),
17}
18
19///User defined functions for nbody
20pub trait Nbody {
21    type T: Aabb<Num = Self::N>;
22    type N: Num;
23    type Mass: Default + Copy + core::fmt::Debug;
24
25    //return the position of the center of mass
26    fn compute_center_of_mass(&mut self, a: &[Self::T]) -> Self::Mass;
27
28    fn is_close(&self, a: &Self::Mass, line: Self::N, a: impl Axis) -> bool;
29
30    fn is_close_half(&self, a: &Self::Mass, line: Self::N, a: impl Axis) -> bool;
31
32    fn gravitate(&mut self, a: GravEnum<Self::T, Self::Mass>, b: GravEnum<Self::T, Self::Mass>);
33
34    fn gravitate_self(&mut self, a: AabbPin<&mut [Self::T]>);
35
36    fn apply_a_mass<'a>(
37        &'a mut self,
38        mass: Self::Mass,
39        i: impl Iterator<Item = AabbPin<&'a mut Self::T>>,
40        len: usize,
41    );
42
43    fn combine_two_masses(&mut self, a: &Self::Mass, b: &Self::Mass) -> Self::Mass;
44}
45
46use compt::dfs_order::PreOrder;
47use compt::dfs_order::VistrMut;
48
49struct NodeWrapper<'a, T: Aabb, M> {
50    node: Node<'a, T, T::Num>,
51    mass: M,
52}
53
54fn build_masses2<N: Nbody>(vistr: NodeWrapperVistr<N::T, N::Mass>, no: &mut N) -> N::Mass {
55    let (nn, rest) = vistr.next();
56    let mass = no.compute_center_of_mass(&nn.node.range);
57    let mass = if let Some([left, right]) = rest {
58        let a = build_masses2(left, no);
59        let b = build_masses2(right, no);
60        let m = no.combine_two_masses(&a, &b);
61        no.combine_two_masses(&m, &mass)
62    } else {
63        mass
64    };
65    nn.mass = mass;
66    mass
67}
68
69fn collect_masses<'a, 'b, N: Nbody>(
70    root_div: N::N,
71    root_axis: impl Axis,
72    vistr: NodeWrapperVistr<'b, 'a, N::T, N::Mass>,
73    no: &mut N,
74    func1: &mut impl FnMut(&'b mut NodeWrapper<'a, N::T, N::Mass>, &mut N),
75    func2: &mut impl FnMut(&'b mut AabbPin<&'a mut [N::T]>, &mut N),
76) {
77    let (nn, rest) = vistr.next();
78
79    if !no.is_close_half(&nn.mass, root_div, root_axis) {
80        func1(nn, no);
81        return;
82    }
83
84    func2(&mut nn.node.range, no);
85
86    if let Some([left, right]) = rest {
87        collect_masses(root_div, root_axis, left, no, func1, func2);
88        collect_masses(root_div, root_axis, right, no, func1, func2);
89    }
90}
91
92fn pre_recc<N: Nbody>(
93    root_div: N::N,
94    root_axis: impl Axis,
95    root: &mut NodeWrapper<N::T, N::Mass>,
96    vistr: NodeWrapperVistr<N::T, N::Mass>,
97    no: &mut N,
98) {
99    let (nn, rest) = vistr.next();
100
101    if !no.is_close(&nn.mass, root_div, root_axis) {
102        no.gravitate(
103            GravEnum::Bot(root.node.range.borrow_mut()),
104            GravEnum::Mass(&mut nn.mass),
105        );
106        return;
107    }
108
109    no.gravitate(
110        GravEnum::Bot(root.node.range.borrow_mut()),
111        GravEnum::Bot(nn.node.range.borrow_mut()),
112    );
113
114    if let Some([left, right]) = rest {
115        pre_recc(root_div, root_axis, root, left, no);
116        pre_recc(root_div, root_axis, root, right, no);
117    }
118}
119
120fn recc_common<'a, 'b, N: Nbody>(
121    axis: impl Axis,
122    vistr: NodeWrapperVistr<'a, 'b, N::T, N::Mass>,
123    no: &mut N,
124) -> Option<[NodeWrapperVistr<'a, 'b, N::T, N::Mass>; 2]> {
125    let (nn, rest) = vistr.next();
126
127    no.gravitate_self(nn.node.range.borrow_mut());
128
129    if let Some([mut left, mut right]) = rest {
130        if let Some(div) = nn.node.div {
131            pre_recc(div, axis, nn, left.borrow_mut(), no);
132            pre_recc(div, axis, nn, right.borrow_mut(), no);
133
134            let mut finished_masses = Vec::new();
135            let mut finished_bots = Vec::new();
136
137            collect_masses(
138                div,
139                axis,
140                left.borrow_mut(),
141                no,
142                &mut |a, _| finished_masses.push(a),
143                &mut |a, _| finished_bots.push(a),
144            );
145
146            let mut finished_masses2 = Vec::new();
147            let mut finished_bots2 = Vec::new();
148
149            collect_masses(
150                div,
151                axis,
152                right.borrow_mut(),
153                no,
154                &mut |a, _| finished_masses2.push(a),
155                &mut |a, _| finished_bots2.push(a),
156            );
157
158            //We have collected masses on both sides.
159            //now gravitate all the ones on the left side with all the ones on the right side.
160
161            for a in finished_masses.into_iter() {
162                for b in finished_masses2.iter_mut() {
163                    no.gravitate(GravEnum::Mass(&mut a.mass), GravEnum::Mass(&mut b.mass));
164                }
165
166                for b in finished_bots2.iter_mut() {
167                    no.gravitate(GravEnum::Mass(&mut a.mass), GravEnum::Bot(b.borrow_mut()));
168                }
169            }
170            for a in finished_bots.into_iter() {
171                for b in finished_masses2.iter_mut() {
172                    no.gravitate(GravEnum::Bot(a.borrow_mut()), GravEnum::Mass(&mut b.mass));
173                }
174                for b in finished_bots2.iter_mut() {
175                    no.gravitate(GravEnum::Bot(a.borrow_mut()), GravEnum::Bot(b.borrow_mut()));
176                }
177            }
178
179            //parallelize this
180            Some([left, right])
181        } else {
182            None
183        }
184    } else {
185        None
186    }
187}
188
189fn recc<N: Nbody>(
190    axis: impl Axis,
191    vistr: VistrMut<NodeWrapper<N::T, N::Mass>, PreOrder>,
192    no: &mut N,
193) {
194    let keep_going = recc_common(axis, vistr, no);
195
196    if let Some([left, right]) = keep_going {
197        recc(axis.next(), left, no);
198        recc(axis.next(), right, no);
199    }
200}
201
202fn apply_tree<N: Nbody>(mut vistr: NodeWrapperVistr<N::T, N::Mass>, no: &mut N) {
203    {
204        let mass = vistr.borrow_mut().next().0.mass;
205
206        let len = vistr
207            .borrow_mut()
208            .dfs_preorder_iter()
209            .map(|x| x.node.range.borrow_mut().len())
210            .sum();
211
212        let it = vistr
213            .borrow_mut()
214            .dfs_preorder_iter()
215            .flat_map(|x| x.node.range.borrow_mut().iter_mut());
216
217        no.apply_a_mass(mass, it, len);
218    }
219
220    let (_, rest) = vistr.next();
221
222    if let Some([left, right]) = rest {
223        apply_tree(left, no);
224        apply_tree(right, no);
225    }
226}
227
228impl<'a, T: Aabb> crate::Tree<'a, T> {
229    pub fn handle_nbody<N: Nbody<T = T>>(&mut self, no: &mut N) {
230        ///Perform nbody
231        ///The tree is taken by value so that its nodes can be expended to include more data.
232        pub fn nbody_mut<'a, N: Nbody<T = T>, T: Aabb>(
233            tree: Box<[Node<'a, T, T::Num>]>,
234            no: &mut N,
235        ) -> Box<[Node<'a, T, T::Num>]> {
236            let mut newnodes: Vec<_> = Vec::from(tree)
237                .into_iter()
238                .map(|x| NodeWrapper {
239                    node: x,
240                    mass: Default::default(),
241                })
242                .collect();
243
244            let tree = compt::dfs_order::CompleteTreeMut::from_preorder_mut(&mut newnodes).unwrap();
245            let mut vistr = tree.vistr_mut();
246
247            //calculate node masses of each node.
248            build_masses2(vistr.borrow_mut(), no);
249
250            recc(default_axis(), vistr.borrow_mut(), no);
251
252            apply_tree(vistr, no);
253
254            newnodes.into_iter().map(|x| x.node).collect()
255        }
256
257        let mut vvv = vec![].into_boxed_slice();
258        std::mem::swap(&mut self.nodes, &mut vvv);
259
260        let mut new = nbody_mut(vvv, no);
261        std::mem::swap(&mut self.nodes, &mut new);
262    }
263}
264
265mod assert {
266    use super::*;
267    impl<'a, T: Aabb> Naive<'a, T> {
268        pub fn handle_nbody<N: Nbody<T = T>>(&mut self, no: &mut N) {
269            ///Naive version simply visits every pair.
270            pub fn naive_nbody_mut<T: Aabb>(
271                bots: AabbPin<&mut [T]>,
272                func: impl FnMut(AabbPin<&mut T>, AabbPin<&mut T>),
273            ) {
274                queries::for_every_pair(bots, func);
275            }
276
277            naive_nbody_mut(self.inner.borrow_mut(), |a, b| {
278                no.gravitate(GravEnum::Bot(a.into_slice()), GravEnum::Bot(b.into_slice()));
279            });
280        }
281    }
282}
broccoli/queries/nbody.rs

broccoli/queries/
nbody.rs
