Skip to main content

comnum_math/
nearest.rs

1//! Nearest-point search in a finite set of complex points.
2//!
3//! [`Slicer`] finds the point of a set nearest to a complex value: by rounding the
4//! coordinates when the points form a regular square (QAM) or linear (PAM) grid, in
5//! `O(1)` per value, and by exhaustive search otherwise. It knows nothing of symbols or
6//! bit labels; the constellations of `comnum-mapping` and the decision-directed loops of
7//! `comnum-dsp` build on it.
8//!
9//! # Examples
10//! ```
11//! use comnum_math::{nearest::Slicer, C64};
12//! let s = Slicer::new(&[C64::new(-1.0, 0.0), C64::new(1.0, 0.0)]);
13//! assert_eq!(s.nearest(C64::new(0.3, 2.0)), 1);
14//! assert!((s.distance2(C64::new(0.5, 0.0)) - 0.25).abs() < 1e-12);
15//! ```
16
17use crate::{C64, SCALE_FLOOR};
18
19/// Nearest-point search in a set of points: by rounding the coordinates on a regular square (QAM) or
20/// linear (PAM) grid, `O(1)` per sample, otherwise by exhaustive search.
21#[derive(Clone, Debug, PartialEq)]
22pub struct Slicer {
23    points: Vec<C64>,
24    grid: Option<Grid>,
25}
26
27/// A regular grid of points: a line of levels (PAM) or a square of `side × side` (QAM).
28#[derive(Clone, Debug, PartialEq)]
29struct Grid {
30    /// The levels per axis.
31    side: usize,
32    /// The spacing of the levels.
33    step: f64,
34    /// The first real and imaginary levels.
35    re0: f64,
36    im0: f64,
37    /// Whether the points lie on the real axis (one row).
38    line: bool,
39    /// The point of every cell, row after row (`row · side + column`).
40    lookup: Vec<usize>,
41}
42
43impl Grid {
44    /// The cell of a coordinate `x` along an axis that starts at `x0`, clamped to the grid.
45    fn cell(&self, x: f64, x0: f64) -> usize {
46        ((x - x0) / self.step).round().clamp(0.0, (self.side - 1) as f64) as usize
47    }
48
49    /// The index of the point nearest to `z`.
50    fn nearest(&self, z: C64) -> usize {
51        let c = self.cell(z.re, self.re0);
52        if self.line { self.lookup[c] } else { self.lookup[self.cell(z.im, self.im0) * self.side + c] }
53    }
54}
55
56impl Slicer {
57    /// Detect a regular grid among the points.
58    pub fn new(points: &[C64]) -> Slicer {
59        Slicer { points: points.to_vec(), grid: Slicer::detect(points) }
60    }
61
62    /// A slicer that always searches exhaustively (the reference of the grid search).
63    pub fn exhaustive(points: &[C64]) -> Slicer {
64        Slicer { points: points.to_vec(), grid: None }
65    }
66
67    /// Whether the points form a regular grid (the `O(1)` search is used).
68    pub fn is_grid(&self) -> bool {
69        self.grid.is_some()
70    }
71
72    /// The points.
73    pub fn points(&self) -> &[C64] {
74        &self.points
75    }
76
77    fn levels(v: impl Iterator<Item = f64>) -> Vec<f64> {
78        let mut l: Vec<f64> = v.collect();
79        l.sort_by(f64::total_cmp);
80        let span = (l[l.len() - 1] - l[0]).abs().max(SCALE_FLOOR);
81        l.dedup_by(|a, b| (*a - *b).abs() <= 1e-9 * span);
82        l
83    }
84
85    fn regular(l: &[f64]) -> Option<f64> {
86        if l.len() < 2 {
87            return None;
88        }
89        let step = (l[l.len() - 1] - l[0]) / (l.len() - 1) as f64;
90        l.iter().enumerate().all(|(i, v)| (v - (l[0] + step * i as f64)).abs() <= 1e-9 * step.abs()).then_some(step)
91    }
92
93    fn detect(points: &[C64]) -> Option<Grid> {
94        let m = points.len();
95        if m < 2 {
96            return None;
97        }
98        let re = Slicer::levels(points.iter().map(|p| p.re));
99        let step = Slicer::regular(&re)?;
100        let tol = 1e-6 * step;
101        let col = |x: f64| ((x - re[0]) / step).round() as usize;
102        if points.iter().all(|p| p.im.abs() <= tol) {
103            // a real grid (PAM): one row
104            if re.len() != m {
105                return None;
106            }
107            let mut lookup = vec![0; m];
108            for (i, p) in points.iter().enumerate() {
109                lookup[col(p.re)] = i;
110            }
111            return Some(Grid { side: m, step, re0: re[0], im0: 0.0, line: true, lookup });
112        }
113        let im = Slicer::levels(points.iter().map(|p| p.im));
114        let step_im = Slicer::regular(&im)?;
115        let side = re.len();
116        if im.len() != side || side * side != m || (step_im - step).abs() > tol {
117            return None;
118        }
119        let mut lookup = vec![usize::MAX; m];
120        for (i, p) in points.iter().enumerate() {
121            let (c, r) = (col(p.re), ((p.im - im[0]) / step).round() as usize);
122            if c >= side || r >= side || lookup[r * side + c] != usize::MAX {
123                return None;
124            }
125            lookup[r * side + c] = i;
126        }
127        Some(Grid { side, step, re0: re[0], im0: im[0], line: false, lookup })
128    }
129
130    /// Index of the nearest point.
131    #[inline]
132    pub fn nearest(&self, z: C64) -> usize {
133        if let Some(grid) = &self.grid {
134            grid.nearest(z)
135        } else {
136            let mut best = (f64::INFINITY, 0usize);
137            for (i, p) in self.points.iter().enumerate() {
138                let d = (z - p).norm_sqr();
139                if d < best.0 {
140                    best = (d, i);
141                }
142            }
143            best.1
144        }
145    }
146
147    /// Squared distance to the nearest point.
148    #[inline]
149    pub fn distance2(&self, z: C64) -> f64 {
150        (z - self.points[self.nearest(z)]).norm_sqr()
151    }
152}