fastiron/geometry/
grid_assignment_object.rs

1//! Code used to locate vectors in the grid
2//!
3//! This module contains the structure used to locate vectors in the
4//! mesh's grid, i.e. locate particles in the problem and operates on them
5//! accordingly to the properties of the space (cell) they belong to.
6
7use 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)] // default value of bool is false
17struct GridCell {
18    pub burned: bool,
19    pub my_centers: Vec<usize>,
20}
21
22/// Structure used to locate vectors, i.e. coordinates, in the grid.
23#[derive(Debug)]
24pub struct GridAssignmentObject<T: CustomFloat> {
25    /// Number of cells along the x-axis.
26    pub nx: usize,
27    /// Number of cells along the y-axis.
28    pub ny: usize,
29    /// Number of cells along the z-axis.
30    pub nz: usize,
31    /// Size of a mesh cell along the x-axis (cm).
32    pub dx: T,
33    /// Size of a mesh cell along the y-axis (cm).
34    pub dy: T,
35    /// Size of a mesh cell along the z-axis (cm).
36    pub dz: T,
37
38    /// List of corners.
39    pub corner: MCVector<T>,
40    /// List of centers.
41    pub centers: Vec<MCVector<T>>,
42
43    /// List of cells.
44    grid: Vec<GridCell>,
45    /// Internal queue used when browsing through the cells.
46    flood_queue: VecDeque<usize>,
47    /// Internal queue used when browsing through the cells.
48    wet_list: VecDeque<usize>,
49}
50
51impl<T: CustomFloat> GridAssignmentObject<T> {
52    /// Constructor.
53    pub fn new(centers: &[MCVector<T>]) -> Self {
54        // sets the length scale of the grid cells
55        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(); // cell per unit of length
74        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    /// Returns the closest center to a given coordinate. This is done by browsing
107    /// through the cell, doing a tweaked breadth-first search on the graph. The
108    /// algorithm is essentially the same, with just one additional condition
109    /// to keep searching "towards" a given direction: for the minimal distance
110    /// to diminish.
111    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            // next cell to check
119            let cell_idx: usize = self.flood_queue.pop_front().unwrap();
120            // if cell is too far away, don't even try
121            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    /// Returns the tuple of the cell the coordinate belongs to.
147    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    /// Returns the index of the cell the coordinate belongs to.
171    fn which_cell(&self, r: MCVector<T>) -> usize {
172        self.tuple_to_index(&self.which_cell_tuple(r))
173    }
174
175    /// Converts a cell tuple to its index.
176    fn tuple_to_index(&self, t: &Tuple3) -> usize {
177        t.0 + self.nx * (t.1 + self.ny * t.2)
178    }
179
180    /// Converts a cell index to its tuple.
181    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    /// Finds a lower bound of the squared distance from the point
190    /// r to the cell with index cell_idx.
191    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    /// ?
208    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    /// Add valid tuple to internal queues.
219    fn add_nbrs_to_queue(&mut self, cell_idx: usize) {
220        let tuple_idx: Tuple3 = self.index_to_tuple(cell_idx);
221        // on x
222        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        // on y
233        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        // on z
244        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//=============
258// Unit tests
259//=============
260
261#[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}