Skip to main content

projective_grid/
graph.rs

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