molframe_spatial/queries/
within.rs

1//! `within` selection reduction: which query atoms lie near a target set.
2
3use super::visit::for_each_pairs_within_unsorted;
4use crate::{
5    NeighborPair, PairQuery, PeriodicBox, SpatialBackend, SpatialError, SpatialSearchOptions,
6};
7use molframe_core::{ExecutionContext, selection::AtomSelection};
8
9/// Selects query atoms within `cutoff` of any target atom.
10///
11/// Target atoms are included when they also belong to `query`, matching the
12/// conventional meaning of `within` rather than `around`.
13///
14/// # Errors
15///
16/// Returns the same errors as [`crate::pairs_within`].
17pub fn within(
18    positions: &[[f32; 3]],
19    query: &AtomSelection,
20    target: &AtomSelection,
21    cutoff: f32,
22    backend: SpatialBackend,
23    periodic: Option<&PeriodicBox>,
24    context: &ExecutionContext,
25) -> Result<AtomSelection, SpatialError> {
26    within_with_options(
27        positions,
28        query,
29        target,
30        cutoff,
31        SpatialSearchOptions::with_backend(backend),
32        periodic,
33        context,
34    )
35}
36
37/// Selects query atoms under a complete, inspectable planning profile.
38///
39/// # Errors
40///
41/// Returns the same errors as [`crate::pairs_within_with_options`].
42pub fn within_with_options(
43    positions: &[[f32; 3]],
44    query: &AtomSelection,
45    target: &AtomSelection,
46    cutoff: f32,
47    options: SpatialSearchOptions,
48    periodic: Option<&PeriodicBox>,
49    context: &ExecutionContext,
50) -> Result<AtomSelection, SpatialError> {
51    let mut matched = vec![false; positions.len()];
52    let mut reduction_error = None;
53    for_each_pairs_within_unsorted(
54        &PairQuery {
55            positions,
56            left: query,
57            right: target,
58            cutoff,
59            options,
60            periodic,
61            context,
62        },
63        |pair| {
64            if reduction_error.is_none()
65                && let Err(error) = mark_pair_matches(&mut matched, query, target, pair)
66            {
67                reduction_error = Some(error);
68            }
69        },
70    )?;
71    if let Some(error) = reduction_error {
72        return Err(error);
73    }
74
75    collect_matches(query, target, &matched)
76}
77
78/// Marks query endpoints participating in query-target neighbour pairs.
79///
80/// The dense bit vector removes the `O(M log M)` sort/dedup stage previously
81/// required by `within`.
82fn mark_pair_matches(
83    matched: &mut [bool],
84    query: &AtomSelection,
85    target: &AtomSelection,
86    pair: NeighborPair,
87) -> Result<(), SpatialError> {
88    if query.contains(pair.first) && target.contains(pair.second) {
89        mark_atom(matched, pair.first)?;
90    }
91    if query.contains(pair.second) && target.contains(pair.first) {
92        mark_atom(matched, pair.second)?;
93    }
94    Ok(())
95}
96
97/// Collects marked atoms in query-selection order.
98///
99/// Query-target overlap is included without requiring self-pairs.
100fn collect_matches(
101    query: &AtomSelection,
102    target: &AtomSelection,
103    matched: &[bool],
104) -> Result<AtomSelection, SpatialError> {
105    let mut selected = Vec::new();
106
107    for atom in query {
108        let index = usize::try_from(atom).map_err(|_| SpatialError::NumericRangeExceeded)?;
109        let spatial_match = matched
110            .get(index)
111            .copied()
112            .ok_or(SpatialError::AtomOutOfBounds(atom))?;
113
114        if spatial_match || target.contains(atom) {
115            selected.push(atom);
116        }
117    }
118
119    Ok(AtomSelection::from_sorted(selected))
120}
121
122/// Marks one atom in a dense match bitmap.
123///
124/// # Errors
125///
126/// Returns an out-of-bounds error for an invalid pair endpoint.
127fn mark_atom(matched: &mut [bool], atom: u32) -> Result<(), SpatialError> {
128    let index = usize::try_from(atom).map_err(|_| SpatialError::NumericRangeExceeded)?;
129    let Some(slot) = matched.get_mut(index) else {
130        return Err(SpatialError::AtomOutOfBounds(atom));
131    };
132
133    *slot = true;
134    Ok(())
135}
molframe_spatial/queries/within.rs

molframe_spatial/queries/
within.rs
