1use 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
15pub enum OctreeType {
17 Regular,
19 BalancedAdaptive,
21 UnbalancedAdaptive,
23}
24
25pub struct Octree<'a, T: RealType> {
27 pub particles: ArrayView2<'a, T>,
29
30 pub max_level: usize,
32
33 pub origin: [f64; 3],
35
36 pub diameter: [f64; 3],
38
39 pub level_keys: HashMap<usize, HashSet<usize>>,
41
42 pub particle_keys: Array1<usize>,
44
45 pub near_field: HashMap<usize, HashSet<usize>>,
47
48 pub interaction_list: HashMap<usize, HashSet<usize>>,
50
51 pub leaf_key_to_particles: HashMap<usize, HashSet<usize>>,
53
54 pub all_keys: HashSet<usize>,
56
57 pub octree_type: OctreeType,
59
60 pub statistics: Statistics,
62}
63
64pub 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 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
105pub 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(¤t_interaction_list);
124 });
125
126 interaction_list_map
127}
128
129pub 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(¤t_near_field);
146 });
147
148 near_field_map
149}
150
151pub 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
169pub 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 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 for current_level in (1..nlevels).rev() {
208 let parent_index = current_level - 1;
209 let current_set: HashSet<usize> = level_keys
210 .get(¤t_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
228pub 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
251pub 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 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}