melvin_ob/objective/
bayesian_set.rs1use crate::util::Vec2D;
2use super::BeaconMeas;
3use fixed::types::I32F32;
4use kiddo::{ImmutableKdTree, SquaredEuclidean};
5use num::traits::FloatConst;
6use std::collections::{HashMap, HashSet};
7use std::num::NonZero;
8use crate::fatal;
9
10#[derive(Debug, Clone, serde::Serialize)]
12pub struct SquareSlice {
13 offset: Vec2D<I32F32>,
15 side_length: Vec2D<I32F32>,
17}
18
19impl SquareSlice {
20 const HEX_PACK_SPACING_FACTOR: f32 = 0.93;
22
23 pub fn new(pos: Vec2D<I32F32>, max_dist: Vec2D<I32F32>) -> Self {
32 let offset = (pos - max_dist).wrap_around_map();
33 Self { offset, side_length: max_dist * I32F32::from_num(2) }
34 }
35
36 pub fn map_right_top(&self, p: Vec2D<I32F32>) -> Vec2D<I32F32> {
44 self.offset + self.offset.unwrapped_to_top_right(&p)
45 }
46
47 pub fn intersect(&self, other: &SquareSlice) -> Option<Self> {
55 let corr_offs = self.map_right_top(other.offset);
56 let start_x = self.offset.x().max(corr_offs.x());
57 let start_y = self.offset.y().max(corr_offs.y());
58
59 let end_x =
60 (self.offset.x() + self.side_length.x()).min(corr_offs.x() + other.side_length.x());
61 let end_y =
62 (self.offset.y() + self.side_length.y()).min(corr_offs.y() + other.side_length.y());
63
64 if end_x - start_x > self.side_length.x() || end_y - start_y > self.side_length.y() {
65 println!("I think wrapping should have occurred here");
66 }
67
68 if start_x >= end_x || start_y >= end_y {
70 return None;
71 }
72
73 Some(Self {
75 offset: Vec2D::new(start_x, start_y),
76 side_length: Vec2D::new(end_x - start_x, end_y - start_y),
77 })
78 }
79
80 pub fn get_coord_set(
90 &self,
91 pos: Vec2D<I32F32>,
92 min_dist: I32F32,
93 max_dist: I32F32,
94 ) -> HashSet<Vec2D<i32>> {
95 let mut coord_set = HashSet::new();
96 let min_dist_sq = min_dist * min_dist;
97 let max_dist_sq = max_dist * max_dist;
98
99 let u_pos = self.map_right_top(pos);
100 let x_start = self.offset.x().to_num::<i32>();
101 let y_start = self.offset.y().to_num::<i32>();
102 let x_end = x_start + self.side_length.x().to_num::<i32>();
103 let y_end = y_start + self.side_length.y().to_num::<i32>();
104
105 for x in x_start..x_end {
106 let x_coord = I32F32::from_num(x);
107 let delt_x_sq = (x_coord - u_pos.x()) * (x_coord - u_pos.x());
108
109 if delt_x_sq > max_dist_sq {
110 continue;
111 }
112
113 let max_y_dist = (max_dist_sq - delt_x_sq).sqrt().ceil().to_num::<i32>();
114 let side_len = self.side_length.y().to_num::<i32>() / 2 - max_y_dist;
115 let y_min = y_start.max(y_start + side_len);
116 let y_max = y_end.min(y_end - side_len);
117
118 for y in y_min..=y_max {
119 let y_coord = I32F32::from_num(y);
120 let delt_y_sq = (y_coord - u_pos.y()) * (y_coord - u_pos.y());
121 let dist_sq = delt_x_sq + delt_y_sq;
122 if dist_sq >= min_dist_sq && dist_sq <= max_dist_sq {
123 coord_set.insert(Vec2D::new(x, y).wrap_around_map());
124 }
125 }
126 }
127 coord_set
128 }
129
130 #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
137 fn generate_hex_grid(&self) -> (ImmutableKdTree<f64, 2>, Vec<Vec2D<f64>>) {
138 let r = BayesianSet::MAX_RES_UNCERTAINTY_RAD;
139 let mut centers = Vec::new();
140 let dx = ((3.0 * r / 2.0) * Self::HEX_PACK_SPACING_FACTOR) as i32;
141 let dy = ((r * 3.0f32.sqrt() / 2.0) * Self::HEX_PACK_SPACING_FACTOR) as i32;
142
143 let side_x = self.side_length.x().to_num::<i32>();
144 let side_y = self.side_length.y().to_num::<i32>();
145
146 for y in (0..side_y).step_by(dy as usize) {
147 for x in (0..side_x).step_by(dx as usize) {
148 let x_offset = if (y / dy) % 2 == 1 { dx / 2 } else { 0 };
150 centers.push([f64::from(x + x_offset), f64::from(y)]);
151 }
152 }
153 (
154 ImmutableKdTree::<f64, 2>::new_from_slice(¢ers),
155 centers.iter().map(Vec2D::from).collect(),
156 )
157 }
158}
159
160#[derive(Debug, Clone, serde::Serialize)]
161pub struct BayesianSet {
167 #[serde(skip)]
169 set: HashSet<Vec2D<i32>>,
170 curr_slice: SquareSlice,
172 measurements: Vec<BeaconMeas>,
174}
175
176impl BayesianSet {
177 const K_FAC_MAX: I32F32 = I32F32::lit("0.9");
179 const K_FAC_MIN: I32F32 = I32F32::lit("1.1");
181 pub const STD_DIST_SAFETY: I32F32 = I32F32::lit("5");
183 pub const MAX_DIST: I32F32 = I32F32::lit("2000");
185 const MIN_DIST: I32F32 = I32F32::lit("0");
187 pub const K_ADD: I32F32 = I32F32::lit("225.1");
189 pub const MAX_RES_UNCERTAINTY_RAD: f32 = 75.0;
191 const MAX_ITEMS: NonZero<usize> = unsafe { NonZero::new_unchecked(6) };
193
194 fn get_dists(d_noisy: I32F32) -> (I32F32, I32F32) {
202 let min_dist = ((d_noisy - Self::K_ADD) / Self::K_FAC_MIN) - Self::STD_DIST_SAFETY;
203 let max_dist = ((d_noisy + Self::K_ADD) / Self::K_FAC_MAX) + Self::STD_DIST_SAFETY;
204 (
205 min_dist.max(Self::MIN_DIST).floor(),
206 max_dist.min(Self::MAX_DIST).ceil(),
207 )
208 }
209
210 pub fn new(meas: BeaconMeas) -> Self {
218 let (min_dist, max_dist) = Self::get_dists(I32F32::from_num(meas.rssi()));
219 let side_len = I32F32::from_num(max_dist);
220 let pos = meas.corr_pos();
221 let slice = SquareSlice::new(pos, Vec2D::new(side_len, side_len));
222 let set = slice.get_coord_set(pos, min_dist, max_dist);
223 Self { set, curr_slice: slice, measurements: vec![meas] }
224 }
225
226 pub fn update(&mut self, meas: &BeaconMeas) {
231 let (min_dist, max_dist) = Self::get_dists(I32F32::from_num(meas.rssi()));
232 let pos = meas.corr_pos();
233 let slice = self
234 .curr_slice
235 .intersect(&SquareSlice::new(pos, Vec2D::new(max_dist, max_dist)))
236 .unwrap_or_else(|| fatal!("No possible intersection found!"));
237 let new_set = slice.get_coord_set(pos, min_dist, max_dist);
238 self.set = self.set.intersection(&new_set).copied().collect();
239 self.curr_slice = slice;
240 }
241
242 pub fn is_in_set(&self, pos: Vec2D<i32>) -> bool { self.set.contains(&pos) }
250
251 #[allow(clippy::cast_sign_loss, clippy::cast_precision_loss, clippy::cast_possible_truncation)]
256 pub fn guess_estimate(&self) -> usize {
257 let len = self.set.len();
258 let max_one_guess_area = Self::MAX_RES_UNCERTAINTY_RAD.powi(2) * f32::PI();
259 (len as f32 / max_one_guess_area).ceil() as usize
260 }
261
262 pub fn pack_perfect_circles(&self) -> Vec<Vec2D<I32F32>> {
267 let (h_c_tree, h_c) = self.curr_slice.generate_hex_grid();
268 let assignments = self.assign_points_to_hexes(&h_c_tree, &h_c);
269 let circles = Self::select_minimal_circles(assignments);
270 circles
271 .iter()
272 .map(|circ| (self.curr_slice.offset + Vec2D::from_real(circ)).wrap_around_map().round())
273 .collect()
274 }
275
276 #[allow(clippy::cast_possible_truncation)]
285 fn assign_points_to_hexes(
286 &self,
287 h_c_tree: &ImmutableKdTree<f64, 2>,
288 h_c: &[Vec2D<f64>],
289 ) -> HashMap<Vec2D<i32>, HashSet<Vec2D<i32>>> {
290 let mut assignments: HashMap<Vec2D<i32>, HashSet<Vec2D<i32>>> = HashMap::new();
291
292 for &p in &self.set {
293 let p_fix = Vec2D::from_real(&p);
294 let p_scaled_fix = self.curr_slice.offset.unwrapped_to(&p_fix);
295 let p_search = [
296 p_scaled_fix.x().to_num::<f64>(),
297 p_scaled_fix.y().to_num::<f64>(),
298 ];
299 let n_res = h_c_tree.nearest_n_within::<SquaredEuclidean>(
300 &p_search,
301 75.0,
302 Self::MAX_ITEMS,
303 true,
304 );
305 for n in &n_res {
306 let center_p = h_c[usize::try_from(n.item).unwrap()];
307 if n.distance > f64::from(BayesianSet::MAX_RES_UNCERTAINTY_RAD).powi(2) {
308 println!(
309 "WARNING: Point {} is {} away: TOO FAR, nearest_point: {}",
310 Vec2D::from(&p_search),
311 n.distance,
312 h_c[usize::try_from(n.item).unwrap()],
313 );
314 }
315 let nearest = Vec2D::new(center_p.x().round() as i32, center_p.y().round() as i32);
316 assignments.entry(nearest).or_default().insert(p);
317 }
318 }
319 assignments
320 }
321
322 fn select_minimal_circles(
330 mut assignments: HashMap<Vec2D<i32>, HashSet<Vec2D<i32>>>,
331 ) -> Vec<Vec2D<i32>> {
332 let mut selected_centers = Vec::new();
333 let mut uncovered: HashSet<Vec2D<i32>> = assignments.values().flatten().copied().collect();
334
335 while !uncovered.is_empty() {
336 let best_center = *assignments
337 .iter()
338 .max_by_key(|(_, covered)| covered.len())
339 .map(|(center, _)| center)
340 .unwrap();
341
342 selected_centers.push(best_center);
343 uncovered = uncovered.difference(&assignments[&best_center]).copied().collect();
344 assignments.remove(&best_center);
345 }
346 selected_centers
347 }
348}