use std::collections::{HashMap, HashSet};
use nalgebra::{Point2, Vector2};
use crate::shared::extension::{extend_via_local_homography, LocalExtensionParams};
use crate::shared::fill::{fill_grid_holes, FillParams};
use crate::shared::grow::{GrowParams, GrowResult, SquareAttachPolicy};
use crate::shared::grow_extend::extend_from_labelled;
use crate::shared::validate::{self as pg_validate, ValidationParams};
#[non_exhaustive]
#[derive(Clone, Debug)]
pub struct RecoveryParams {
pub local_extension: LocalExtensionParams,
pub bfs_extension: GrowParams,
pub fill: FillParams,
pub max_sweeps: u32,
pub apply_wrong_label_drops: bool,
pub apply_largest_component: bool,
}
impl Default for RecoveryParams {
fn default() -> Self {
Self {
local_extension: LocalExtensionParams::default(),
bfs_extension: GrowParams::default(),
fill: FillParams::default(),
max_sweeps: 4,
apply_wrong_label_drops: true,
apply_largest_component: true,
}
}
}
#[non_exhaustive]
#[derive(Clone, Debug, Default)]
pub enum RecoverySchedule {
#[default]
Auto,
Off,
On(RecoveryParams),
}
impl RecoverySchedule {
pub(crate) fn resolve(&self, synthesized_axes: bool) -> Option<RecoveryParams> {
match self {
RecoverySchedule::Auto if synthesized_axes => Some(RecoveryParams::default()),
RecoverySchedule::Auto => None,
RecoverySchedule::Off => None,
RecoverySchedule::On(p) => Some(p.clone()),
}
}
}
#[derive(Clone, Debug, Default)]
pub struct RecoveryStats {
pub sweeps: u32,
pub net_added: i64,
pub attached: usize,
pub dropped: usize,
}
pub fn run_schedule<V: SquareAttachPolicy>(
positions: &[Point2<f32>],
grow: &mut GrowResult,
cell_size: f32,
policy: &V,
params: &RecoveryParams,
validate_params: &ValidationParams,
) -> RecoveryStats {
let mut stats = RecoveryStats::default();
if grow.labelled.len() < 4 {
return stats;
}
ensure_axes(grow, positions);
let start = grow.labelled.len() as i64;
for _sweep in 0..params.max_sweeps.max(1) {
stats.sweeps += 1;
let before = grow.labelled.len();
let local = extend_via_local_homography(
positions,
grow,
cell_size,
¶ms.local_extension,
policy,
);
stats.attached += local.attached;
let bfs = extend_from_labelled(positions, grow, cell_size, ¶ms.bfs_extension, policy);
stats.attached += bfs.attached;
let fill = fill_grid_holes(positions, grow, cell_size, ¶ms.fill, policy);
stats.attached += fill.added;
let dropped = revalidate_and_filter(positions, grow, cell_size, params, validate_params);
stats.dropped += dropped;
if grow.labelled.len() == before {
break;
}
}
stats.net_added = grow.labelled.len() as i64 - start;
stats
}
fn revalidate_and_filter(
positions: &[Point2<f32>],
grow: &mut GrowResult,
cell_size: f32,
params: &RecoveryParams,
validate_params: &ValidationParams,
) -> usize {
let result = pg_validate::wrong_label_filters::drop_set(
&grow.labelled,
|idx| positions[idx],
cell_size,
validate_params,
params.apply_wrong_label_drops,
params.apply_largest_component,
);
if result.drop.is_empty() {
return 0;
}
let mut removed = 0usize;
grow.labelled.retain(|_, &mut idx| {
if result.drop.contains(&idx) {
removed += 1;
false
} else {
true
}
});
grow.by_corner.retain(|idx, _| !result.drop.contains(idx));
removed
}
fn ensure_axes(grow: &mut GrowResult, positions: &[Point2<f32>]) {
let needs = grow.axis_i.norm() < 1e-3 || grow.axis_j.norm() < 1e-3;
if !needs {
return;
}
let (mut sum_i, mut n_i) = (Vector2::<f32>::zeros(), 0u32);
let (mut sum_j, mut n_j) = (Vector2::<f32>::zeros(), 0u32);
for (&(i, j), &idx) in &grow.labelled {
let here = positions[idx];
if let Some(&n) = grow.labelled.get(&(i + 1, j)) {
sum_i += positions[n] - here;
n_i += 1;
}
if let Some(&n) = grow.labelled.get(&(i, j + 1)) {
sum_j += positions[n] - here;
n_j += 1;
}
}
if n_i > 0 {
let v = sum_i / n_i as f32;
if v.norm() > 1e-3 {
grow.axis_i = v.normalize();
}
}
if n_j > 0 {
let v = sum_j / n_j as f32;
if v.norm() > 1e-3 {
grow.axis_j = v.normalize();
}
}
if grow.axis_i.norm() < 1e-3 {
grow.axis_i = Vector2::new(1.0, 0.0);
}
if grow.axis_j.norm() < 1e-3 {
grow.axis_j = Vector2::new(0.0, 1.0);
}
}
#[derive(Clone, Copy)]
pub(crate) struct RecoveryInputs<'a> {
pub features: &'a [crate::feature::OrientedFeature<2>],
pub positions: &'a [Point2<f32>],
pub local_pitch: &'a [f32],
pub params: &'a RecoveryParams,
pub validate_params: &'a ValidationParams,
}
pub(crate) fn recover_components(
merged: Vec<HashMap<(i32, i32), usize>>,
inputs: RecoveryInputs<'_>,
) -> Vec<HashMap<(i32, i32), usize>> {
let RecoveryInputs { positions, .. } = inputs;
let mut order: Vec<usize> = (0..merged.len()).collect();
order.sort_by(|&a, &b| {
merged[b]
.len()
.cmp(&merged[a].len())
.then_with(|| min_index(&merged[a]).cmp(&min_index(&merged[b])))
});
let mut claimed: HashSet<usize> = HashSet::new();
let mut recovered_by_slot: Vec<HashMap<(i32, i32), usize>> = vec![HashMap::new(); merged.len()];
for &k in &order {
let comp: HashMap<(i32, i32), usize> = merged[k]
.iter()
.filter(|(_, &idx)| !claimed.contains(&idx))
.map(|(&k, &v)| (k, v))
.collect();
if comp.len() < 4 {
for &idx in comp.values() {
claimed.insert(idx);
}
recovered_by_slot[k] = if comp.is_empty() {
HashMap::new()
} else {
rebase_to_origin(&comp)
};
continue;
}
let cell_size = cell_size_of(&comp, positions);
let own: HashSet<usize> = comp.values().copied().collect();
let masked: HashSet<usize> = claimed.difference(&own).copied().collect();
let recovered = recover_positions_component(&comp, &masked, cell_size, inputs);
claimed.extend(recovered.values().copied());
recovered_by_slot[k] = rebase_to_origin(&recovered);
}
recovered_by_slot.retain(|m| !m.is_empty());
recovered_by_slot
}
fn min_index(labelled: &HashMap<(i32, i32), usize>) -> usize {
labelled.values().copied().min().unwrap_or(usize::MAX)
}
const LOCAL_PITCH_NEIGHBOURS: usize = 5;
pub(crate) fn local_pitch_of(positions: &[Point2<f32>]) -> Vec<f32> {
use kiddo::{KdTree, SquaredEuclidean};
let n = positions.len();
if n < 2 {
return vec![0.0; n];
}
let mut tree: KdTree<f32, 2> = KdTree::new();
for (i, p) in positions.iter().enumerate() {
tree.add(&[p.x, p.y], i as u64);
}
positions
.iter()
.enumerate()
.map(|(i, p)| {
let hits = tree.nearest_n::<SquaredEuclidean>(&[p.x, p.y], LOCAL_PITCH_NEIGHBOURS + 1);
let mut dists: Vec<f32> = hits
.into_iter()
.filter(|nn| nn.item as usize != i)
.map(|nn| nn.distance.sqrt())
.filter(|d| d.is_finite() && *d > 1e-3)
.collect();
if dists.is_empty() {
return 0.0;
}
dists.sort_by(|a, b| a.total_cmp(b));
dists[dists.len() / 2]
})
.collect()
}
fn cell_size_of(labelled: &HashMap<(i32, i32), usize>, positions: &[Point2<f32>]) -> f32 {
let mut sum = 0.0_f32;
let mut count = 0usize;
for (&(i, j), &idx) in labelled {
let here = positions[idx];
for (di, dj) in [(1, 0), (0, 1), (-1, 0), (0, -1)] {
if let Some(&n) = labelled.get(&(i + di, j + dj)) {
sum += (positions[n] - here).norm();
count += 1;
}
}
}
if count == 0 {
1.0
} else {
sum / count as f32
}
}
fn rebase_to_origin(labelled: &HashMap<(i32, i32), usize>) -> HashMap<(i32, i32), usize> {
let min_i = labelled.keys().map(|&(i, _)| i).min().unwrap_or(0);
let min_j = labelled.keys().map(|&(_, j)| j).min().unwrap_or(0);
if min_i == 0 && min_j == 0 {
return labelled.clone();
}
labelled
.iter()
.map(|(&(i, j), &idx)| ((i - min_i, j - min_j), idx))
.collect()
}
pub(crate) fn recover_positions_component(
labelled: &HashMap<(i32, i32), usize>,
masked: &HashSet<usize>,
cell_size: f32,
inputs: RecoveryInputs<'_>,
) -> HashMap<(i32, i32), usize> {
use crate::shared::positions_policy::{PositionsAttachPolicy, PositionsTolerances};
let tol = PositionsTolerances {
soft_axis_tol_rad: 0.872_664_6,
edge_length_tol: 0.40,
cell_size,
};
let inner =
PositionsAttachPolicy::new(inputs.features, inputs.positions, inputs.local_pitch, tol);
let policy = MaskedPolicy {
inner: &inner,
masked,
};
let mut grow = grow_result_from_labelled(labelled, inputs.positions);
run_schedule(
inputs.positions,
&mut grow,
cell_size,
&policy,
inputs.params,
inputs.validate_params,
);
grow.labelled
}
struct MaskedPolicy<'a, V: SquareAttachPolicy> {
inner: &'a V,
masked: &'a HashSet<usize>,
}
impl<V: SquareAttachPolicy> SquareAttachPolicy for MaskedPolicy<'_, V> {
fn is_eligible(&self, idx: usize) -> bool {
!self.masked.contains(&idx) && self.inner.is_eligible(idx)
}
fn required_label_at(&self, i: i32, j: i32) -> Option<u8> {
self.inner.required_label_at(i, j)
}
fn label_of(&self, idx: usize) -> Option<u8> {
self.inner.label_of(idx)
}
fn accept_candidate(
&self,
idx: usize,
at: (i32, i32),
prediction: Point2<f32>,
neighbours: &[crate::shared::grow::LabelledNeighbour],
) -> crate::shared::grow::Admit {
self.inner.accept_candidate(idx, at, prediction, neighbours)
}
fn edge_ok(&self, c: usize, n: usize, ac: (i32, i32), an: (i32, i32)) -> bool {
self.inner.edge_ok(c, n, ac, an)
}
}
pub(crate) fn grow_result_from_labelled(
labelled: &HashMap<(i32, i32), usize>,
positions: &[Point2<f32>],
) -> GrowResult {
let by_corner: HashMap<usize, (i32, i32)> = labelled.iter().map(|(&k, &v)| (v, k)).collect();
let mut grow = GrowResult {
labelled: labelled.clone(),
by_corner,
..Default::default()
};
ensure_axes(&mut grow, positions);
grow
}
#[cfg(test)]
mod tests {
use super::*;
use crate::shared::grow::{Admit, LabelledNeighbour};
struct OpenPolicy<'a> {
positions: &'a [Point2<f32>],
cell_size: f32,
}
impl SquareAttachPolicy for OpenPolicy<'_> {
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
}
fn edge_ok(&self, c: usize, n: usize, _ac: (i32, i32), _an: (i32, i32)) -> bool {
let len = (self.positions[c] - self.positions[n]).norm();
let r = len / self.cell_size;
(0.6..=1.4).contains(&r)
}
}
type SyntheticGrid = (Vec<Point2<f32>>, HashMap<(i32, i32), usize>);
fn grid(rows: i32, cols: i32, s: f32) -> SyntheticGrid {
let mut pos = Vec::new();
let mut map = HashMap::new();
let mut idx = 0usize;
for j in 0..rows {
for i in 0..cols {
pos.push(Point2::new(i as f32 * s + 40.0, j as f32 * s + 40.0));
map.insert((i, j), idx);
idx += 1;
}
}
(pos, map)
}
#[test]
fn fills_interior_holes_and_extends_boundary() {
let s = 30.0_f32;
let (pos, full) = grid(7, 7, s);
let mut seed: HashMap<(i32, i32), usize> = HashMap::new();
for j in 2..5 {
for i in 2..5 {
seed.insert((i, j), full[&(i, j)]);
}
}
let mut grow = grow_result_from_labelled(&seed, &pos);
let policy = OpenPolicy {
positions: &pos,
cell_size: s,
};
let params = RecoveryParams::default();
let vp = ValidationParams::default();
let stats = run_schedule(&pos, &mut grow, s, &policy, ¶ms, &vp);
assert!(
grow.labelled.len() >= 45,
"recovered only {}/49 (sweeps {})",
grow.labelled.len(),
stats.sweeps
);
for (&cell, &idx) in &grow.labelled {
assert_eq!(
full.get(&cell),
Some(&idx),
"cell {cell:?} mislabelled to index {idx}"
);
}
}
#[test]
fn decoys_off_lattice_are_never_labelled() {
let s = 30.0_f32;
let (mut pos, full) = grid(6, 6, s);
let grid_n = pos.len();
let decoys = [
Point2::new(40.0 + 0.5 * s, 40.0 + 0.5 * s), Point2::new(40.0 + 2.5 * s, 40.0 + 1.5 * s),
Point2::new(40.0 - 3.0 * s, 40.0 + 2.0 * s), Point2::new(40.0 + 9.0 * s, 40.0 + 9.0 * s),
];
for d in decoys {
pos.push(d);
}
let mut seed: HashMap<(i32, i32), usize> = HashMap::new();
for j in 1..4 {
for i in 1..4 {
seed.insert((i, j), full[&(i, j)]);
}
}
let mut grow = grow_result_from_labelled(&seed, &pos);
let policy = OpenPolicy {
positions: &pos,
cell_size: s,
};
let stats = run_schedule(
&pos,
&mut grow,
s,
&policy,
&RecoveryParams::default(),
&ValidationParams::default(),
);
for &(_, idx) in grow
.by_corner
.iter()
.map(|(idx, c)| (c, idx))
.collect::<Vec<_>>()
.iter()
{
assert!(
*idx < grid_n,
"a decoy (index {idx} ≥ {grid_n}) was labelled (sweeps {})",
stats.sweeps
);
}
for (&cell, &idx) in &grow.labelled {
assert_eq!(full.get(&cell), Some(&idx), "cell {cell:?} mislabelled");
}
}
#[test]
fn idempotent_on_clean_full_grid() {
let s = 30.0_f32;
let (pos, full) = grid(5, 5, s);
let mut grow = grow_result_from_labelled(&full, &pos);
let policy = OpenPolicy {
positions: &pos,
cell_size: s,
};
let before = grow.labelled.len();
let stats = run_schedule(
&pos,
&mut grow,
s,
&policy,
&RecoveryParams::default(),
&ValidationParams::default(),
);
assert_eq!(grow.labelled.len(), before, "schedule altered a full grid");
assert_eq!(stats.net_added, 0);
}
}