spatial_neighbors/
grid.rs

1use std::ops::Range;
2use crate::SpatialPartitioner;
3use crate::util::in_range;
4
5pub struct Grid<Data: Copy> {
6    cells: Vec<Vec<((f64, f64), Data)>>,
7
8    cell_count: (usize, usize),
9    cell_size: (f64, f64),
10
11    x: Range<f64>,
12    y: Range<f64>,
13
14    count: usize,
15}
16
17impl<Data: Copy> Grid<Data> {
18    ///
19    /// # Arguments
20    ///
21    /// * `x`: min_x..max_x defines the area in wich data can be inserted
22    /// * `y`: min_y..max_y defines the area in wich data can be inserted
23    /// * `cell_count`: (count_x, count_y) defines how many cell should be present
24    ///
25    pub fn with_cell_count(x: Range<f64>, y: Range<f64>, cell_count: (usize, usize)) -> Grid<Data> {
26        let mut cells = Vec::new();
27
28        for i in 0..(cell_count.0 * cell_count.1) {
29            cells.insert(i as usize, Vec::new());
30        };
31
32        Grid {
33            cells,
34            cell_count,
35            cell_size: ((x.end - x.start) as f64 / cell_count.0 as f64, (y.end - y.start) as f64 / cell_count.1 as f64),
36            x,
37            y,
38            count: 0,
39        }
40    }
41
42    ///
43    /// # Arguments
44    ///
45    /// * `indexes`: (x,y) cord of the cell
46    ///
47    /// returns: &Vec<((f64, f64), Data), Global> the data stored in the cell
48    ///
49    pub fn cell(&self, indexes: (usize, usize)) -> &Vec<((f64, f64), Data)> {
50        let index = indexes.0 + (indexes.1 * self.cell_count.0);
51
52        self.cells.get(index).unwrap()
53    }
54
55    ///
56    /// # Arguments
57    ///
58    /// returns: (usize, usize) the defined cell_count
59    ///
60    pub fn cell_count(&self) -> (usize, usize) {
61        self.cell_count
62    }
63
64    ///
65    /// # Arguments
66    ///
67    /// returns: (f64, f64) width and height of each cell
68    ///
69    pub fn cell_size(&self) -> (f64, f64) {
70        self.cell_size
71    }
72
73    fn pos_to_index(&self, position: (f64, f64)) -> (usize, usize) {
74        let x = ((position.0 - self.x.start) / self.cell_size.0).floor() as usize;
75        let y = ((position.1 - self.y.start) / self.cell_size.1).floor() as usize;
76
77        (x, y)
78    }
79
80    fn index_to_pos(&self, position: (usize, usize)) -> (f64, f64) {
81        let x = self.cell_size.0 * position.0 as f64 + self.x.start;
82        let y = self.cell_size.1 * position.1 as f64 + self.y.start;
83
84        (x, y)
85    }
86}
87
88impl<Data: Copy> SpatialPartitioner<Data> for Grid<Data> {
89    ///
90    /// create a Grid with a default size of (100,100). More info here [`Grid::new()`]
91    ///
92    fn new(x: Range<f64>, y: Range<f64>) -> Self {
93        Grid::with_cell_count(x, y, (100, 100))
94    }
95
96    fn insert(&mut self, position: (f64, f64), data: Data) {
97        if position.0 < self.x.start || position.0 >= self.x.end || position.1 < self.y.start || position.1 >= self.y.end {
98            panic!("tried to insert position into SpatialHash which was out of bounce")
99        }
100
101        self.insert_unchecked(position, data);
102    }
103
104    fn insert_unchecked(&mut self, position: (f64, f64), data: Data) {
105        let index_position = self.pos_to_index(position);
106
107        self.cells.get_mut((index_position.0 + (index_position.1 * self.cell_count.0)) as usize).unwrap().push((position, data));
108        self.count += 1;
109    }
110
111    fn count(&self) -> usize {
112        self.count
113    }
114
115    fn clear(&mut self) {
116        self.count = 0;
117        self.cells.iter_mut().for_each(|cell| cell.clear())
118    }
119
120    fn in_circle(&self, position: (f64, f64), radius: f64) -> Vec<Data> {
121        let radius_x = (radius / ((self.x.end - self.x.start) as f64 / self.cell_count.0 as f64)).ceil().min(self.cell_count.0 as f64) as i32;
122        let radius_y = (radius / ((self.y.end - self.y.start) as f64 / self.cell_count.1 as f64)).ceil().min(self.cell_count.1 as f64) as i32;
123
124        let index_position = self.pos_to_index(position);
125
126        let mut data = Vec::new();
127
128        for x in -radius_x..(radius_x + 1) {
129            for y in -radius_y..(radius_y + 1) {
130                let x = index_position.0 as i32 + x;
131                let y = index_position.1 as i32 + y;
132
133                if x < 0 || x >= self.cell_count.0 as i32 || y < 0 || y >= self.cell_count.1 as i32 {
134                    continue;
135                }
136
137                let index = x + (y * self.cell_count.0 as i32);
138
139                match self.cells.get(index as usize) {
140                    None => {}
141                    Some(elements) => {
142                        if elements.len() > 4 {
143                            let pos = self.index_to_pos((x as usize, y as usize));
144
145                            if in_range(position, (pos.0 + self.cell_size.0, pos.1), radius)
146                                && in_range(position, (pos.0 + self.cell_size.0, pos.1 + self.cell_size.1), radius)
147                                && in_range(position, (pos.0, pos.1), radius)
148                                && in_range(position, (pos.0, pos.1 + self.cell_size.1), radius) {
149                                data.extend(elements.iter().map(|x| x.1));
150
151                                continue;
152                            }
153                        }
154
155                        for element in elements {
156                            if in_range(element.0, position, radius) {
157                                data.push(element.1);
158                            }
159                        }
160                    }
161                };
162            }
163        }
164
165        data
166    }
167}