fastiron/geometry/
grid_assignment_object.rs1use std::collections::VecDeque;
8
9use num::{zero, FromPrimitive};
10
11use crate::{
12 constants::{CustomFloat, Tuple3},
13 data::mc_vector::MCVector,
14};
15
16#[derive(Debug, Clone, Default)] struct GridCell {
18 pub burned: bool,
19 pub my_centers: Vec<usize>,
20}
21
22#[derive(Debug)]
24pub struct GridAssignmentObject<T: CustomFloat> {
25 pub nx: usize,
27 pub ny: usize,
29 pub nz: usize,
31 pub dx: T,
33 pub dy: T,
35 pub dz: T,
37
38 pub corner: MCVector<T>,
40 pub centers: Vec<MCVector<T>>,
42
43 grid: Vec<GridCell>,
45 flood_queue: VecDeque<usize>,
47 wet_list: VecDeque<usize>,
49}
50
51impl<T: CustomFloat> GridAssignmentObject<T> {
52 pub fn new(centers: &[MCVector<T>]) -> Self {
54 let centers_per_cell: T = FromPrimitive::from_usize(5).unwrap();
56 let n_centers: T = FromPrimitive::from_usize(centers.len()).unwrap();
57 let one: T = FromPrimitive::from_f64(1.0).unwrap();
58
59 let mut min_coords = centers[0];
60 let mut max_coords = centers[0];
61 centers.iter().for_each(|vv| {
62 min_coords.x = min_coords.x.min(vv.x);
63 min_coords.y = min_coords.y.min(vv.y);
64 min_coords.z = min_coords.z.min(vv.z);
65 max_coords.x = max_coords.x.max(vv.x);
66 max_coords.y = max_coords.y.max(vv.y);
67 max_coords.z = max_coords.z.max(vv.z);
68 });
69
70 let lx = one.max(max_coords.x - min_coords.x);
71 let ly = one.max(max_coords.y - min_coords.y);
72 let lz = one.max(max_coords.z - min_coords.z);
73 let dim: T = (n_centers / (centers_per_cell * lx * ly * lz)).cbrt(); let nx = one.max((dim * lx).floor());
75 let ny = one.max((dim * ly).floor());
76 let nz = one.max((dim * lz).floor());
77 let dx = lx / nx;
78 let dy = ly / ny;
79 let dz = lz / nz;
80
81 let n_cells: usize = (nx * ny * nz).to_usize().unwrap();
82 let grid: Vec<GridCell> = vec![GridCell::default(); n_cells];
83
84 let mut gao: GridAssignmentObject<T> = Self {
85 nx: nx.to_usize().unwrap(),
86 ny: ny.to_usize().unwrap(),
87 nz: nz.to_usize().unwrap(),
88 dx,
89 dy,
90 dz,
91 corner: min_coords,
92 centers: centers.to_vec(),
93 grid,
94 flood_queue: Default::default(),
95 wet_list: Default::default(),
96 };
97
98 (0..centers.len()).for_each(|center_idx| {
99 let cell_idx = gao.which_cell(centers[center_idx]);
100 gao.grid[cell_idx].my_centers.push(center_idx);
101 });
102
103 gao
104 }
105
106 pub fn nearest_center(&mut self, rr: MCVector<T>) -> usize {
112 let mut r2_min: T = FromPrimitive::from_f64(1e300).unwrap();
113 let mut center_min: Option<usize> = None;
114
115 self.add_tuple_to_queue(self.which_cell_tuple(rr));
116
117 while !self.flood_queue.is_empty() {
118 let cell_idx: usize = self.flood_queue.pop_front().unwrap();
120 if self.min_dist2(rr, cell_idx) > r2_min {
122 continue;
123 }
124 for center_idx in &self.grid[cell_idx].my_centers {
125 let center_r = self.centers[*center_idx];
126 let r2: T = (rr - center_r).dot(&(rr - center_r));
127 if r2 == r2_min {
128 center_min.map(|m| m.min(*center_idx));
129 }
130 if r2 < r2_min {
131 r2_min = r2;
132 center_min = Some(*center_idx);
133 }
134 }
135 self.add_nbrs_to_queue(cell_idx);
136 }
137
138 while !self.wet_list.is_empty() {
139 self.grid[self.wet_list.pop_front().unwrap()].burned = false;
140 }
141
142 assert!(center_min.is_some());
143 center_min.unwrap()
144 }
145
146 fn which_cell_tuple(&self, r: MCVector<T>) -> Tuple3 {
148 let mut ix: usize = ((r.x - self.corner.x) / self.dx)
149 .floor()
150 .max(zero())
151 .to_usize()
152 .unwrap();
153 let mut iy: usize = ((r.y - self.corner.y) / self.dy)
154 .floor()
155 .max(zero())
156 .to_usize()
157 .unwrap();
158 let mut iz: usize = ((r.z - self.corner.z) / self.dz)
159 .floor()
160 .max(zero())
161 .to_usize()
162 .unwrap();
163 ix = ix.min(self.nx - 1);
164 iy = iy.min(self.ny - 1);
165 iz = iz.min(self.nz - 1);
166
167 (ix, iy, iz)
168 }
169
170 fn which_cell(&self, r: MCVector<T>) -> usize {
172 self.tuple_to_index(&self.which_cell_tuple(r))
173 }
174
175 fn tuple_to_index(&self, t: &Tuple3) -> usize {
177 t.0 + self.nx * (t.1 + self.ny * t.2)
178 }
179
180 fn index_to_tuple(&self, idx: usize) -> Tuple3 {
182 let ix: usize = idx % self.nx;
183 let tmp: usize = idx / self.nx;
184 let iy: usize = tmp % self.ny;
185 let iz: usize = tmp / self.ny;
186 (ix, iy, iz)
187 }
188
189 fn min_dist2(&self, r: MCVector<T>, cell_idx: usize) -> T {
192 let r_idx: Tuple3 = self.which_cell_tuple(r);
193 let tuple_idx: Tuple3 = self.index_to_tuple(cell_idx);
194
195 let rx: T = self.dx
196 * (FromPrimitive::from_usize(tuple_idx.0.abs_diff(r_idx.0).saturating_sub(1))).unwrap();
197 let ry: T = self.dy
198 * (FromPrimitive::from_usize(tuple_idx.1.abs_diff(r_idx.1).saturating_sub(1))).unwrap();
199 let rz: T = self.dz
200 * (FromPrimitive::from_usize(tuple_idx.2.abs_diff(r_idx.2).saturating_sub(1))).unwrap();
201 assert!(rx >= zero());
202 assert!(ry >= zero());
203 assert!(rz >= zero());
204 rx * rx + ry * ry + rz * rz
205 }
206
207 fn add_tuple_to_queue(&mut self, t: Tuple3) {
209 let idx: usize = self.tuple_to_index(&t);
210 if self.grid[idx].burned {
211 return;
212 }
213 self.flood_queue.push_back(idx);
214 self.wet_list.push_back(idx);
215 self.grid[idx].burned = true
216 }
217
218 fn add_nbrs_to_queue(&mut self, cell_idx: usize) {
220 let tuple_idx: Tuple3 = self.index_to_tuple(cell_idx);
221 if tuple_idx.0 + 1 < self.nx {
223 let mut tmp: Tuple3 = tuple_idx;
224 tmp.0 += 1;
225 self.add_tuple_to_queue(tmp);
226 }
227 if tuple_idx.0 > 0 {
228 let mut tmp: Tuple3 = tuple_idx;
229 tmp.0 -= 1;
230 self.add_tuple_to_queue(tmp);
231 }
232 if tuple_idx.1 + 1 < self.ny {
234 let mut tmp: Tuple3 = tuple_idx;
235 tmp.1 += 1;
236 self.add_tuple_to_queue(tmp);
237 }
238 if tuple_idx.1 > 0 {
239 let mut tmp: Tuple3 = tuple_idx;
240 tmp.1 -= 1;
241 self.add_tuple_to_queue(tmp);
242 }
243 if tuple_idx.2 + 1 < self.nz {
245 let mut tmp: Tuple3 = tuple_idx;
246 tmp.2 += 1;
247 self.add_tuple_to_queue(tmp);
248 }
249 if tuple_idx.2 > 0 {
250 let mut tmp: Tuple3 = tuple_idx;
251 tmp.2 -= 1;
252 self.add_tuple_to_queue(tmp);
253 }
254 }
255}
256
257#[cfg(test)]
262mod tests {
263 use num::Float;
264
265 use super::*;
266 #[test]
267 fn basic() {
268 let c1 = MCVector {
269 x: 0.0,
270 y: 0.0,
271 z: 0.0,
272 };
273 let c2 = MCVector {
274 x: 2.0,
275 y: 1.0,
276 z: 0.0,
277 };
278 let c3 = MCVector {
279 x: 1.0,
280 y: 2.0,
281 z: 4.0,
282 };
283 let c4 = MCVector {
284 x: 0.0,
285 y: 3.0,
286 z: 1.0,
287 };
288 let assigner = GridAssignmentObject::new(&[c1, c2, c3, c4]);
289
290 let lx = 2.0;
291 let ly = 3.0;
292 let lz = 4.0;
293 let center_per_cell = 5.0;
294 let cell_per_length = (4.0 / (center_per_cell * lx * ly * lz)).cbrt();
295 let nx = 1.0.max((lx * cell_per_length).floor());
296 let ny = 1.0.max((ly * cell_per_length).floor());
297 let nz = 1.0.max((lz * cell_per_length).floor());
298 let dx = lx / nx;
299 let dy = ly / ny;
300 let dz = lz / nz;
301 assert_eq!(dx, assigner.dx);
302 assert_eq!(dy, assigner.dy);
303 assert_eq!(dz, assigner.dz);
304 }
305}