Skip to main content

molframe_spatial/
cross.rs

1//! Neighbours between coordinate sets with independent index spaces.
2
3use molframe_core::{AtomSelection, ExecutionContext};
4
5use crate::{SpatialBackend, SpatialError, pairs_within};
6
7/// A neighbour relation with indices local to each input coordinate set.
8#[derive(Clone, Copy, Debug, PartialEq)]
9pub struct CrossPair {
10    /// Index in the first coordinate set.
11    pub first: u32,
12    /// Index in the second coordinate set.
13    pub second: u32,
14    /// Squared separation in ångström squared.
15    pub distance_squared: f32,
16}
17
18/// Find cross-set neighbours without producing within-set pairs.
19///
20/// The coordinate workspace respects the shared execution memory budget. Output
21/// order follows the first index and then the second index deterministically.
22///
23/// # Errors
24/// Returns an error for invalid coordinates or cutoff, overflowing indices,
25/// cancellation, or insufficient execution memory.
26pub fn cross_pairs(
27    first: &[[f32; 3]],
28    second: &[[f32; 3]],
29    cutoff: f32,
30    backend: SpatialBackend,
31    context: &ExecutionContext,
32) -> Result<Vec<CrossPair>, SpatialError> {
33    let left_count = u32::try_from(first.len()).map_err(|_| SpatialError::NumericRangeExceeded)?;
34    let count = first
35        .len()
36        .checked_add(second.len())
37        .ok_or(SpatialError::NumericRangeExceeded)?;
38    let total = u32::try_from(count).map_err(|_| SpatialError::NumericRangeExceeded)?;
39    let bytes = count
40        .checked_mul(size_of::<[f32; 3]>())
41        .ok_or(SpatialError::NumericRangeExceeded)?;
42    let _reservation = context.try_reserve(bytes)?;
43    let mut positions = Vec::with_capacity(count);
44    positions.extend_from_slice(first);
45    positions.extend_from_slice(second);
46    let pairs = pairs_within(
47        &positions,
48        &AtomSelection::range(0..left_count),
49        &AtomSelection::range(left_count..total),
50        cutoff,
51        backend,
52        None,
53        context,
54    )?;
55    Ok(pairs
56        .into_iter()
57        .map(|pair| CrossPair {
58            first: pair.first,
59            second: pair.second - left_count,
60            distance_squared: pair.distance_squared,
61        })
62        .collect())
63}
64
65#[cfg(test)]
66#[path = "cross_tests.rs"]
67mod tests;