molframe_spatial/backends/
kd.rs

1//! Implicit, bucketed, cache-friendly three-dimensional k-d tree.
2//!
3//! # Layout
4//!
5//! The tree is a perfect binary tree over *index medians*: a node covering the
6//! slot range `lo..hi` splits at `lo + (hi - lo) / 2` on the axis of widest
7//! extent, after `select_nth_unstable` has placed the median there. Both the
8//! shape and every child range therefore follow from the point count alone, so
9//! no child pointers are stored. Per node the tree keeps one split value and
10//! one axis (heap order, children `2i + 1` and `2i + 2`); per slot it keeps
11//! structure-of-arrays coordinates and the atom index. Leaves are buckets of
12//! `LEAF_MAX / 2 ..= LEAF_MAX` points scanned linearly, which the compiler
13//! vectorises. Splitting by count rather than by value keeps the depth at
14//! `ceil(log2(n / LEAF_MAX))` for duplicates, collinear layouts and outliers.
15//!
16//! Left subtrees hold coordinates `<=` the split value and right subtrees
17//! `>=` it, so equal coordinates may sit on either side; searches treat a
18//! plane at distance zero as intersecting both.
19//!
20//! # Concurrency
21//!
22//! Construction splits the top levels serially and the independent subtrees on
23//! the context's workers. Queries are independent and run in blocks on the
24//! same workers; no query allocates (scratch is owned per block).
25
26use super::errors::{validate_cutoff, validate_indices};
27use crate::brute::finite;
28use crate::{KdPeriodicOptions, NeighborPair, PeriodicBox, SpatialError};
29use molframe_core::ExecutionContext;
30use molframe_core::parallel::{BlockPlan, map_blocks_in};
31use std::sync::Mutex;
32
33mod batch;
34mod heap;
35mod periodic;
36mod search;
37mod stream;
38
39pub use batch::{KnnOptions, KnnResult, NO_NEIGHBOR};
40use periodic::{for_each_shift, periodic_image_limits, validate_image_budget};
41
42/// Largest leaf bucket. Leaves hold `LEAF_MAX / 2 ..= LEAF_MAX` points.
43pub(super) const LEAF_MAX: usize = 32;
44
45#[derive(Clone, Copy, Debug)]
46struct Entry {
47    position: [f32; 3],
48    atom: u32,
49}
50
51/// A balanced static k-dimensional tree over selected atoms.
52///
53/// Construction uses linear-time median partitioning at each level, for
54/// O(n log n) total work with `n / LEAF_MAX` split nodes and four flat
55/// per-point arrays.
56#[derive(Debug)]
57pub struct KdTree<'a> {
58    positions: &'a [[f32; 3]],
59    xs: Vec<f32>,
60    ys: Vec<f32>,
61    zs: Vec<f32>,
62    ids: Vec<u32>,
63    split_value: Vec<f32>,
64    split_axis: Vec<u8>,
65    bounds: Option<([f32; 3], [f32; 3])>,
66    periodic: Option<PeriodicBox>,
67    periodic_options: KdPeriodicOptions,
68    reservation: Option<molframe_core::MemoryReservation>,
69}
70
71impl<'a> KdTree<'a> {
72    /// Builds a balanced tree over finite target positions.
73    ///
74    /// # Errors
75    ///
76    /// Returns an error when a target index is outside `positions`.
77    pub fn build(
78        positions: &'a [[f32; 3]],
79        targets: &[u32],
80        periodic: Option<PeriodicBox>,
81    ) -> Result<Self, SpatialError> {
82        Self::build_with_options(positions, targets, periodic, KdPeriodicOptions::default())
83    }
84
85    /// Builds a balanced tree with an explicit periodic image budget.
86    ///
87    /// # Errors
88    ///
89    /// Returns an error for invalid options or target indices.
90    pub fn build_with_options(
91        positions: &'a [[f32; 3]],
92        targets: &[u32],
93        periodic: Option<PeriodicBox>,
94        periodic_options: KdPeriodicOptions,
95    ) -> Result<Self, SpatialError> {
96        Self::build_parallel(positions, targets, periodic, periodic_options, None)
97    }
98
99    pub(super) fn build_parallel(
100        positions: &'a [[f32; 3]],
101        targets: &[u32],
102        periodic: Option<PeriodicBox>,
103        periodic_options: KdPeriodicOptions,
104        context: Option<&ExecutionContext>,
105    ) -> Result<Self, SpatialError> {
106        validate_indices(targets, positions.len())?;
107        periodic_options.validate()?;
108
109        let mut entries = finite_entries(positions, targets, periodic.as_ref());
110        let bounds = bounds_of(&entries);
111        let depth = depth_for(entries.len());
112        let node_count = (1_usize << depth) - 1;
113        let mut split_value = vec![0.0_f32; node_count];
114        let mut split_axis = vec![0_u8; node_count];
115        for (node, axis, value) in build_splits(&mut entries, depth, context)? {
116            split_value[node] = value;
117            split_axis[node] = axis;
118        }
119        let mut xs = Vec::with_capacity(entries.len());
120        let mut ys = Vec::with_capacity(entries.len());
121        let mut zs = Vec::with_capacity(entries.len());
122        let mut ids = Vec::with_capacity(entries.len());
123        for entry in &entries {
124            xs.push(entry.position[0]);
125            ys.push(entry.position[1]);
126            zs.push(entry.position[2]);
127            ids.push(entry.atom);
128        }
129
130        Ok(Self {
131            positions,
132            xs,
133            ys,
134            zs,
135            ids,
136            split_value,
137            split_axis,
138            bounds,
139            periodic,
140            periodic_options,
141            reservation: None,
142        })
143    }
144
145    /// Number of finite indexed points.
146    #[must_use]
147    pub fn len(&self) -> usize {
148        self.ids.len()
149    }
150
151    /// Whether the tree indexes no finite point.
152    #[must_use]
153    pub fn is_empty(&self) -> bool {
154        self.ids.is_empty()
155    }
156
157    /// Bytes retained by the finished tree.
158    pub(crate) fn retained_bytes(&self) -> usize {
159        self.ids.capacity() * (3 * size_of::<f32>() + size_of::<u32>())
160            + self.split_value.capacity() * (size_of::<f32>() + size_of::<u8>())
161    }
162
163    /// Bytes a build over `targets` may hold at its peak (entries plus result).
164    pub(crate) fn build_bytes(targets: usize) -> Option<usize> {
165        let depth = depth_for(targets);
166        targets
167            .checked_mul(size_of::<Entry>() + 3 * size_of::<f32>() + size_of::<u32>())?
168            .checked_add(((1_usize << depth) - 1) * (size_of::<f32>() + size_of::<u8>()))
169    }
170
171    /// Finds all indexed pairs within `cutoff`.
172    ///
173    /// # Errors
174    ///
175    /// Returns an error for an invalid cutoff or query index.
176    pub fn pairs(&self, query: &[u32], cutoff: f32) -> Result<Vec<NeighborPair>, SpatialError> {
177        validate_cutoff(cutoff)?;
178        validate_indices(query, self.positions.len())?;
179        let mut found = Vec::new();
180        let mut collector = self.radius_collector(cutoff)?;
181        for &atom in query {
182            collector.collect(self, atom, |target, squared| {
183                found.push(NeighborPair::new(atom, target, squared));
184            })?;
185        }
186        crate::brute::canonicalise(&mut found);
187        Ok(found)
188    }
189
190    /// Returns the nearest `count` targets in increasing distance order.
191    ///
192    /// Distances are squared. Equal distances order by atom index.
193    ///
194    /// # Errors
195    ///
196    /// Returns an error when `query` is outside the coordinate array.
197    pub fn k_nearest(&self, query: u32, count: usize) -> Result<Vec<(u32, f32)>, SpatialError> {
198        let index = usize::try_from(query).map_err(|_| SpatialError::NumericRangeExceeded)?;
199        let Some(point) = self.positions.get(index).copied() else {
200            return Err(SpatialError::AtomOutOfBounds(query));
201        };
202        if count == 0 || !finite(point) {
203            return Ok(Vec::new());
204        }
205        let mut heap = heap::KnnHeap::new(count);
206        self.knn_into(&mut heap, point, Some(query), f32::INFINITY)?;
207        Ok(heap
208            .sorted()
209            .into_iter()
210            .map(|(squared, atom)| (atom, squared))
211            .collect())
212    }
213}
214
215/// Returns finite target entries in the tree coordinate system.
216///
217/// Runtime is `O(T)` with `O(F)` output space for `F` finite targets.
218fn finite_entries(
219    positions: &[[f32; 3]],
220    targets: &[u32],
221    periodic: Option<&PeriodicBox>,
222) -> Vec<Entry> {
223    let mut entries = Vec::with_capacity(targets.len());
224    entries.extend(targets.iter().filter_map(|&atom| {
225        let index = usize::try_from(atom).ok()?;
226        let position = positions.get(index).copied()?;
227        if !finite(position) {
228            return None;
229        }
230        let position = match periodic {
231            Some(periodic) => periodic.wrap(position),
232            None => position,
233        };
234        Some(Entry { position, atom })
235    }));
236    entries
237}
238
239fn bounds_of(entries: &[Entry]) -> Option<([f32; 3], [f32; 3])> {
240    let first = entries.first()?;
241    let mut low = first.position;
242    let mut high = first.position;
243    for entry in entries {
244        for axis in 0..3 {
245            low[axis] = low[axis].min(entry.position[axis]);
246            high[axis] = high[axis].max(entry.position[axis]);
247        }
248    }
249    Some((low, high))
250}
251
252/// Smallest depth whose leaves all hold at most `LEAF_MAX` points.
253fn depth_for(count: usize) -> usize {
254    let mut depth = 0;
255    while count.div_ceil(1 << depth) > LEAF_MAX {
256        depth += 1;
257    }
258    depth
259}
260
261type Split = (usize, u8, f32);
262
263/// Splits one node's slice at its count median on the widest axis.
264fn split_slice(slice: &mut [Entry]) -> (u8, f32, usize) {
265    let mut low = [f32::INFINITY; 3];
266    let mut high = [f32::NEG_INFINITY; 3];
267    for entry in slice.iter() {
268        for axis in 0..3 {
269            low[axis] = low[axis].min(entry.position[axis]);
270            high[axis] = high[axis].max(entry.position[axis]);
271        }
272    }
273    let mut axis = 0_u8;
274    for candidate in 1..3_u8 {
275        let (current, next) = (usize::from(axis), usize::from(candidate));
276        if high[next] - low[next] > high[current] - low[current] {
277            axis = candidate;
278        }
279    }
280    let middle = slice.len() / 2;
281    let key = usize::from(axis);
282    slice.select_nth_unstable_by(middle, |left, right| {
283        left.position[key].total_cmp(&right.position[key])
284    });
285    (axis, slice[middle].position[key], middle)
286}
287
288type Task<'s> = Mutex<Option<(usize, &'s mut [Entry])>>;
289
290fn build_subtree(slice: &mut [Entry], node: usize, remaining: usize, out: &mut Vec<Split>) {
291    if remaining == 0 {
292        return;
293    }
294    let (axis, value, middle) = split_slice(slice);
295    out.push((node, axis, value));
296    let (left, right) = slice.split_at_mut(middle);
297    build_subtree(left, 2 * node + 1, remaining - 1, out);
298    build_subtree(right, 2 * node + 2, remaining - 1, out);
299}
300
301/// Computes every split node, parallelising the independent subtrees.
302fn build_splits(
303    entries: &mut [Entry],
304    depth: usize,
305    context: Option<&ExecutionContext>,
306) -> Result<Vec<Split>, SpatialError> {
307    let workers = context.map_or(1, ExecutionContext::worker_budget);
308    let parallel = workers > 1 && entries.len() >= 16_384 && depth > 0;
309    let mut splits = Vec::new();
310    if !parallel {
311        build_subtree(entries, 0, depth, &mut splits);
312        return Ok(splits);
313    }
314    // Enough subtrees to balance dynamic scheduling, few enough that the
315    // serial top levels stay a small fraction of the work.
316    let top = depth.min((workers * 4).next_power_of_two().trailing_zeros() as usize);
317    let mut frontier = vec![(0_usize, entries)];
318    for _ in 0..top {
319        let mut next = Vec::with_capacity(frontier.len() * 2);
320        for (node, slice) in frontier {
321            let (axis, value, middle) = split_slice(slice);
322            splits.push((node, axis, value));
323            let (left, right) = slice.split_at_mut(middle);
324            next.push((2 * node + 1, left));
325            next.push((2 * node + 2, right));
326        }
327        frontier = next;
328    }
329    let remaining = depth - top;
330    let tasks: Vec<Task> = frontier
331        .into_iter()
332        .map(|task| Mutex::new(Some(task)))
333        .collect();
334    let Some(context) = context else {
335        return Err(SpatialError::WorkerPanicked);
336    };
337    let produced = map_blocks_in(BlockPlan::new(tasks.len(), 1), context, |index, _| {
338        let mut out = Vec::new();
339        let task = tasks.get(index).and_then(|slot| slot.lock().ok()?.take());
340        if let Some((node, slice)) = task {
341            build_subtree(slice, node, remaining, &mut out);
342        }
343        out
344    })
345    .map_err(|_| SpatialError::WorkerPanicked)?;
346    for mut block in produced {
347        splits.append(&mut block);
348    }
349    Ok(splits)
350}
351
352#[cfg(test)]
353#[path = "kd_tests.rs"]
354mod tests;
355
356#[cfg(test)]
357#[path = "kd_adversarial_tests.rs"]
358mod adversarial_tests;
molframe_spatial/backends/kd.rs

molframe_spatial/backends/
kd.rs
