molframe_spatial/queries/
count.rs

1//! Scalar contact counts with bounded, ordered partial reductions.
2
3use super::{PairQuery, indices::IndexWorkspace, search::validate_cutoff, visit};
4use crate::{CellList, SpatialBackend, SpatialError};
5use molframe_core::parallel::{BlockExecutionError, BlockPlan, try_for_each_block_in};
6
7/// Counts unique unordered pairs without storing pairs or all block results.
8///
9/// Cell queries share one charged index across dynamically scheduled blocks.
10/// Fixed 64-cell or 256-query boundaries are independent of worker count. Other
11/// backends use their bounded visitor. Memory scales with the index and the
12/// admitted worker window, independently of the number of matching pairs.
13///
14/// # Errors
15///
16/// Returns input, memory, cancellation or worker errors before returning a count.
17pub fn count_pairs_within(query: &PairQuery<'_>) -> Result<u64, SpatialError> {
18    validate_cutoff(query.cutoff)?;
19    let indices = IndexWorkspace::new(
20        query.left,
21        query.right,
22        query.positions.len(),
23        query.context,
24    )?;
25    let left = indices.left();
26    let right = indices.right();
27    let plan = query.options.plan(left.len(), right.len(), query.cutoff)?;
28    let mut total = 0_u64;
29    if !matches!(
30        plan.backend,
31        SpatialBackend::CellList | SpatialBackend::NeighborList
32    ) {
33        visit::visit_indices(query, left, right, |_| total += 1)?;
34        return Ok(total);
35    }
36    let index = CellList::build_in(
37        query.positions,
38        right,
39        query.cutoff,
40        query.periodic.copied(),
41        query.options.cell_grid,
42        query.context,
43    )?;
44    let same = left == right;
45    let owned_cells = if same { index.owned_cell_count() } else { None };
46    let blocks = owned_cells.map_or_else(
47        || BlockPlan::new(left.len(), if same { 64 } else { 256 }),
48        |count| BlockPlan::new(count, 64),
49    );
50    try_for_each_block_in(
51        blocks,
52        query.context,
53        size_of::<u64>(),
54        |_, range| {
55            let mut count = 0_u64;
56            if owned_cells.is_some() {
57                index.for_each_owned_cell(range, query.cutoff, |_| count += 1)?;
58                return Ok(count);
59            }
60            if same {
61                index.for_each_upper_candidate(&left[range], query.cutoff, |_| count += 1)?;
62                return Ok(count);
63            }
64            index.for_each_candidate(&left[range], query.cutoff, |a, b, distance| {
65                visit::emit_unique_cross_candidate(a, b, distance, left, right, &mut |_| {
66                    count += 1;
67                });
68            })?;
69            Ok(count)
70        },
71        |count| {
72            total = total
73                .checked_add(count)
74                .ok_or(SpatialError::NumericRangeExceeded)?;
75            Ok(())
76        },
77    )
78    .map_err(|error| match error {
79        BlockExecutionError::Memory(error) => SpatialError::Memory(error),
80        BlockExecutionError::Cancelled => SpatialError::Cancelled,
81        BlockExecutionError::Worker(_) => SpatialError::WorkerPanicked,
82        BlockExecutionError::Operation(error) => error,
83    })?;
84    Ok(total)
85}
86
87#[cfg(test)]
88#[path = "count_tests.rs"]
89mod tests;
molframe_spatial/queries/count.rs

molframe_spatial/queries/
count.rs
