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::{PeriodicBox, 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    cross_pairs_periodic(first, second, cutoff, backend, None, context)
34}
35
36/// Cross-set pairs within `cutoff`, measured through the periodic box when one is given.
37///
38/// # Errors
39/// As [`cross_pairs`], plus the periodic search's own refusals.
40pub fn cross_pairs_periodic(
41    first: &[[f32; 3]],
42    second: &[[f32; 3]],
43    cutoff: f32,
44    backend: SpatialBackend,
45    periodic: Option<&PeriodicBox>,
46    context: &ExecutionContext,
47) -> Result<Vec<CrossPair>, SpatialError> {
48    let left_count = u32::try_from(first.len()).map_err(|_| SpatialError::NumericRangeExceeded)?;
49    let count = first
50        .len()
51        .checked_add(second.len())
52        .ok_or(SpatialError::NumericRangeExceeded)?;
53    let total = u32::try_from(count).map_err(|_| SpatialError::NumericRangeExceeded)?;
54    let bytes = count
55        .checked_mul(size_of::<[f32; 3]>())
56        .ok_or(SpatialError::NumericRangeExceeded)?;
57    let _reservation = context.try_reserve(bytes)?;
58    let mut positions = Vec::with_capacity(count);
59    positions.extend_from_slice(first);
60    positions.extend_from_slice(second);
61    let pairs = pairs_within(
62        &positions,
63        &AtomSelection::range(0..left_count),
64        &AtomSelection::range(left_count..total),
65        cutoff,
66        backend,
67        periodic,
68        context,
69    )?;
70    Ok(pairs
71        .into_iter()
72        .map(|pair| CrossPair {
73            first: pair.first,
74            second: pair.second - left_count,
75            distance_squared: pair.distance_squared,
76        })
77        .collect())
78}
79
80#[cfg(test)]
81#[path = "cross_tests.rs"]
82mod tests;