Skip to main content

molframe_spatial/queries/
knn_between.rs

1//! Nearest neighbours of a separate query cloud among a target cloud.
2
3use super::knn_within;
4use crate::{KnnOptions, KnnResult, PeriodicBox, SpatialBackend, SpatialError};
5use molframe_core::{AtomSelection, ExecutionContext};
6
7/// Finds the `k` nearest `targets` of every point of `queries`.
8///
9/// Rows follow `queries`; neighbour indices refer to `targets`. Semantics
10/// otherwise match [`knn_within`].
11///
12/// # Errors
13///
14/// Returns an error when the combined point count exceeds a 32-bit index, or
15/// for any reason [`knn_within`] does.
16pub fn knn_between(
17    targets: &[[f32; 3]],
18    queries: &[[f32; 3]],
19    k: usize,
20    options: KnnOptions,
21    backend: SpatialBackend,
22    periodic: Option<&PeriodicBox>,
23    context: &ExecutionContext,
24) -> Result<KnnResult, SpatialError> {
25    let total = targets
26        .len()
27        .checked_add(queries.len())
28        .and_then(|total| u32::try_from(total).ok())
29        .ok_or(SpatialError::AtomOutOfBounds(u32::MAX))?;
30    let split =
31        u32::try_from(targets.len()).map_err(|_| SpatialError::AtomOutOfBounds(u32::MAX))?;
32    let mut joined = Vec::with_capacity(targets.len() + queries.len());
33    joined.extend_from_slice(targets);
34    joined.extend_from_slice(queries);
35    knn_within(
36        &joined,
37        &AtomSelection::range(split..total),
38        &AtomSelection::range(0..split),
39        k,
40        options,
41        backend,
42        periodic,
43        context,
44    )
45}