melvin_ob/objective/
bayesian_set.rs

1use 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/// A structure representing a square-shaped slice of a 2D map.
11#[derive(Debug, Clone, serde::Serialize)]
12pub struct SquareSlice {
13    /// The offset of the square slice on the map.
14    offset: Vec2D<I32F32>,
15    /// The side length of the square slice.
16    side_length: Vec2D<I32F32>,
17}
18
19impl SquareSlice {
20    /// The spacing factor used for hexagonal packing within the square slice.
21    const HEX_PACK_SPACING_FACTOR: f32 = 0.93;
22
23    /// Creates a new [`SquareSlice`] given a position and maximum distance.
24    ///
25    /// # Arguments
26    /// * `pos` - The position vector on the map.
27    /// * `max_dist` - The maximum distance vector defining the slice size.
28    ///
29    /// # Returns
30    /// A new [`SquareSlice`] instance.
31    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    /// Maps a given point to the right-top corner of the slice, considering wrapping.
37    ///
38    /// # Arguments
39    /// * `p` - The point to map.
40    ///
41    /// # Returns
42    /// The mapped point as a `Vec2D<I32F32>`.
43    pub fn map_right_top(&self, p: Vec2D<I32F32>) -> Vec2D<I32F32> {
44        self.offset + self.offset.unwrapped_to_top_right(&p)
45    }
46
47    /// Calculates the intersection of the current slice with another slice.
48    ///
49    /// # Arguments
50    /// * `other` - The other [`SquareSlice`] to intersect with.
51    ///
52    /// # Returns
53    /// An `Option` containing the resulting `SquareSlice` if an intersection exists, otherwise `None`.
54    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 there's no overlap, return None
69        if start_x >= end_x || start_y >= end_y {
70            return None;
71        }
72
73        // Create a new square slice representing the intersection
74        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    /// Generates a set of coordinates within the slice that fall within a given distance range.
81    ///
82    /// # Arguments
83    /// * `pos` - The central position for distance measurement.
84    /// * `min_dist` - The minimum distance from the center.
85    /// * `max_dist` - The maximum distance from the center.
86    ///
87    /// # Returns
88    /// A `HashSet` of map coordinates satisfying the distance condition.
89    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    /// Generates a hexagonal grid of points within the square slice.
131    ///
132    /// # Returns
133    /// A tuple containing:
134    /// 1. A `ImmutableKdTree` for efficient spatial querying.
135    /// 2. A vector of hexagonal points represented as `Vec2D<f64>`.
136    #[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                // Offset every other row
149                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(&centers),
155            centers.iter().map(Vec2D::from).collect(),
156        )
157    }
158}
159
160#[derive(Debug, Clone, serde::Serialize)]
161/// Represents a discrete binary Bayesian set used for probabilistic mapping and spatial estimation.
162///
163/// Maintains a collection of coordinates (`set`) within a certain region (`curr_slice`) 
164/// that satisfy constraints derived from beacon measurements. Utilizes these measurements to 
165/// estimate positions and optimize spatial packing of regions.
166pub struct BayesianSet {
167    /// The current set of coordinates that satisfy constraints.
168    #[serde(skip)]
169    set: HashSet<Vec2D<i32>>,
170    /// The square slice of the map currently being evaluated.
171    curr_slice: SquareSlice,
172    /// A collection of beacon measurements contributing to the set's constraints.
173    measurements: Vec<BeaconMeas>,
174}
175
176impl BayesianSet {
177    /// Maximum scale factor for noise calculations.
178    const K_FAC_MAX: I32F32 = I32F32::lit("0.9");
179    /// Minimum scale factor for noise calculations.
180    const K_FAC_MIN: I32F32 = I32F32::lit("1.1");
181    /// Safety buffer for standard deviation in distance calculations.
182    pub const STD_DIST_SAFETY: I32F32 = I32F32::lit("5");
183    /// Maximum allowable distance for measurements.
184    pub const MAX_DIST: I32F32 = I32F32::lit("2000");
185    /// Minimum allowable distance for measurements.
186    const MIN_DIST: I32F32 = I32F32::lit("0");
187    /// Constant noise offset.
188    pub const K_ADD: I32F32 = I32F32::lit("225.1");
189    /// Maximum resolution for uncertainty radius, used in hexagonal packing.
190    pub const MAX_RES_UNCERTAINTY_RAD: f32 = 75.0;
191    /// Maximum number of items to retrieve during a nearest neighbor search.
192    const MAX_ITEMS: NonZero<usize> = unsafe { NonZero::new_unchecked(6) };
193
194    /// Computes the minimum and maximum distances for a given noisy distance value.
195    ///
196    /// # Arguments
197    /// * `d_noisy` - The noisy distance value (`I32F32`).
198    ///
199    /// # Returns
200    /// A tuple containing the minimum and maximum distances.
201    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    /// Creates a new [`BayesianSet`] from an initial beacon measurement.
211    ///
212    /// # Arguments
213    /// * `meas` - The initial beacon measurement.
214    ///
215    /// # Returns
216    /// A new `BayesianSet` instance.
217    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    /// Updates the current Bayesian set based on a new beacon measurement.
227    ///
228    /// # Arguments
229    /// * `meas` - The new beacon measurement to incorporate.
230    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    /// Checks if a given position is part of the current set.
243    ///
244    /// # Arguments
245    /// * `pos` - The position to check.
246    ///
247    /// # Returns
248    /// `true` if the position is in the set, otherwise `false`.
249    pub fn is_in_set(&self, pos: Vec2D<i32>) -> bool { self.set.contains(&pos) }
250
251    /// Estimates the number of 75px guesses required to cover the current coordinate set.
252    ///
253    /// # Returns
254    /// An estimate of regions needed.
255    #[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    /// Packs the set's coordinates into circular regions with minimal overlap.
263    ///
264    /// # Returns
265    /// A `Vec` of circle centers represented as `Vec2D<I32F32>`.
266    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    /// Assigns points in the set to the nearest hexagonal grid center.
277    ///
278    /// # Arguments
279    /// * `h_c_tree` - Immutable k-d tree of hex centers.
280    /// * `h_c` - Hexagonal centers as a vector.
281    ///
282    /// # Returns
283    /// A `HashMap` mapping each hex center to the points that it covers.
284    #[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    /// Selects minimal circle centers to cover the assigned points.
323    ///
324    /// # Arguments
325    /// * `assignments` - A map of hex centers to their assigned points.
326    ///
327    /// # Returns
328    /// A vector of selected hex centers that cover the points.
329    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}