rusty_tree/
octree.rs

1//! Data structures and functions to create regular and adaptive Octrees.
2
3use ndarray::{Array1, ArrayView1, ArrayView2, Axis};
4use rusty_kernel_tools::RealType;
5use std::collections::{HashMap, HashSet};
6use std::fmt;
7use std::time::Duration;
8
9pub mod adaptive_octree;
10pub mod regular_octree;
11
12pub use adaptive_octree::*;
13pub use regular_octree::*;
14
15/// The type of the octree.
16pub enum OctreeType {
17    /// Regular octree.
18    Regular,
19    /// Use for balanced adaptive octrees.
20    BalancedAdaptive,
21    /// Use for unbalanced adaptive octrees.
22    UnbalancedAdaptive,
23}
24
25/// The basic Octree data structure
26pub struct Octree<'a, T: RealType> {
27    /// A (3, N) array of N particles.
28    pub particles: ArrayView2<'a, T>,
29
30    /// The maximum level in the tree.
31    pub max_level: usize,
32
33    /// The origin of the bounding box for the particles.
34    pub origin: [f64; 3],
35
36    /// The diameter across each dimension of the bounding box.
37    pub diameter: [f64; 3],
38
39    /// The non-empty keys for each level of the tree.
40    pub level_keys: HashMap<usize, HashSet<usize>>,
41
42    /// The keys associated with the particles.
43    pub particle_keys: Array1<usize>,
44
45    /// The set of near-field keys for each non-empty key.
46    pub near_field: HashMap<usize, HashSet<usize>>,
47
48    /// The set of keys in the interaction list for each non-empty key.
49    pub interaction_list: HashMap<usize, HashSet<usize>>,
50
51    /// The index set of particles associated with each leaf key.
52    pub leaf_key_to_particles: HashMap<usize, HashSet<usize>>,
53
54    /// The set of all non-empty keys in the tree.
55    pub all_keys: HashSet<usize>,
56
57    /// The type of the Octree.
58    pub octree_type: OctreeType,
59
60    /// Statistics for the tree.
61    pub statistics: Statistics,
62}
63
64/// A structure that stores various statistics for a tree.
65pub struct Statistics {
66    pub number_of_particles: usize,
67    pub max_level: usize,
68    pub number_of_leafs: usize,
69    pub number_of_keys: usize,
70    pub creation_time: Duration,
71    pub minimum_number_of_particles_in_leaf: usize,
72    pub maximum_number_of_particles_in_leaf: usize,
73    pub average_number_of_particles_in_leaf: f64,
74}
75
76impl std::fmt::Display for Statistics {
77    /// Create an output string for tree statistics.
78    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
79        write!(
80            f,
81            "\n\nOctree Statistics\n\
82                   ==============================\n\
83                   Number of Particles: {}\n\
84                   Maximum level: {}\n\
85                   Number of leaf keys: {}\n\
86                   Number of keys in tree: {}\n\
87                   Creation time [s]: {}\n\
88                   Minimum number of particles in leaf node: {}\n\
89                   Maximum number of particles in leaf node: {}\n\
90                   Average number of particles in leaf node: {:.2}\n\
91                   ==============================\n\n
92                   ",
93            self.number_of_particles,
94            self.max_level,
95            self.number_of_leafs,
96            self.number_of_keys,
97            (self.creation_time.as_millis() as f64) / 1000.0,
98            self.minimum_number_of_particles_in_leaf,
99            self.maximum_number_of_particles_in_leaf,
100            self.average_number_of_particles_in_leaf
101        )
102    }
103}
104
105/// Given the set of all keys, compute the interaction list for each key.
106///
107/// Returns a map from keys to the corresponding interaction list, represented by
108/// a set of keys.
109pub fn compute_interaction_list_map(all_keys: &HashSet<usize>) -> HashMap<usize, HashSet<usize>> {
110    use crate::morton::compute_interaction_list;
111    use rayon::prelude::*;
112
113    let mut interaction_list_map = HashMap::<usize, HashSet<usize>>::new();
114
115    for &key in all_keys {
116        interaction_list_map.insert(key, HashSet::<usize>::new());
117    }
118
119    interaction_list_map
120        .par_iter_mut()
121        .for_each(|(&key, hash_set)| {
122            let current_interaction_list = compute_interaction_list(key);
123            hash_set.extend(&current_interaction_list);
124        });
125
126    interaction_list_map
127}
128
129/// Given the set of all keys, compute the near field for each key.
130///
131/// Returns a map from keys to the corresponding near field, represented by
132/// a set of keys.
133pub fn compute_near_field_map(all_keys: &HashSet<usize>) -> HashMap<usize, HashSet<usize>> {
134    use crate::morton::compute_near_field;
135    use rayon::prelude::*;
136
137    let mut near_field_map = HashMap::<usize, HashSet<usize>>::new();
138
139    for &key in all_keys {
140        near_field_map.insert(key, HashSet::<usize>::new());
141    }
142
143    near_field_map.par_iter_mut().for_each(|(&key, hash_set)| {
144        let current_near_field = compute_near_field(key);
145        hash_set.extend(&current_near_field);
146    });
147
148    near_field_map
149}
150
151/// Compute the leaf map.
152///
153/// Returns a map from leaf keys to associated particle indices.
154pub fn compute_leaf_map(particle_keys: ArrayView1<usize>) -> HashMap<usize, HashSet<usize>> {
155    use itertools::Itertools;
156
157    let mut leaf_key_to_particles = HashMap::<usize, HashSet<usize>>::new();
158    for &key in particle_keys.iter().unique() {
159        leaf_key_to_particles.insert(key, HashSet::<usize>::new());
160    }
161
162    for (index, key) in particle_keys.iter().enumerate() {
163        leaf_key_to_particles.get_mut(key).unwrap().insert(index);
164    }
165
166    leaf_key_to_particles
167}
168
169/// Given an array of keys. Return the level information of the tree.
170///
171/// The function returns a 3-tuple `(max_level, all_keys, level_keys)`.
172/// `max_level` us a `usize` that contains the maximum level of the keys.
173/// The set `all_keys` contains all keys from the tree by completing the tree
174/// from the leaf onwards to the top and storing all parent keys along the way.
175/// The map `level_keys` is a map from a given level to the set of all keys contained
176/// in the level.
177pub fn compute_level_information(
178    particle_keys: ArrayView1<usize>,
179) -> (usize, HashSet<usize>, HashMap<usize, HashSet<usize>>) {
180    use crate::morton::{find_level, find_parent};
181    use std::iter::FromIterator;
182
183    let max_level = particle_keys
184        .iter()
185        .map(|&item| find_level(item))
186        .max()
187        .unwrap();
188
189    let nlevels = max_level + 1;
190    let leaf_keys: HashSet<usize> = particle_keys.iter().copied().collect();
191    let mut level_keys = HashMap::<usize, HashSet<usize>>::new();
192
193    // Distribute the keys among different sets for each level
194    for current_level in 0..nlevels {
195        level_keys.insert(
196            current_level,
197            HashSet::from_iter(
198                leaf_keys
199                    .iter()
200                    .filter(|&&key| find_level(key) == current_level)
201                    .copied(),
202            ),
203        );
204    }
205
206    // Now fill up the sets with all the various parent keys.
207    for current_level in (1..nlevels).rev() {
208        let parent_index = current_level - 1;
209        let current_set: HashSet<usize> = level_keys
210            .get(&current_level)
211            .unwrap()
212            .iter()
213            .copied()
214            .collect();
215        let parent_set = level_keys.get_mut(&parent_index).unwrap();
216        parent_set.extend(current_set.iter().map(|&key| find_parent(key)));
217    }
218
219    let mut all_keys = HashSet::<usize>::new();
220
221    for (_, current_keys) in level_keys.iter() {
222        all_keys.extend(current_keys.iter());
223    }
224
225    (max_level, all_keys, level_keys)
226}
227
228/// Create a regular octree from an array of particles.
229///
230/// Returns the set of all non-empty keys associated with a regular octree created
231/// from an array of particles.
232/// # Arguments
233/// * `particles` - A (3, N) array of particles
234/// * `max_level` - The deepest level of the tree.
235/// * `origin` - The origin of the bounding box.
236/// * `diameter` - The diameter of the bounding box in each dimension.
237pub fn compute_complete_regular_tree<T: RealType>(
238    particles: ArrayView2<T>,
239    max_level: usize,
240    origin: &[f64; 3],
241    diameter: &[f64; 3],
242) -> HashSet<usize> {
243    use super::morton::encode_points;
244
245    let keys = encode_points(particles, max_level, origin, diameter);
246    let (_, all_keys, _) = compute_level_information(keys.view());
247
248    all_keys
249}
250
251/// Export an octree to vtk
252pub fn export_to_vtk<T: RealType>(tree: &Octree<T>, filename: &str) {
253    use super::morton::{serialize_box_from_key};
254    use std::iter::FromIterator;
255    use vtkio::model::*;
256    use std::path::PathBuf;
257
258    let filename = filename.to_owned();
259
260    // We use a vtk voxel type, which has
261    // 8 points per cell, i.e. 24 float numbers
262    // per cell.
263    const POINTS_PER_CELL: usize = 8;
264
265    let num_particles = tree.particles.len_of(Axis(1));
266
267    let all_keys = &tree.all_keys;
268    let num_keys = all_keys.len();
269
270    let num_floats = 3 * POINTS_PER_CELL * num_keys;
271
272    let mut cell_points = Vec::<f64>::with_capacity(num_floats);
273
274    for &key in all_keys {
275        let serialized = serialize_box_from_key(key, &tree.origin, &tree.diameter);
276        cell_points.extend(serialized);
277    }
278
279    for index in 0..num_particles {
280        cell_points.push(tree.particles[[0, index]].to_f64().unwrap());
281        cell_points.push(tree.particles[[1, index]].to_f64().unwrap());
282        cell_points.push(tree.particles[[2, index]].to_f64().unwrap());
283    }
284
285    let num_points = 8 * (num_keys as u64) + (num_particles as u64);
286
287    let connectivity = Vec::<u64>::from_iter(0..num_points);
288    let mut offsets = Vec::<u64>::from_iter((0..(num_keys as u64)).map(|item| 8 * item + 8));
289    offsets.push(num_points);
290
291    let mut types = vec![CellType::Voxel; num_keys];
292    types.push(CellType::PolyVertex);
293
294    let mut cell_data = Vec::<i32>::with_capacity(num_points as usize);
295
296    for _ in 0..num_keys {
297        cell_data.push(0);
298    }
299    cell_data.push(1);
300
301
302    let model = Vtk {
303        version: Version { major: 1, minor: 0 },
304        title: String::new(),
305        byte_order: ByteOrder::BigEndian,
306        file_path: Some(PathBuf::from(&filename)),
307        data: DataSet::inline(UnstructuredGridPiece {
308            points: IOBuffer::F64(cell_points),
309            cells: Cells {
310                cell_verts: VertexNumbers::XML {
311                    connectivity: connectivity,
312                    offsets: offsets,
313                },
314                types: types,
315            },
316            data: Attributes {
317                point: vec![],
318                cell: vec![Attribute::DataArray(DataArrayBase {
319                    name: String::from("colors"),
320                    elem: ElementType::Scalars {
321                        num_comp: 1,
322                        lookup_table: None,
323                    },
324                    data: IOBuffer::I32(cell_data),
325                })],
326            },
327        }),
328    };
329
330
331    model.export(filename).unwrap();
332}