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 = Matched::new(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/// One bit per atom: which query atoms some pair has already matched.
79///
80/// Packed words rather than `bool`s, so a 100,000-atom structure needs 12.5 kB
81/// instead of 100 kB; plain words rather than a shared bit vector, so marking a
82/// bit costs no copy-on-write check.
83struct Matched {
84    words: Vec<u64>,
85    len: usize,
86}
87
88impl Matched {
89    fn new(len: usize) -> Self {
90        Self {
91            words: vec![0; len.div_ceil(64)],
92            len,
93        }
94    }
95
96    /// Marks one atom.
97    ///
98    /// # Errors
99    ///
100    /// Returns an out-of-bounds error for an invalid pair endpoint.
101    fn mark(&mut self, atom: u32) -> Result<(), SpatialError> {
102        let index = self.index(atom)?;
103        self.words[index / 64] |= 1 << (index % 64);
104        Ok(())
105    }
106
107    fn get(&self, atom: u32) -> Result<bool, SpatialError> {
108        let index = self.index(atom)?;
109        Ok(self.words[index / 64] & (1 << (index % 64)) != 0)
110    }
111
112    fn index(&self, atom: u32) -> Result<usize, SpatialError> {
113        let index = usize::try_from(atom).map_err(|_| SpatialError::NumericRangeExceeded)?;
114        if index >= self.len {
115            return Err(SpatialError::AtomOutOfBounds(atom));
116        }
117        Ok(index)
118    }
119}
120
121/// Marks query endpoints participating in query-target neighbour pairs.
122///
123/// The dense bitmap removes the `O(M log M)` sort/dedup stage previously
124/// required by `within`.
125fn mark_pair_matches(
126    matched: &mut Matched,
127    query: &AtomSelection,
128    target: &AtomSelection,
129    pair: NeighborPair,
130) -> Result<(), SpatialError> {
131    if query.contains(pair.first) && target.contains(pair.second) {
132        matched.mark(pair.first)?;
133    }
134    if query.contains(pair.second) && target.contains(pair.first) {
135        matched.mark(pair.second)?;
136    }
137    Ok(())
138}
139
140/// Collects marked atoms in query-selection order.
141///
142/// Query-target overlap is included without requiring self-pairs.
143fn collect_matches(
144    query: &AtomSelection,
145    target: &AtomSelection,
146    matched: &Matched,
147) -> Result<AtomSelection, SpatialError> {
148    let mut selected = Vec::new();
149
150    for atom in query {
151        if matched.get(atom)? || target.contains(atom) {
152            selected.push(atom);
153        }
154    }
155
156    Ok(AtomSelection::from_sorted(selected))
157}
molframe_spatial/queries/within.rs

molframe_spatial/queries/
within.rs
