molframe_spatial/backends/kd/
batch.rs

1//! Parallel batch queries over a k-d tree: nearest neighbours and pair lists.
2
3use super::KdTree;
4use crate::backends::errors::{block_error, validate_cutoff, validate_indices};
5use crate::brute::canonicalise;
6use crate::{NeighborPair, SpatialError};
7use molframe_core::ExecutionContext;
8use molframe_core::parallel::{BlockPlan, try_for_each_block_in};
9
10/// Queries handled per scheduled block.
11const QUERY_BLOCK: usize = 1024;
12/// Index reported for a neighbour slot that no target filled.
13pub const NO_NEIGHBOR: u32 = u32::MAX;
14
15/// Options for a batch nearest-neighbour query.
16#[derive(Clone, Copy, Debug, Default, PartialEq)]
17pub struct KnnOptions {
18    /// Largest Euclidean distance reported; farther targets leave the slot empty.
19    pub max_distance: Option<f32>,
20    /// Skip the query atom itself (and, when periodic, its own images).
21    pub exclude_self: bool,
22}
23
24/// Row-major `(query_count, k)` nearest-neighbour table.
25///
26/// Rows are ordered by increasing distance, ties by atom index. Slots no
27/// target fills hold [`NO_NEIGHBOR`] and an infinite squared distance.
28#[derive(Clone, Debug, PartialEq)]
29pub struct KnnResult {
30    /// Neighbours per query.
31    pub k: usize,
32    /// Target atom indices, `query_count * k` entries.
33    pub indices: Vec<u32>,
34    /// Squared distances, `query_count * k` entries.
35    pub distances_squared: Vec<f32>,
36}
37
38impl KnnOptions {
39    /// Options for a table query: `exclude_self` defaults to true when the
40    /// points query themselves and to false when separate queries are given.
41    #[must_use]
42    pub fn for_queries(
43        max_distance: Option<f32>,
44        exclude_self: Option<bool>,
45        separate_queries: bool,
46    ) -> Self {
47        let exclude_self = match exclude_self {
48            Some(choice) => choice,
49            None => !separate_queries,
50        };
51        Self {
52            max_distance,
53            exclude_self,
54        }
55    }
56}
57
58impl KnnResult {
59    /// Consumes the table into its indices and Euclidean distances.
60    #[must_use]
61    pub fn into_columns(self) -> (Vec<u32>, Vec<f32>) {
62        let mut distances = self.distances_squared;
63        for value in &mut distances {
64            *value = value.sqrt();
65        }
66        (self.indices, distances)
67    }
68
69    /// Number of query rows.
70    #[must_use]
71    pub fn rows(&self) -> usize {
72        match self.indices.len().checked_div(self.k) {
73            Some(rows) => rows,
74            None => 0,
75        }
76    }
77}
78
79impl KdTree<'_> {
80    /// Finds the `k` nearest targets of every query atom in parallel.
81    ///
82    /// Results are identical at any worker count. Equal distances order by
83    /// atom index. Periodic trees report exact minimum-image distances.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error for an invalid `max_distance`, a query atom outside
88    /// the coordinate array, memory admission failure or cancellation.
89    pub fn knn(
90        &self,
91        query: &[u32],
92        k: usize,
93        options: KnnOptions,
94        context: &ExecutionContext,
95    ) -> Result<KnnResult, SpatialError> {
96        if let Some(distance) = options.max_distance {
97            validate_cutoff(distance)?;
98        }
99        validate_indices(query, self.positions.len())?;
100        let slots = query
101            .len()
102            .checked_mul(k)
103            .ok_or(SpatialError::NumericRangeExceeded)?;
104        let bytes = slots
105            .checked_mul(size_of::<u32>() + size_of::<f32>())
106            .ok_or(SpatialError::NumericRangeExceeded)?;
107        let _reservation = context.try_reserve(bytes)?;
108        let mut indices = Vec::new();
109        let mut distances = Vec::new();
110        indices
111            .try_reserve_exact(slots)
112            .map_err(|_| SpatialError::NumericRangeExceeded)?;
113        distances
114            .try_reserve_exact(slots)
115            .map_err(|_| SpatialError::NumericRangeExceeded)?;
116        if k == 0 || query.is_empty() {
117            return Ok(KnnResult {
118                k,
119                indices,
120                distances_squared: distances,
121            });
122        }
123        self.searcher_check(k, options)?;
124
125        let plan = BlockPlan::new(query.len(), QUERY_BLOCK);
126        try_for_each_block_in(
127            plan,
128            context,
129            QUERY_BLOCK * k * (size_of::<u32>() + size_of::<f32>()),
130            |_, range| {
131                let mut searcher = self.knn_searcher(k, options.max_distance)?;
132                let chunk = match query.get(range) {
133                    Some(chunk) => chunk,
134                    None => &[],
135                };
136                let mut ids = Vec::with_capacity(chunk.len() * k);
137                let mut squared = Vec::with_capacity(chunk.len() * k);
138                for &atom in chunk {
139                    searcher.search(self, atom, options.exclude_self, options.max_distance);
140                    let found = searcher.heap.sort_in_place();
141                    let limit = options.max_distance.map_or(f32::INFINITY, |d| d * d);
142                    let mut filled = 0;
143                    for &(distance, id) in found.iter().take(k) {
144                        if distance > limit {
145                            break;
146                        }
147                        ids.push(id);
148                        squared.push(distance);
149                        filled += 1;
150                    }
151                    for _ in filled..k {
152                        ids.push(NO_NEIGHBOR);
153                        squared.push(f32::INFINITY);
154                    }
155                }
156                Ok((ids, squared))
157            },
158            |(ids, squared): (Vec<u32>, Vec<f32>)| {
159                indices.extend_from_slice(&ids);
160                distances.extend_from_slice(&squared);
161                Ok(())
162            },
163        )
164        .map_err(|error| block_error(&error))?;
165        Ok(KnnResult {
166            k,
167            indices,
168            distances_squared: distances,
169        })
170    }
171
172    /// Fails early when the periodic image budget cannot serve the query.
173    fn searcher_check(&self, k: usize, options: KnnOptions) -> Result<(), SpatialError> {
174        self.knn_searcher(k, options.max_distance).map(|_| ())
175    }
176
177    /// Enumerates unique pairs within `cutoff` for the query atoms, in parallel.
178    ///
179    /// With `same_selection`, the tree must index the same atoms as `query`;
180    /// each pair is then produced once, from its smaller atom, and the output
181    /// is globally sorted without a global sort when `query` ascends. Cross
182    /// selections go through [`canonicalise`] so overlapping selections stay
183    /// duplicate-free.
184    pub(crate) fn pairs_in(
185        &self,
186        query: &[u32],
187        cutoff: f32,
188        same_selection: bool,
189        sorted: bool,
190        context: &ExecutionContext,
191    ) -> Result<Vec<NeighborPair>, SpatialError> {
192        validate_cutoff(cutoff)?;
193        validate_indices(query, self.positions.len())?;
194        // Surface image-budget errors before scheduling any work.
195        self.radius_collector(cutoff)?;
196
197        let block = if self.periodic.is_some() { 256 } else { 2048 };
198        let plan = BlockPlan::new(query.len(), block);
199        let expected = 32 * size_of::<NeighborPair>();
200        let mut found: Vec<NeighborPair> = Vec::new();
201        try_for_each_block_in(
202            plan,
203            context,
204            block.saturating_mul(expected),
205            |_, range| {
206                let mut collector = self.radius_collector(cutoff)?;
207                let mut out = Vec::new();
208                let mut hits: Vec<(u32, f32)> = Vec::new();
209                let chunk = match query.get(range) {
210                    Some(chunk) => chunk,
211                    None => &[],
212                };
213                for &atom in chunk {
214                    hits.clear();
215                    collector.collect(self, atom, |target, squared| {
216                        if !same_selection || target > atom {
217                            hits.push((target, squared));
218                        }
219                    })?;
220                    if sorted || same_selection {
221                        hits.sort_unstable_by_key(|hit| hit.0);
222                    }
223                    out.extend(
224                        hits.iter()
225                            .map(|&(target, squared)| NeighborPair::new(atom, target, squared)),
226                    );
227                }
228                Ok(out)
229            },
230            |mut block: Vec<NeighborPair>| {
231                found.append(&mut block);
232                Ok(())
233            },
234        )
235        .map_err(|error| block_error(&error))?;
236
237        let ascending = query.windows(2).all(|pair| pair[0] < pair[1]);
238        if !same_selection || (sorted && !ascending) {
239            canonicalise(&mut found);
240        }
241        Ok(found)
242    }
243}
molframe_spatial/backends/kd/batch.rs

molframe_spatial/backends/kd/
batch.rs
