molframe_spatial/queries/
knn.rs

1//! First-class parallel nearest-neighbour queries.
2
3use super::indices::IndexWorkspace;
4use crate::brute::distance_squared;
5use crate::{
6    KdPeriodicOptions, KdTree, KnnOptions, KnnResult, NO_NEIGHBOR, PeriodicBox, SpatialBackend,
7    SpatialError,
8};
9use molframe_core::{AtomSelection, ExecutionContext};
10
11/// Finds the `k` nearest target atoms of every query atom.
12///
13/// Rows follow the query selection's order. Within a row neighbours increase in
14/// distance, equal distances ordering by atom index, and unfilled slots hold
15/// [`NO_NEIGHBOR`] with an infinite squared distance. Periodic queries report
16/// exact minimum-image distances; an unbounded periodic search looks through
17/// two lattice shells of images, which covers every wrapped target of a cell
18/// that is not extremely skewed, and a `max_distance` widens it exactly.
19///
20/// `BruteForce` scans every target serially and exists as a reference; every
21/// other backend uses the parallel k-d tree, since a cell list has no
22/// radius-free nearest query. Distances are f32; a target whose f64 distance
23/// ties the k-th at the f32 rounding boundary may be ranked either way.
24///
25/// # Errors
26///
27/// Returns an error for an invalid `max_distance`, out-of-range atom indices,
28/// memory admission failure or cancellation.
29pub fn knn_within(
30    positions: &[[f32; 3]],
31    query: &AtomSelection,
32    targets: &AtomSelection,
33    k: usize,
34    options: KnnOptions,
35    backend: SpatialBackend,
36    periodic: Option<&PeriodicBox>,
37    context: &ExecutionContext,
38) -> Result<KnnResult, SpatialError> {
39    let indices = IndexWorkspace::new(query, targets, positions.len(), context)?;
40    if backend == SpatialBackend::BruteForce {
41        return brute_knn(
42            positions,
43            indices.left(),
44            indices.right(),
45            k,
46            options,
47            periodic,
48        );
49    }
50    KdTree::build_in(
51        positions,
52        indices.right(),
53        periodic.copied(),
54        KdPeriodicOptions::default(),
55        context,
56    )?
57    .knn(indices.left(), k, options, context)
58}
59
60fn brute_knn(
61    positions: &[[f32; 3]],
62    query: &[u32],
63    targets: &[u32],
64    k: usize,
65    options: KnnOptions,
66    periodic: Option<&PeriodicBox>,
67) -> Result<KnnResult, SpatialError> {
68    if let Some(distance) = options.max_distance
69        && !(distance.is_finite() && distance >= 0.0)
70    {
71        return Err(SpatialError::InvalidCutoff);
72    }
73    let limit = options.max_distance.map_or(f32::INFINITY, |d| d * d);
74    let mut indices = Vec::with_capacity(query.len() * k);
75    let mut distances = Vec::with_capacity(query.len() * k);
76    let finite = |point: [f32; 3]| point.iter().all(|value| value.is_finite());
77    let mut row: Vec<(f32, u32)> = Vec::new();
78    for &atom in query {
79        let point = positions
80            .get(atom as usize)
81            .copied()
82            .ok_or(SpatialError::AtomOutOfBounds(atom))?;
83        row.clear();
84        if finite(point) {
85            for &target in targets {
86                let other = positions
87                    .get(target as usize)
88                    .copied()
89                    .ok_or(SpatialError::AtomOutOfBounds(target))?;
90                if (options.exclude_self && target == atom) || !finite(other) {
91                    continue;
92                }
93                let squared = distance_squared(point, other, periodic);
94                if squared <= limit {
95                    row.push((squared, target));
96                }
97            }
98        }
99        row.sort_unstable_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
100        for slot in 0..k {
101            let (id, squared) = row
102                .get(slot)
103                .map_or((NO_NEIGHBOR, f32::INFINITY), |&(d, id)| (id, d));
104            indices.push(id);
105            distances.push(squared);
106        }
107    }
108    Ok(KnnResult {
109        k,
110        indices,
111        distances_squared: distances,
112    })
113}
molframe_spatial/queries/knn.rs

molframe_spatial/queries/
knn.rs
