use crate::Float;
use kiddo::{KdTree, SquaredEuclidean};
use nalgebra::{Point2, RealField};
#[derive(Clone, Copy, Debug)]
pub struct GlobalStepEstimate<F: Float = f32> {
pub cell_size: F,
pub support: u32,
pub sample_count: u32,
pub confidence: F,
}
#[derive(Clone, Copy, Debug)]
pub struct GlobalStepParams<F: Float = f32> {
pub bandwidth_rel: F,
pub max_iters: u32,
pub convergence_rel: F,
}
impl<F: Float> Default for GlobalStepParams<F> {
fn default() -> Self {
Self {
bandwidth_rel: F::from_subset(&0.15),
max_iters: 20,
convergence_rel: F::from_subset(&1e-3),
}
}
}
pub fn estimate_global_cell_size<F: Float + kiddo::float::kdtree::Axis>(
positions: &[Point2<F>],
params: &GlobalStepParams<F>,
) -> Option<GlobalStepEstimate<F>> {
if positions.len() < 2 {
return None;
}
let coords: Vec<[F; 2]> = positions.iter().map(|p| [p.x, p.y]).collect();
let tree: KdTree<F, 2> = (&coords).into();
let mut nn_distances: Vec<F> = Vec::with_capacity(positions.len());
for (i, p) in positions.iter().enumerate() {
let hits = tree.nearest_n::<SquaredEuclidean>(&[p.x, p.y], 2);
for hit in hits {
let j = hit.item as usize;
if j == i {
continue;
}
let d2 = hit.distance;
if d2 > F::zero() {
nn_distances.push(d2.sqrt());
}
break;
}
}
if nn_distances.is_empty() {
return None;
}
nn_distances.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let sample_count = nn_distances.len() as u32;
let seeds = [
percentile_sorted(&nn_distances, F::from_subset(&0.25)),
percentile_sorted(&nn_distances, F::from_subset(&0.5)),
percentile_sorted(&nn_distances, F::from_subset(&0.75)),
];
let mut best: Option<(F, u32, F)> = None; for seed in seeds {
if let Some((mode, support)) = mean_shift_mode(&nn_distances, seed, params) {
if support == 0 {
continue;
}
let score = F::from_subset(&(support as f64)) * mode;
if best.map(|b: (F, u32, F)| score > b.2).unwrap_or(true) {
best = Some((mode, support, score));
}
}
}
let (cell_size, support, _) = best?;
let confidence = RealField::max(
RealField::min(
F::from_subset(&(support as f64)) / F::from_subset(&(sample_count as f64)),
F::one(),
),
F::zero(),
);
Some(GlobalStepEstimate {
cell_size,
support,
sample_count,
confidence,
})
}
fn percentile_sorted<F: Float>(sorted: &[F], q: F) -> F {
let len = sorted.len();
if len == 0 {
return F::zero();
}
let idx_f = q * F::from_subset(&((len - 1) as f64));
let idx = idx_f.floor();
let i = idx.to_subset().unwrap_or(0.0) as usize;
let i = i.min(len - 1);
sorted[i]
}
fn mean_shift_mode<F: Float>(
sorted: &[F],
seed: F,
params: &GlobalStepParams<F>,
) -> Option<(F, u32)> {
if seed <= F::zero() {
return None;
}
let bandwidth = seed * params.bandwidth_rel;
if bandwidth <= F::zero() {
return Some((seed, 0));
}
let convergence = bandwidth * params.convergence_rel;
let mut center = seed;
for _ in 0..params.max_iters {
let mut sum = F::zero();
let mut weight = F::zero();
let mut count_in_band = 0u32;
for &v in sorted {
let diff = v - center;
if diff.abs() > bandwidth {
continue;
}
let t = diff / bandwidth;
let w = F::one() - t * t;
let w = if w < F::zero() { F::zero() } else { w };
if w > F::zero() {
sum += v * w;
weight += w;
count_in_band += 1;
}
}
if weight <= F::zero() {
return Some((center, 0));
}
let next = sum / weight;
if (next - center).abs() <= convergence {
return Some((next, count_in_band));
}
center = next;
}
let mut in_band = 0u32;
for &v in sorted {
if (v - center).abs() <= bandwidth {
in_band += 1;
}
}
Some((center, in_band))
}
#[cfg(test)]
mod tests {
use super::*;
fn rectangular_grid(rows: u32, cols: u32, spacing: f32) -> Vec<Point2<f32>> {
let mut out = Vec::new();
for j in 0..rows {
for i in 0..cols {
out.push(Point2::new(i as f32 * spacing, j as f32 * spacing));
}
}
out
}
#[test]
fn recovers_regular_grid_spacing() {
let params = GlobalStepParams::<f32>::default();
for &spacing in &[10.0_f32, 24.0, 50.0] {
let pts = rectangular_grid(5, 5, spacing);
let est = estimate_global_cell_size(&pts, ¶ms).expect("estimate");
assert!(
(est.cell_size - spacing).abs() / spacing < 0.02,
"spacing {spacing}: estimate {} off >2 %",
est.cell_size
);
assert!(est.confidence > 0.9, "confidence {}", est.confidence);
}
}
#[test]
fn sparse_noise_does_not_drag_mode() {
let mut pts = rectangular_grid(5, 5, 24.0);
for (dx, dy) in [(6.0, 9.0), (43.0, 9.0), (9.0, 43.0), (81.0, 81.0)] {
pts.push(Point2::new(dx, dy));
}
let est =
estimate_global_cell_size(&pts, &GlobalStepParams::<f32>::default()).expect("estimate");
assert!(
(est.cell_size - 24.0).abs() < 2.0,
"expected board step ~24 but got {}",
est.cell_size
);
assert!(est.support >= 10); }
#[test]
fn bimodal_density_weights_by_cell_size() {
let mut pts = Vec::new();
for j in 0..4 {
for i in 0..4 {
pts.push(Point2::new(i as f32 * 4.0, j as f32 * 4.0));
}
}
for j in 0..4 {
for i in 0..4 {
pts.push(Point2::new(
1000.0 + i as f32 * 40.0,
1000.0 + j as f32 * 40.0,
));
}
}
let est =
estimate_global_cell_size(&pts, &GlobalStepParams::<f32>::default()).expect("estimate");
assert!(
(est.cell_size - 40.0).abs() < 4.0,
"expected larger-grid cell ~40 but got {}",
est.cell_size
);
}
#[test]
fn too_small_input_returns_none() {
let pts: Vec<Point2<f32>> = vec![];
assert!(estimate_global_cell_size(&pts, &GlobalStepParams::<f32>::default()).is_none());
let pts = vec![Point2::new(0.0, 0.0)];
assert!(estimate_global_cell_size(&pts, &GlobalStepParams::<f32>::default()).is_none());
}
#[test]
fn degenerate_duplicate_points_are_skipped() {
let pts = vec![
Point2::new(0.0, 0.0),
Point2::new(0.0, 0.0),
Point2::new(10.0, 0.0),
Point2::new(0.0, 10.0),
Point2::new(10.0, 10.0),
];
let est =
estimate_global_cell_size(&pts, &GlobalStepParams::<f32>::default()).expect("estimate");
assert!((est.cell_size - 10.0).abs() < 1.0);
}
#[test]
fn mild_jitter_still_recovers_mode() {
let pts: Vec<Point2<f32>> = rectangular_grid(5, 5, 24.0)
.into_iter()
.enumerate()
.map(|(i, p)| {
let jitter_x = ((i * 17 % 7) as f32 - 3.0) * 0.4;
let jitter_y = ((i * 23 % 9) as f32 - 4.0) * 0.4;
Point2::new(p.x + jitter_x, p.y + jitter_y)
})
.collect();
let est =
estimate_global_cell_size(&pts, &GlobalStepParams::<f32>::default()).expect("estimate");
assert!(
(est.cell_size - 24.0).abs() < 2.0,
"expected ~24 got {}",
est.cell_size
);
}
}