molframe-spatial 0.6.4

Deterministic spatial indices and neighbour search for molframe
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

//! Streaming cross-selection queries over a Cartesian cell grid.

use super::{Grid, NEIGHBOUR_OFFSETS, cell_of, members_of_cell, neighbour_cell};
use crate::SpatialError;
use crate::brute::finite;

/// Evaluates a cell-grid query and emits matches as they are found.
///
/// Each finite query atom visits at most 27 cells. Pairwise work inside those
/// cells depends on local spatial density.
pub(super) fn grid_for_each_pair<const UPPER: bool>(
    positions: &[[f32; 3]],
    grid: &Grid,
    query: &[u32],
    cutoff_squared: f32,
    emit: &mut impl FnMut(u32, u32, f32),
) -> Result<(), SpatialError> {
    for &atom in query {
        let Ok(index) = usize::try_from(atom) else {
            return Err(SpatialError::NumericRangeExceeded);
        };
        let Some(position) = positions.get(index).copied() else {
            continue;
        };

        if !finite(position) {
            continue;
        }

        append_grid_pairs::<UPPER>(positions, grid, atom, position, cutoff_squared, emit)?;
    }
    Ok(())
}

/// Searches the 27 neighbouring cells for one query atom.
fn append_grid_pairs<const UPPER: bool>(
    positions: &[[f32; 3]],
    grid: &Grid,
    atom: u32,
    position: [f32; 3],
    cutoff_squared: f32,
    emit: &mut impl FnMut(u32, u32, f32),
) -> Result<(), SpatialError> {
    visit_members(grid, position, |members| {
        super::super::brute_simd::append_pairs::<UPPER>(
            positions,
            atom,
            position,
            members,
            cutoff_squared,
            emit,
        );
    })
}

pub(super) fn grid_for_each_neighbor(
    positions: &[[f32; 3]],
    grid: &Grid,
    position: [f32; 3],
    cutoff_squared: f32,
    emit: &mut impl FnMut(u32, f32),
) -> Result<(), SpatialError> {
    visit_members(grid, position, |members| {
        for &target in members {
            let squared =
                crate::brute::distance_squared(position, positions[target as usize], None);
            if squared <= cutoff_squared {
                emit(target, squared);
            }
        }
    })
}

fn visit_members(
    grid: &Grid,
    position: [f32; 3],
    mut visit: impl FnMut(&[u32]),
) -> Result<(), SpatialError> {
    let centre = cell_of(grid, position)?;
    for delta in NEIGHBOUR_OFFSETS {
        let Some(cell) = neighbour_cell(grid.dims, centre, delta) else {
            continue;
        };
        visit(members_of_cell(grid, cell).ok_or(SpatialError::NumericRangeExceeded)?);
    }
    Ok(())
}