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)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn cross_indices_remain_local_even_when_the_right_index_is_smaller() {
71        let first = [[10.0, 0.0, 0.0], [0.0, 0.0, 0.0]];
72        let second = [[0.5, 0.0, 0.0]];
73        let pairs = cross_pairs(
74            &first,
75            &second,
76            1.0,
77            SpatialBackend::BruteForce,
78            &ExecutionContext::default(),
79        )
80        .expect("valid independent coordinate sets");
81        assert_eq!(
82            pairs,
83            vec![CrossPair {
84                first: 1,
85                second: 0,
86                distance_squared: 0.25
87            }]
88        );
89    }
90
91    #[test]
92    fn dense_sets_do_not_emit_within_set_pairs() {
93        let first = [[0.0; 3]; 4];
94        let second = [[0.0; 3]; 3];
95        let pairs = cross_pairs(
96            &first,
97            &second,
98            1.0,
99            SpatialBackend::CellList,
100            &ExecutionContext::default(),
101        )
102        .expect("valid dense coordinate sets");
103        assert_eq!(pairs.len(), first.len() * second.len());
104        assert!(pairs.iter().all(|pair| pair.first < 4 && pair.second < 3));
105    }
106}
molframe_spatial/cross.rs

molframe_spatial/
cross.rs
