projective_grid/
graph.rs

1use crate::Float;
2use crate::{NeighborDirection, NodeNeighbor};
3use kiddo::{KdTree, SquaredEuclidean};
4use nalgebra::{Point2, Vector2};
5
6/// A spatially close candidate found by KD-tree search.
7pub struct NeighborCandidate<F: Float = f32> {
8    /// Index of the candidate in the input point array.
9    pub index: usize,
10    /// Vector from source point to this candidate.
11    pub offset: Vector2<F>,
12    /// Euclidean distance.
13    pub distance: F,
14}
15
16/// Extension point for pattern-specific neighbor validation.
17///
18/// Implementors decide whether a spatially close point is a valid
19/// grid neighbor, and if so, assign it a direction and quality score.
20pub trait NeighborValidator<F: Float = f32> {
21    /// Per-point data beyond position (e.g., orientation, cluster label).
22    /// Use `()` if no extra data is needed.
23    type PointData;
24
25    /// Validate whether `candidate` is a valid grid neighbor of the point
26    /// at `source_index`. Returns `(direction, score)` where lower score
27    /// is better, or `None` to reject.
28    fn validate(
29        &self,
30        source_index: usize,
31        source_data: &Self::PointData,
32        candidate: &NeighborCandidate<F>,
33        candidate_data: &Self::PointData,
34    ) -> Option<(NeighborDirection, F)>;
35}
36
37/// Parameters for grid graph construction.
38#[derive(Clone, Debug)]
39pub struct GridGraphParams<F: Float = f32> {
40    /// Number of nearest neighbors to query from the KD-tree.
41    pub k_neighbors: usize,
42    /// Maximum distance (pixels) for the KD-tree pre-filter.
43    pub max_distance: F,
44}
45
46impl<F: Float> Default for GridGraphParams<F> {
47    fn default() -> Self {
48        Self {
49            k_neighbors: 8,
50            max_distance: F::max_value().unwrap_or_else(|| F::from_subset(&1e30)),
51        }
52    }
53}
54
55/// A 4-connected grid graph over 2D points.
56///
57/// Each node has at most one neighbor per cardinal direction (Right, Left, Up, Down),
58/// selected as the best-scoring candidate from spatial proximity search.
59pub struct GridGraph<F: Float = f32> {
60    /// Per-node adjacency list. `neighbors[i]` contains up to 4 validated neighbors
61    /// of the point at index `i`.
62    pub neighbors: Vec<Vec<NodeNeighbor<F>>>,
63}
64
65impl<F: Float + kiddo::float::kdtree::Axis> GridGraph<F> {
66    /// Build a grid graph from 2D points using a caller-supplied validator.
67    ///
68    /// - `positions`: 2D point positions for spatial search.
69    /// - `point_data`: per-point data passed to the validator (same length as `positions`).
70    /// - `validator`: determines which spatial neighbors are valid grid neighbors.
71    /// - `params`: controls KD-tree search parameters.
72    pub fn build<V: NeighborValidator<F>>(
73        positions: &[Point2<F>],
74        point_data: &[V::PointData],
75        validator: &V,
76        params: &GridGraphParams<F>,
77    ) -> Self {
78        assert_eq!(
79            positions.len(),
80            point_data.len(),
81            "positions and point_data must have the same length"
82        );
83
84        let coords: Vec<[F; 2]> = positions.iter().map(|p| [p.x, p.y]).collect();
85        let tree: KdTree<F, 2> = (&coords).into();
86        let max_dist_sq = params.max_distance * params.max_distance;
87
88        let mut neighbors = Vec::with_capacity(positions.len());
89
90        for (i, pos) in positions.iter().enumerate() {
91            let query = [pos.x, pos.y];
92            let results = tree.nearest_n::<SquaredEuclidean>(&query, params.k_neighbors);
93
94            let mut candidates = Vec::new();
95
96            for nn in results {
97                let j = nn.item as usize;
98                if j == i {
99                    continue;
100                }
101
102                let dist_sq = nn.distance;
103                if dist_sq > max_dist_sq {
104                    continue;
105                }
106
107                let neighbor_pos = positions[j];
108                let offset = Vector2::new(neighbor_pos.x - pos.x, neighbor_pos.y - pos.y);
109                let distance = dist_sq.sqrt();
110
111                let candidate = NeighborCandidate {
112                    index: j,
113                    offset,
114                    distance,
115                };
116
117                if let Some((direction, score)) =
118                    validator.validate(i, &point_data[i], &candidate, &point_data[j])
119                {
120                    candidates.push(NodeNeighbor {
121                        direction,
122                        index: j,
123                        distance,
124                        score,
125                    });
126                }
127            }
128
129            neighbors.push(select_neighbors(candidates));
130        }
131
132        Self { neighbors }
133    }
134}
135
136/// Keep at most one neighbor per direction, choosing the lowest-score candidate.
137fn select_neighbors<F: Float>(candidates: Vec<NodeNeighbor<F>>) -> Vec<NodeNeighbor<F>> {
138    let mut best: [Option<NodeNeighbor<F>>; 4] = [None, None, None, None];
139
140    for candidate in candidates {
141        let slot = match candidate.direction {
142            NeighborDirection::Right => &mut best[0],
143            NeighborDirection::Left => &mut best[1],
144            NeighborDirection::Up => &mut best[2],
145            NeighborDirection::Down => &mut best[3],
146        };
147
148        let replace = match slot {
149            None => true,
150            Some(current) => {
151                candidate.score < current.score
152                    || (candidate.score == current.score && candidate.distance < current.distance)
153            }
154        };
155
156        if replace {
157            *slot = Some(candidate);
158        }
159    }
160
161    best.into_iter().flatten().collect()
162}
projective_grid/graph.rs

projective_grid/
graph.rs
