use kiddo::{KdTree, SquaredEuclidean};
use nalgebra::{Point2, Vector2};
use std::collections::{HashMap, HashSet, VecDeque};
pub use crate::square::seed::Seed;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Admit {
Accept,
Reject,
}
#[derive(Clone, Copy, Debug)]
pub struct LabelledNeighbour {
pub idx: usize,
pub at: (i32, i32),
pub position: Point2<f32>,
}
pub trait GrowValidator {
fn is_eligible(&self, idx: usize) -> bool;
fn required_label_at(&self, i: i32, j: i32) -> Option<u8>;
fn label_of(&self, idx: usize) -> Option<u8>;
fn accept_candidate(
&self,
idx: usize,
at: (i32, i32),
prediction: Point2<f32>,
neighbours: &[LabelledNeighbour],
) -> Admit;
fn edge_ok(
&self,
_candidate_idx: usize,
_neighbour_idx: usize,
_at_candidate: (i32, i32),
_at_neighbour: (i32, i32),
) -> bool {
true
}
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug)]
pub struct GrowParams {
pub attach_search_rel: f32,
pub attach_ambiguity_factor: f32,
pub boundary_search_factor: f32,
}
impl Default for GrowParams {
fn default() -> Self {
Self {
attach_search_rel: 0.35,
attach_ambiguity_factor: 1.5,
boundary_search_factor: 2.0,
}
}
}
impl GrowParams {
pub fn new(attach_search_rel: f32, attach_ambiguity_factor: f32) -> Self {
Self {
attach_search_rel,
attach_ambiguity_factor,
..Self::default()
}
}
}
#[derive(Debug, Default)]
pub struct GrowResult {
pub labelled: HashMap<(i32, i32), usize>,
pub by_corner: HashMap<usize, (i32, i32)>,
pub ambiguous: HashSet<(i32, i32)>,
pub holes: HashSet<(i32, i32)>,
pub grid_u: Vector2<f32>,
pub grid_v: Vector2<f32>,
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
level = "info",
skip_all,
fields(num_corners = positions.len(), cell_size = cell_size),
)
)]
pub fn bfs_grow<V: GrowValidator>(
positions: &[Point2<f32>],
seed: Seed,
cell_size: f32,
params: &GrowParams,
validator: &V,
) -> GrowResult {
let grid_u = {
let raw = positions[seed.b] - positions[seed.a];
let n = raw.norm().max(1e-6);
raw / n
};
let grid_v = {
let raw = positions[seed.c] - positions[seed.a];
let n = raw.norm().max(1e-6);
raw / n
};
let mut tree: KdTree<f32, 2> = KdTree::new();
let mut tree_slot_to_corner: Vec<usize> = Vec::new();
for (idx, pos) in positions.iter().enumerate() {
if validator.is_eligible(idx) {
tree.add(&[pos.x, pos.y], tree_slot_to_corner.len() as u64);
tree_slot_to_corner.push(idx);
}
}
let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
let mut by_corner: HashMap<usize, (i32, i32)> = HashMap::new();
let mut ambiguous: HashSet<(i32, i32)> = HashSet::new();
let mut holes: HashSet<(i32, i32)> = HashSet::new();
for (ij, idx) in [
((0, 0), seed.a),
((1, 0), seed.b),
((0, 1), seed.c),
((1, 1), seed.d),
] {
labelled.insert(ij, idx);
by_corner.insert(idx, ij);
}
let mut boundary: VecDeque<(i32, i32)> = VecDeque::new();
let mut seen_boundary: HashSet<(i32, i32)> = HashSet::new();
for ij in labelled.keys().copied().collect::<Vec<_>>() {
enqueue_cardinal_neighbours(ij, &labelled, &mut boundary, &mut seen_boundary);
}
while let Some(pos) = boundary.pop_front() {
if labelled.contains_key(&pos) {
continue;
}
let (decision, _neighbours) = process_boundary_cell(
pos,
positions,
&labelled,
&by_corner,
&tree,
&tree_slot_to_corner,
grid_u,
grid_v,
cell_size,
params,
validator,
);
match decision {
BoundaryDecision::Hole | BoundaryDecision::EdgeRejected => {
holes.insert(pos);
}
BoundaryDecision::Ambiguous => {
ambiguous.insert(pos);
}
BoundaryDecision::Attach(c_idx) => {
labelled.insert(pos, c_idx);
by_corner.insert(c_idx, pos);
enqueue_cardinal_neighbours(pos, &labelled, &mut boundary, &mut seen_boundary);
}
}
}
let (min_i, min_j) = labelled
.keys()
.fold((i32::MAX, i32::MAX), |(a, b), &(i, j)| (a.min(i), b.min(j)));
if min_i != 0 || min_j != 0 {
let rebased: HashMap<(i32, i32), usize> = labelled
.into_iter()
.map(|((i, j), idx)| ((i - min_i, j - min_j), idx))
.collect();
let rebased_by_corner: HashMap<usize, (i32, i32)> =
rebased.iter().map(|(&ij, &idx)| (idx, ij)).collect();
labelled = rebased;
by_corner = rebased_by_corner;
}
let rebase_pos = |(i, j)| (i - min_i, j - min_j);
let ambiguous: HashSet<(i32, i32)> = ambiguous.into_iter().map(rebase_pos).collect();
let holes: HashSet<(i32, i32)> = holes.into_iter().map(rebase_pos).collect();
GrowResult {
labelled,
by_corner,
ambiguous,
holes,
grid_u,
grid_v,
}
}
pub(super) fn enqueue_cardinal_neighbours(
pos: (i32, i32),
labelled: &HashMap<(i32, i32), usize>,
boundary: &mut VecDeque<(i32, i32)>,
seen: &mut HashSet<(i32, i32)>,
) {
for (di, dj) in [(1, 0), (-1, 0), (0, 1), (0, -1)] {
let neigh = (pos.0 + di, pos.1 + dj);
if !labelled.contains_key(&neigh) && seen.insert(neigh) {
boundary.push_back(neigh);
}
}
}
pub(super) fn collect_labelled_neighbours(
pos: (i32, i32),
window_half: i32,
labelled: &HashMap<(i32, i32), usize>,
positions: &[Point2<f32>],
) -> Vec<LabelledNeighbour> {
let mut out = Vec::new();
for dj in -window_half..=window_half {
for di in -window_half..=window_half {
if di == 0 && dj == 0 {
continue;
}
let at = (pos.0 + di, pos.1 + dj);
if let Some(&idx) = labelled.get(&at) {
out.push(LabelledNeighbour {
idx,
at,
position: positions[idx],
});
}
}
}
out
}
pub fn predict_from_neighbours(
target: (i32, i32),
neighbours: &[LabelledNeighbour],
u: Vector2<f32>,
v: Vector2<f32>,
cell_size: f32,
labelled: &HashMap<(i32, i32), usize>,
positions: &[Point2<f32>],
) -> Point2<f32> {
debug_assert!(!neighbours.is_empty());
let global_i_step = u * cell_size;
let global_j_step = v * cell_size;
let mut sum_x = 0.0_f32;
let mut sum_y = 0.0_f32;
let mut sum_w = 0.0_f32;
for n in neighbours {
let di = (target.0 - n.at.0) as f32;
let dj = (target.1 - n.at.1) as f32;
let d2 = di * di + dj * dj;
let w = if d2 > 0.0 { 1.0 / d2 } else { 1.0 };
let i_step = local_step_at(n.at, (1, 0), labelled, positions).unwrap_or(global_i_step);
let j_step = local_step_at(n.at, (0, 1), labelled, positions).unwrap_or(global_j_step);
let off = i_step * di + j_step * dj;
sum_x += w * (n.position.x + off.x);
sum_y += w * (n.position.y + off.y);
sum_w += w;
}
Point2::new(sum_x / sum_w, sum_y / sum_w)
}
pub(super) fn is_extrapolating(target: (i32, i32), neighbours: &[LabelledNeighbour]) -> bool {
let mut has_neg_di = false;
let mut has_pos_di = false;
let mut has_neg_dj = false;
let mut has_pos_dj = false;
for n in neighbours {
let di = target.0 - n.at.0;
let dj = target.1 - n.at.1;
if di > 0 {
has_neg_di = true; } else if di < 0 {
has_pos_di = true;
}
if dj > 0 {
has_neg_dj = true;
} else if dj < 0 {
has_pos_dj = true;
}
}
!(has_neg_di && has_pos_di && has_neg_dj && has_pos_dj)
}
pub(super) fn local_step_at(
at: (i32, i32),
step: (i32, i32),
labelled: &HashMap<(i32, i32), usize>,
positions: &[Point2<f32>],
) -> Option<Vector2<f32>> {
let here = labelled.get(&at).map(|&i| positions[i])?;
let fwd = (at.0 + step.0, at.1 + step.1);
let bwd = (at.0 - step.0, at.1 - step.1);
let fwd_pos = labelled.get(&fwd).map(|&i| positions[i]);
let bwd_pos = labelled.get(&bwd).map(|&i| positions[i]);
match (fwd_pos, bwd_pos) {
(Some(f), Some(b)) => {
let v = (f - b) * 0.5;
Some(v)
}
(Some(f), None) => Some(f - here),
(None, Some(b)) => Some(here - b),
(None, None) => None,
}
}
pub(super) fn collect_candidates<V: GrowValidator>(
tree: &KdTree<f32, 2>,
slot_to_corner: &[usize],
prediction: Point2<f32>,
search_r: f32,
validator: &V,
required_label: Option<u8>,
by_corner: &HashMap<usize, (i32, i32)>,
) -> Vec<(usize, f32)> {
let r2 = search_r * search_r;
let mut out: Vec<(usize, f32)> = Vec::new();
for nn in tree
.within_unsorted::<SquaredEuclidean>(&[prediction.x, prediction.y], r2)
.into_iter()
{
let idx = slot_to_corner[nn.item as usize];
if by_corner.contains_key(&idx) {
continue;
}
if let Some(req) = required_label {
let Some(got) = validator.label_of(idx) else {
continue;
};
if got != req {
continue;
}
}
let d = nn.distance.sqrt();
out.push((idx, d));
}
out.sort_by(|a, b| a.1.total_cmp(&b.1));
out
}
pub(super) enum CandidateChoice {
None,
Ambiguous,
Unique(usize),
}
pub(super) fn choose_unambiguous<V: GrowValidator>(
candidates: &[(usize, f32)],
ambiguity_factor: f32,
prediction: Point2<f32>,
positions: &[Point2<f32>],
validator: &V,
at: (i32, i32),
neighbours: &[LabelledNeighbour],
) -> CandidateChoice {
if candidates.is_empty() {
return CandidateChoice::None;
}
if candidates.len() >= 2 {
let (_, d0) = candidates[0];
let (_, d1) = candidates[1];
if d0 <= f32::EPSILON {
return CandidateChoice::Ambiguous;
}
if d1 / d0 < ambiguity_factor {
return CandidateChoice::Ambiguous;
}
}
for &(idx, _dist) in candidates {
let pos = positions[idx];
let _ = pos; match validator.accept_candidate(idx, at, prediction, neighbours) {
Admit::Accept => return CandidateChoice::Unique(idx),
Admit::Reject => continue,
}
}
CandidateChoice::None
}
pub(super) fn any_cardinal_edge_ok<V: GrowValidator>(
c_idx: usize,
pos: (i32, i32),
labelled: &HashMap<(i32, i32), usize>,
validator: &V,
) -> bool {
let mut found_any = false;
for (di, dj) in [(1, 0), (-1, 0), (0, 1), (0, -1)] {
let neigh = (pos.0 + di, pos.1 + dj);
if let Some(&n_idx) = labelled.get(&neigh) {
found_any = true;
if validator.edge_ok(c_idx, n_idx, pos, neigh) {
return true;
}
}
}
!found_any
}
pub(super) enum BoundaryDecision {
Hole,
Ambiguous,
EdgeRejected,
Attach(usize),
}
#[allow(clippy::too_many_arguments)]
pub(super) fn process_boundary_cell<V: GrowValidator>(
pos: (i32, i32),
positions: &[Point2<f32>],
labelled: &HashMap<(i32, i32), usize>,
by_corner: &HashMap<usize, (i32, i32)>,
tree: &KdTree<f32, 2>,
tree_slot_to_corner: &[usize],
grid_u: Vector2<f32>,
grid_v: Vector2<f32>,
cell_size: f32,
params: &GrowParams,
validator: &V,
) -> (BoundaryDecision, Vec<LabelledNeighbour>) {
let neighbours = collect_labelled_neighbours(pos, 1, labelled, positions);
if neighbours.is_empty() {
return (BoundaryDecision::Hole, neighbours);
}
let prediction = predict_from_neighbours(
pos,
&neighbours,
grid_u,
grid_v,
cell_size,
labelled,
positions,
);
let search_r = params.attach_search_rel * cell_size;
let extrapolating = is_extrapolating(pos, &neighbours);
let local_search_r = if extrapolating {
search_r * params.boundary_search_factor
} else {
search_r
};
let required_label = validator.required_label_at(pos.0, pos.1);
let candidates = collect_candidates(
tree,
tree_slot_to_corner,
prediction,
local_search_r,
validator,
required_label,
by_corner,
);
let choice = choose_unambiguous(
&candidates,
params.attach_ambiguity_factor,
prediction,
positions,
validator,
pos,
&neighbours,
);
let decision = match choice {
CandidateChoice::None => BoundaryDecision::Hole,
CandidateChoice::Ambiguous => BoundaryDecision::Ambiguous,
CandidateChoice::Unique(c_idx) => {
if !any_cardinal_edge_ok(c_idx, pos, labelled, validator) {
BoundaryDecision::EdgeRejected
} else {
BoundaryDecision::Attach(c_idx)
}
}
};
(decision, neighbours)
}
#[cfg(test)]
mod tests {
use super::*;
struct OpenValidator;
impl GrowValidator for OpenValidator {
fn is_eligible(&self, _idx: usize) -> bool {
true
}
fn required_label_at(&self, _i: i32, _j: i32) -> Option<u8> {
None
}
fn label_of(&self, _idx: usize) -> Option<u8> {
None
}
fn accept_candidate(
&self,
_idx: usize,
_at: (i32, i32),
_prediction: Point2<f32>,
_neighbours: &[LabelledNeighbour],
) -> Admit {
Admit::Accept
}
}
#[test]
fn predict_weights_diagonal_less_than_cardinal() {
let s = 10.0_f32;
let u = Vector2::new(1.0, 0.0);
let v = Vector2::new(0.0, 1.0);
let target = (5, 5);
let cardinal = LabelledNeighbour {
idx: 0,
at: (5, 4),
position: Point2::new(50.0, 40.0),
};
let diagonal = LabelledNeighbour {
idx: 1,
at: (3, 3),
position: Point2::new(30.0, 34.0),
};
let positions = vec![cardinal.position, diagonal.position];
let mut labelled = HashMap::new();
labelled.insert(cardinal.at, 0usize);
labelled.insert(diagonal.at, 1usize);
let pred = predict_from_neighbours(
target,
&[cardinal, diagonal],
u,
v,
s,
&labelled,
&positions,
);
let expected_y = (50.0 + 0.125 * 54.0) / 1.125;
assert!(
(pred.x - 50.0).abs() < 1e-4,
"predicted x {} should equal 50",
pred.x
);
assert!(
(pred.y - expected_y).abs() < 1e-4,
"predicted y {} should equal {} (1/d² weighted)",
pred.y,
expected_y
);
let equal_weight_y = (50.0 + 54.0) * 0.5;
assert!(
(pred.y - 50.0) < (equal_weight_y - 50.0),
"weighted bias {} should be smaller than equal-weight bias {}",
pred.y - 50.0,
equal_weight_y - 50.0,
);
}
#[test]
fn predict_with_only_cardinal_recovers_exact_offset() {
let s = 12.0_f32;
let u = Vector2::new(1.0, 0.0);
let v = Vector2::new(0.0, 1.0);
let target = (2, 2);
let neighbour = LabelledNeighbour {
idx: 0,
at: (1, 2),
position: Point2::new(s, 2.0 * s),
};
let positions = vec![neighbour.position];
let mut labelled = HashMap::new();
labelled.insert(neighbour.at, 0usize);
let pred = predict_from_neighbours(target, &[neighbour], u, v, s, &labelled, &positions);
assert!((pred.x - 2.0 * s).abs() < 1e-4);
assert!((pred.y - 2.0 * s).abs() < 1e-4);
}
#[test]
fn predict_uses_local_step_when_neighbour_has_own_neighbours() {
let u = Vector2::new(1.0, 0.0);
let v = Vector2::new(0.0, 1.0);
let global_cell_size = 50.0_f32;
let neighbour = LabelledNeighbour {
idx: 0,
at: (3, 0),
position: Point2::new(300.0, 0.0),
};
let mut positions = vec![neighbour.position];
let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
labelled.insert((3, 0), 0);
positions.push(Point2::new(310.0, 0.0));
labelled.insert((4, 0), 1);
positions.push(Point2::new(320.0, 0.0));
labelled.insert((5, 0), 2);
let pred = predict_from_neighbours(
(2, 0),
&[neighbour],
u,
v,
global_cell_size,
&labelled,
&positions,
);
assert!(
(pred.x - 290.0).abs() < 1e-3,
"expected adaptive prediction at x=290, got {}",
pred.x
);
assert!((pred.y - 0.0).abs() < 1e-3);
}
#[test]
fn predict_falls_back_to_global_when_no_local_steps() {
let u = Vector2::new(1.0, 0.0);
let v = Vector2::new(0.0, 1.0);
let s = 25.0_f32;
let neighbour = LabelledNeighbour {
idx: 0,
at: (4, 4),
position: Point2::new(100.0, 100.0),
};
let positions = vec![neighbour.position];
let mut labelled: HashMap<(i32, i32), usize> = HashMap::new();
labelled.insert((4, 4), 0);
let pred = predict_from_neighbours((5, 4), &[neighbour], u, v, s, &labelled, &positions);
assert!((pred.x - (100.0 + s)).abs() < 1e-3);
assert!((pred.y - 100.0).abs() < 1e-3);
}
#[test]
fn open_validator_grows_clean_grid() {
let s = 20.0_f32;
let rows = 6_i32;
let cols = 6_i32;
let mut positions = Vec::new();
let mut seed_idx = [0usize; 4];
for j in 0..rows {
for i in 0..cols {
let x = i as f32 * s + 50.0;
let y = j as f32 * s + 50.0;
let k = positions.len();
positions.push(Point2::new(x, y));
if (i, j) == (0, 0) {
seed_idx[0] = k;
}
if (i, j) == (1, 0) {
seed_idx[1] = k;
}
if (i, j) == (0, 1) {
seed_idx[2] = k;
}
if (i, j) == (1, 1) {
seed_idx[3] = k;
}
}
}
let seed = Seed {
a: seed_idx[0],
b: seed_idx[1],
c: seed_idx[2],
d: seed_idx[3],
};
let res = bfs_grow(&positions, seed, s, &GrowParams::default(), &OpenValidator);
assert_eq!(res.labelled.len(), (rows * cols) as usize);
let (mi, mj) = res
.labelled
.keys()
.fold((i32::MAX, i32::MAX), |(a, b), &(i, j)| (a.min(i), b.min(j)));
assert_eq!((mi, mj), (0, 0));
}
}