hoomd_spatial/
vec_cell.rs

1// Copyright (c) 2024-2026 The Regents of the University of Michigan.
2// Part of hoomd-rs, released under the BSD 3-Clause License.
3
4#![expect(
5    clippy::cast_possible_truncation,
6    reason = "the necessary conversions are necessary and have been checked"
7)]
8#![expect(
9    clippy::cast_sign_loss,
10    reason = "the necessary conversions are necessary and have been checked"
11)]
12
13//! Implement `VecCell`
14
15use serde::{Deserialize, Serialize};
16use serde_with::serde_as;
17use std::{array, cmp::Eq, fmt, hash::Hash, iter, marker::PhantomData, mem};
18
19use log::trace;
20use rustc_hash::FxHashMap;
21
22use hoomd_utility::valid::PositiveReal;
23use hoomd_vector::Cartesian;
24
25use super::{PointUpdate, PointsNearBall, WithSearchRadius};
26
27use crate::{IndexFromPosition, hash_cell::CellIndex};
28
29/// Bucket sort points into cubes with [`Vec`]-backed storage
30///
31/// [`VecCell`] is a spatial data structure that can efficiently update one
32/// point at a time (see [`PointUpdate`]) and find all points that are near a
33/// position in space (see [`PointsNearBall`]).
34///
35/// Points are stored in [`VecCell`] by key. The caller can choose a generic key
36/// type `K` appropriately. For example, `K` could be `usize` and map to an
37/// array index. Or it could be an enum that distinguishes between primary and
38/// ghost site positions.
39///
40/// Internally, [`VecCell`] stores a `D`-dimensional vector of cells that
41/// partition space. Each bucket is a hypercube with an edge length equal to
42/// the *nominal search radius*. When searching for points near a position in
43/// space, `VecCell` first determines which bucket that point is in. Then it
44/// iterates over all the keys in all the cells within the given search radius.
45/// To perform this iteration efficiently, [`VecCell`] pre-computes a set of
46/// stencils up to the given *maximum search radius* given at const
47///
48/// In systems with low density fluctuations, the time complexity of [`VecCell`]
49/// is O(1). More generally, and important when density fluctuations are
50/// present, the time complexity scales with the maximum number of points in any
51/// of the cells.
52///
53/// [`VecCell`] performs best when the search radius is equal to the nominal
54/// search radius. In that case, 2D searches loop over 9 cells and 3D searches
55/// loop over 27. The caller must filter the output of this iteration because
56/// the resulting hypercube covers more volume than the given hypersphere.
57/// When searching with a radius larger than the nominal, many more cells must
58/// be iterated leading to slower performance. At the same time, the stencil
59/// covers less excess volume and therefore fewer points must be filtered.
60///
61/// [`VecCell`] stores each cell at a unique index backed by storage in [`Vec`].
62/// This performs very well in dense, compact systems. However, the memory
63/// utilization scales with the *volume* of the cube that spans from the most
64/// negative point to the most positive point. In dilute systems, [`VecCell`]
65/// can therefore waste a lot of memory storing empty cells. Use [`HashCell`]
66/// and occasionally call its `shrink_to_fit` method to limit the memory
67/// required to model dilute systems.
68///
69/// [`HashCell`]: crate::HashCell
70///
71/// # Examples
72///
73/// The default [`VecCell`] set both the *nominal* and *maximum* search radii to 1.0.
74/// ```
75/// use hoomd_spatial::VecCell;
76///
77/// let vec_cell = VecCell::<usize, 3>::default();
78/// ```
79///
80/// Use the builder API to set any or all parameters:
81///
82/// ```
83/// use hoomd_spatial::VecCell;
84///
85/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
86/// let vec_cell = VecCell::<usize, 3>::builder()
87///     .nominal_search_radius(2.5.try_into()?)
88///     .maximum_search_radius(7.5)
89///     .build();
90/// # Ok(())
91/// # }
92/// ```
93#[serde_as]
94#[derive(Clone, Debug, Serialize, Deserialize)]
95pub struct VecCell<K, const D: usize>
96where
97    K: Eq + Hash,
98{
99    /// The width of each cell.
100    cell_width: PositiveReal,
101
102    /// A map from cell indices to cell contents.
103    keys_map: Vec<Vec<K>>,
104
105    /// A map from particle indices to cell indices.
106    cell_index: FxHashMap<K, CellIndex<D>>,
107
108    /// The shape of `keys_map` is `(half_extent * 2 + 1).powi(D)`.
109    half_extent: u32,
110
111    /// Pre-computed stencils.
112    #[serde_as(as = "Vec<Vec<[_; D]>>")]
113    stencils: Vec<Vec<[i64; D]>>,
114}
115
116/// Iterate over cell indices in row major order.
117struct CellIndexIterator<const D: usize> {
118    /// The current cell index.
119    cell_index: [i64; D],
120    /// The half extent of the cube.
121    half_extent: u32,
122}
123
124impl<const D: usize> CellIndexIterator<D> {
125    /// Iterate over a cube.
126    ///
127    /// The cube extends from `[-half_extent, -half_extent, ..., -half_extent]` to
128    /// `[half_extent, half_extent, ..., half_extent]`.
129    fn cube(half_extent: u32) -> Self {
130        let mut cell_index = [-i64::from(half_extent); D];
131        cell_index[D - 1] -= 1;
132        Self {
133            cell_index,
134            half_extent,
135        }
136    }
137
138    /// Increment the cell index.
139    #[inline]
140    fn increment_cell_index(&mut self) -> Option<[i64; D]> {
141        self.cell_index[D - 1] += 1;
142
143        for i in (0..D).rev() {
144            if self.cell_index[i] > self.half_extent.into() {
145                if i == 0 {
146                    return None;
147                }
148
149                self.cell_index[i] = -(i64::from(self.half_extent));
150                self.cell_index[i - 1] += 1;
151            }
152        }
153
154        Some(self.cell_index)
155    }
156}
157
158impl<const D: usize> Iterator for CellIndexIterator<D> {
159    type Item = [i64; D];
160
161    fn next(&mut self) -> Option<Self::Item> {
162        self.increment_cell_index()
163    }
164}
165
166/// Generate the stencil for a given radius.
167fn generate_stencil<const D: usize>(radius: u32) -> Vec<[i64; D]> {
168    assert!(radius >= 1, "cell list must have a minimum radius of 1");
169
170    let mut result = CellIndexIterator::cube(radius).collect::<Vec<_>>();
171    result.sort_by(|a, b| {
172        let r_a = a.iter().map(|x| x.pow(2));
173        let r_b = b.iter().map(|x| x.pow(2));
174        r_a.cmp(r_b)
175    });
176    result
177}
178
179/// Generate the stencils up to a given radius.
180pub(crate) fn generate_all_stencils<const D: usize>(max_radius: u32) -> Vec<Vec<[i64; D]>> {
181    assert!(max_radius >= 1, "cell list must have a minimum radius of 1");
182    let mut result = Vec::new();
183
184    // A cell list will never use radius=0, so stencil[i] stores the stencil needed for
185    // radius i+1.
186    for radius in 0..max_radius {
187        result.push(generate_stencil(radius + 1));
188    }
189    result
190}
191
192/// Construct a [`VecCell`] with given parameters.
193///
194/// # Example
195///
196/// ```
197/// use hoomd_spatial::VecCell;
198///
199/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
200/// let vec_cell = VecCell::<usize, 3>::builder()
201///     .nominal_search_radius(2.5.try_into()?)
202///     .maximum_search_radius(7.5)
203///     .build();
204/// # Ok(())
205/// # }
206/// ```
207pub struct VecCellBuilder<K, const D: usize> {
208    /// Most commonly used search radius.
209    nominal_search_radius: PositiveReal,
210
211    /// Largest possible search radius.
212    maximum_search_radius: f64,
213
214    /// Track the key type.
215    phantom_key: PhantomData<K>,
216}
217
218impl<K, const D: usize> VecCellBuilder<K, D>
219where
220    K: Copy + Eq + Hash,
221{
222    /// Choose the most commonly used search radius.
223    ///
224    /// [`VecCell`] performs the best when searching for points within the
225    /// *nominal search radius* of a given position.
226    ///
227    /// # Example
228    /// ```
229    /// use hoomd_spatial::VecCell;
230    ///
231    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
232    /// let vec_cell = VecCell::<usize, 3>::builder()
233    ///     .nominal_search_radius(2.5.try_into()?)
234    ///     .build();
235    /// # Ok(())
236    /// # }
237    /// ```
238    #[inline]
239    #[must_use]
240    pub fn nominal_search_radius(mut self, nominal_search_radius: PositiveReal) -> Self {
241        self.nominal_search_radius = nominal_search_radius;
242        self
243    }
244
245    /// Choose the largest search radius.
246    ///
247    /// The maximum radius is rounded up to the nearest integer multiple of the
248    /// *nominal search radius*. [`VecCell`] will panic when asked to search for
249    /// points within a radius larger than the maximum.
250    ///
251    /// # Example
252    /// ```
253    /// use hoomd_spatial::VecCell;
254    ///
255    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
256    /// let vec_cell = VecCell::<usize, 3>::builder()
257    ///     .nominal_search_radius(2.5.try_into()?)
258    ///     .maximum_search_radius(7.5)
259    ///     .build();
260    /// # Ok(())
261    /// # }
262    /// ```
263    #[inline]
264    #[must_use]
265    pub fn maximum_search_radius(mut self, maximum_search_radius: f64) -> Self {
266        self.maximum_search_radius = maximum_search_radius;
267        self
268    }
269
270    /// Construct the [`VecCell`] with the chosen parameters.
271    ///
272    /// # Example
273    /// ```
274    /// use hoomd_spatial::VecCell;
275    ///
276    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
277    /// let vec_cell = VecCell::<usize, 3>::builder()
278    ///     .nominal_search_radius(2.5.try_into()?)
279    ///     .maximum_search_radius(7.5)
280    ///     .build();
281    /// # Ok(())
282    /// # }
283    /// ```
284    #[inline]
285    #[must_use]
286    pub fn build(self) -> VecCell<K, D> {
287        let maximum_stencil_radius =
288            (self.maximum_search_radius / self.nominal_search_radius.get()).ceil() as u32;
289        let half_extent: u32 = 1;
290
291        VecCell {
292            cell_width: self.nominal_search_radius,
293            keys_map: iter::repeat_n(Vec::new(), (half_extent * 2 + 1).pow(D as u32) as usize)
294                .collect(),
295            cell_index: FxHashMap::default(),
296            half_extent,
297            stencils: generate_all_stencils(maximum_stencil_radius.max(1)),
298        }
299    }
300}
301
302impl<K, const D: usize> Default for VecCell<K, D>
303where
304    K: Copy + Eq + Hash,
305{
306    /// Construct a default [`VecCell`].
307    ///
308    /// The default sets both the *nominal* and *maximum* search radii to 1.0.
309    ///
310    /// # Example
311    ///
312    /// ```
313    /// use hoomd_spatial::VecCell;
314    ///
315    /// let vec_cell = VecCell::<usize, 3>::default();
316    /// ```
317    #[inline]
318    fn default() -> Self {
319        Self::builder().build()
320    }
321}
322
323impl<K, const D: usize> WithSearchRadius for VecCell<K, D>
324where
325    K: Copy + Eq + Hash,
326{
327    /// Construct a [`VecCell`] with the given search radius.
328    ///
329    /// Set both the *nominal* and *maximum* search radii to `radius`.
330    ///
331    /// # Example
332    ///
333    /// ```
334    /// use hoomd_spatial::{VecCell, WithSearchRadius};
335    ///
336    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
337    /// let vec_cell = VecCell::<usize, 3>::with_search_radius(2.5.try_into()?);
338    /// # Ok(())
339    /// # }
340    /// ```
341    #[inline]
342    fn with_search_radius(radius: PositiveReal) -> Self {
343        Self::builder()
344            .nominal_search_radius(radius)
345            .maximum_search_radius(radius.get())
346            .build()
347    }
348}
349
350impl<K, const D: usize> VecCell<K, D>
351where
352    K: Eq + Hash,
353{
354    /// Compute the cell index given a position in space.
355    #[inline]
356    fn cell_index_from_position(&self, position: &Cartesian<D>) -> [i64; D] {
357        std::array::from_fn(|j| (position.coordinates[j] / self.cell_width.get()).floor() as i64)
358    }
359
360    /// Compute the vector index from a cell index
361    ///
362    /// Returns `None` when the index is out of bounds.
363    #[inline]
364    fn map_index_from_cell(half_extent: u32, cell_index: &[i64; D]) -> Option<usize> {
365        assert!(D > 1);
366
367        let mut vec_index: usize = 0;
368        let mut width = 1;
369
370        for i in (0..D).rev() {
371            let needed_extent = cell_index[i].unsigned_abs();
372            if needed_extent > u64::from(half_extent) {
373                return None;
374            }
375            let v: usize = (cell_index[i] + i64::from(half_extent))
376                .try_into()
377                .expect("cell index should be in bounds");
378
379            vec_index += v * width;
380            width *= (half_extent * 2 + 1) as usize;
381        }
382        Some(vec_index)
383    }
384
385    /// Get the keys in a given cell index
386    #[cfg(test)]
387    #[inline]
388    fn get_keys(&self, cell_index: &[i64; D]) -> &[K] {
389        let index = Self::map_index_from_cell(self.half_extent, cell_index)
390            .expect("cell_index should be in bounds");
391        &self.keys_map[index]
392    }
393}
394
395impl<K, const D: usize> VecCell<K, D>
396where
397    K: Copy + Eq + Hash,
398{
399    /// Construct a `VecCell` builder.
400    ///
401    /// Use the builder to set any or all parameters and construct a [`VecCell`].
402    ///
403    /// # Example
404    ///
405    /// ```
406    /// use hoomd_spatial::VecCell;
407    ///
408    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
409    /// let vec_cell = VecCell::<usize, 3>::builder()
410    ///     .nominal_search_radius(2.5.try_into()?)
411    ///     .maximum_search_radius(7.5)
412    ///     .build();
413    /// # Ok(())
414    /// # }
415    /// ```
416    #[expect(
417        clippy::missing_panics_doc,
418        reason = "hard-coded constant will never panic"
419    )]
420    #[inline]
421    #[must_use]
422    pub fn builder() -> VecCellBuilder<K, D> {
423        VecCellBuilder {
424            nominal_search_radius: 1.0
425                .try_into()
426                .expect("hard-coded constant is a positive real"),
427            maximum_search_radius: 1.0,
428            phantom_key: PhantomData,
429        }
430    }
431
432    /// Remove excess capacity from dynamically allocated arrays.
433    ///
434    /// At this time, `shrink_to_fit` only reduces the memory utilized by the
435    /// cell contents. It does not shrink the `D`-dimensional storage to match
436    /// the range spanned by points currently in the data structure.
437    #[inline]
438    pub fn shrink_to_fit(&mut self) {
439        for keys in &mut self.keys_map {
440            keys.shrink_to_fit();
441        }
442        self.keys_map.shrink_to_fit();
443        self.cell_index.shrink_to_fit();
444    }
445
446    /// Double the number of cells stored along each axis until it includes the target.
447    fn expand_to(&mut self, target: u32) {
448        if self.half_extent >= target {
449            return;
450        }
451
452        let mut new_half_extent = self.half_extent.min(1) * 2;
453
454        while new_half_extent < target {
455            new_half_extent *= 2;
456        }
457
458        trace!("Expanding to {}^{} cells", new_half_extent * 2 + 1, D);
459
460        let mut new_keys_map: Vec<Vec<K>> =
461            iter::repeat_n(Vec::new(), (new_half_extent * 2 + 1).pow(D as u32) as usize).collect();
462        let old_half_extent = self.half_extent;
463        let old_keys_map = &mut self.keys_map;
464
465        for old_cell_index in CellIndexIterator::cube(old_half_extent) {
466            let old_vec_index = Self::map_index_from_cell(old_half_extent, &old_cell_index)
467                .expect("cell_index should be consistent with keys_map");
468            let new_vec_index = Self::map_index_from_cell(new_half_extent, &old_cell_index)
469                .expect("old_cell_index should be inside the new keys_map");
470            new_keys_map[new_vec_index] = mem::take(&mut old_keys_map[old_vec_index]);
471        }
472
473        self.half_extent = new_half_extent;
474        self.keys_map = new_keys_map;
475    }
476}
477
478impl<K, const D: usize> PointUpdate<Cartesian<D>, K> for VecCell<K, D>
479where
480    K: Copy + Eq + Hash,
481{
482    /// Insert or update a point identified by a key.
483    ///
484    /// # Example
485    /// ```
486    /// use hoomd_spatial::{PointUpdate, VecCell};
487    ///
488    /// let mut vec_cell = VecCell::default();
489    ///
490    /// vec_cell.insert(0, [1.25, 2.5].into());
491    /// ```
492    #[inline]
493    fn insert(&mut self, key: K, position: Cartesian<D>) {
494        let cell_index = self.cell_index_from_position(&position);
495        let old_cell_index = self.cell_index.insert(key, CellIndex(cell_index));
496        let map_index =
497            Self::map_index_from_cell(self.half_extent, &cell_index).unwrap_or_else(|| {
498                let max_half_extent = cell_index
499                    .iter()
500                    .map(|x| x.unsigned_abs())
501                    .reduce(u64::max)
502                    .expect("D should be greater than 1");
503                self.expand_to(
504                    max_half_extent
505                        .try_into()
506                        .expect("max extent cannot exceed u32::MAX"),
507                );
508                Self::map_index_from_cell(self.half_extent, &cell_index)
509                    .expect("cell_index should be in the expanded VecCell")
510            });
511
512        // This checks if old_cell_index is None or if it is different from the new cell index.
513        if old_cell_index != Some(CellIndex(cell_index)) {
514            // Add the particle index to the new cell index vector.
515            self.keys_map[map_index].push(key);
516
517            if let Some(old_cell_index) = old_cell_index {
518                // If the particle was in a different cell, we need to remove it from the old cell.
519                let old_map_index = Self::map_index_from_cell(self.half_extent, &old_cell_index.0)
520                    .expect("cell_index and keys_map should agree");
521                let old_keys = &mut self.keys_map[old_map_index];
522                if let Some(pos) = old_keys.iter().position(|x| *x == key) {
523                    old_keys.swap_remove(pos);
524                }
525            }
526        }
527    }
528
529    /// Remove the point with the given key.
530    ///
531    /// # Example
532    /// ```
533    /// use hoomd_spatial::{PointUpdate, VecCell};
534    ///
535    /// let mut vec_cell = VecCell::default();
536    /// vec_cell.insert(0, [1.25, 2.5].into());
537    ///
538    /// vec_cell.remove(&0)
539    /// ```
540    #[inline]
541    fn remove(&mut self, key: &K) {
542        let cell_index = self.cell_index.remove(key);
543        if let Some(cell_index) = cell_index {
544            let map_index = Self::map_index_from_cell(self.half_extent, &cell_index.0);
545            if let Some(map_index) = map_index {
546                let keys = &mut self.keys_map[map_index];
547                if let Some(idx) = keys.iter().position(|x| x == key) {
548                    keys.swap_remove(idx);
549                }
550            }
551        }
552    }
553
554    /// Get the number of points in the spatial data structure.
555    ///
556    /// # Example
557    /// ```
558    /// use hoomd_spatial::{PointUpdate, VecCell};
559    ///
560    /// let mut vec_cell = VecCell::default();
561    /// vec_cell.insert(0, [1.25, 2.5].into());
562    ///
563    /// assert_eq!(vec_cell.len(), 1)
564    /// ```
565    #[inline]
566    fn len(&self) -> usize {
567        self.cell_index.len()
568    }
569
570    /// Test if the spatial data structure is empty.
571    ///
572    /// # Example
573    /// ```
574    /// use hoomd_spatial::{PointUpdate, VecCell};
575    ///
576    /// let mut vec_cell = VecCell::<usize, 2>::default();
577    ///
578    /// assert!(vec_cell.is_empty());
579    /// ```
580    #[inline]
581    fn is_empty(&self) -> bool {
582        self.cell_index.is_empty()
583    }
584
585    /// Test if the spatial data structure contains a key.
586    /// ```
587    /// use hoomd_spatial::{PointUpdate, VecCell};
588    ///
589    /// let mut vec_cell = VecCell::default();
590    /// vec_cell.insert(0, [1.25, 2.5].into());
591    ///
592    /// assert!(vec_cell.contains_key(&0));
593    /// ```
594    #[inline]
595    fn contains_key(&self, key: &K) -> bool {
596        self.cell_index.contains_key(key)
597    }
598
599    /// Remove all points.
600    ///
601    /// # Example
602    /// ```
603    /// use hoomd_spatial::{PointUpdate, VecCell};
604    ///
605    /// let mut vec_cell = VecCell::default();
606    /// vec_cell.insert(0, [1.25, 2.5].into());
607    ///
608    /// vec_cell.clear();
609    /// ```
610    #[inline]
611    fn clear(&mut self) {
612        self.cell_index.clear();
613        for keys in &mut self.keys_map {
614            keys.clear();
615        }
616    }
617}
618
619/// Iterate over keys in the cell list around a given center cell.
620struct PointsIterator<'a, K, const D: usize>
621where
622    K: Eq + Hash,
623{
624    /// Keys of the current cell iteration (None if the cell is empty)
625    keys: Option<&'a Vec<K>>,
626
627    /// The cell list we are iterating in.
628    cell_list: &'a VecCell<K, D>,
629
630    /// Current location of the iteration in the cell.
631    index_in_current_cell: usize,
632
633    /// Current location of the iteration in the stencil.
634    current_stencil: usize,
635
636    /// Cell offsets to iterate over.
637    stencil: &'a [[i64; D]],
638
639    /// The cell at the center of the iteration.
640    center: [i64; D],
641}
642
643impl<K, const D: usize> Iterator for PointsIterator<'_, K, D>
644where
645    K: Copy + Eq + Hash,
646{
647    type Item = K;
648
649    #[inline]
650    fn next(&mut self) -> Option<Self::Item> {
651        loop {
652            if let Some(keys) = self.keys
653                && self.index_in_current_cell < keys.len()
654            {
655                let last_index = self.index_in_current_cell;
656                self.index_in_current_cell += 1;
657                return Some(keys[last_index]);
658            }
659
660            self.index_in_current_cell = 0;
661            self.current_stencil += 1;
662
663            if self.current_stencil >= self.stencil.len() {
664                return None;
665            }
666
667            let cell_index =
668                array::from_fn(|i| self.center[i] + self.stencil[self.current_stencil][i]);
669            let map_index =
670                VecCell::<K, D>::map_index_from_cell(self.cell_list.half_extent, &cell_index);
671            self.keys = map_index.map(|index| &self.cell_list.keys_map[index]);
672        }
673    }
674}
675
676impl<const D: usize, K> PointsNearBall<Cartesian<D>, K> for VecCell<K, D>
677where
678    K: Copy + Eq + Hash,
679{
680    /// Find all the points that *might* be in the given ball.
681    ///
682    /// `points_near_ball` will iterate over all points in the given ball *and
683    /// possibly others as well*. [`VecCell`] may iterate over the points in
684    /// any order.
685    ///
686    /// # Example
687    /// ```
688    /// use hoomd_spatial::{PointUpdate, PointsNearBall, VecCell};
689    ///
690    /// let mut vec_cell = VecCell::default();
691    /// vec_cell.insert(0, [1.25, 0.0].into());
692    /// vec_cell.insert(1, [3.25, 0.75].into());
693    /// vec_cell.insert(2, [-10.0, 12.0].into());
694    ///
695    /// for key in vec_cell.points_near_ball(&[2.0, 0.0].into(), 1.0) {
696    ///     println!("{key}");
697    /// }
698    /// ```
699    /// Prints (in any order):
700    /// ```text
701    /// 0
702    /// 1
703    /// ```
704    ///
705    /// # Panics
706    ///
707    /// Panics when `radius` is larger than the *maximum search radius*
708    /// provided at construction, rounded up to the nearest integer multiple
709    /// of the *nominal search radius*.
710    #[inline]
711    fn points_near_ball(&self, position: &Cartesian<D>, radius: f64) -> impl Iterator<Item = K> {
712        let stencil_index = (radius / self.cell_width.get()).ceil() as usize - 1;
713        assert!(
714            stencil_index < self.stencils.len(),
715            "search radius must be less than or equal to the maximum search radius"
716        );
717
718        let center = self.cell_index_from_position(position);
719        let stencil = &self.stencils[stencil_index];
720        let map_index = Self::map_index_from_cell(
721            self.half_extent,
722            &array::from_fn(|i| center[i] + stencil[0][i]),
723        );
724
725        PointsIterator {
726            keys: map_index.map(|index| &self.keys_map[index]),
727            cell_list: self,
728            index_in_current_cell: 0,
729            current_stencil: 0,
730            stencil,
731            center,
732        }
733    }
734}
735
736impl<K, const D: usize> fmt::Display for VecCell<K, D>
737where
738    K: Eq + Hash,
739{
740    /// Summarize the contents of the cell list.
741    ///
742    /// This is a slow operation. It is meant to be printed to logs only
743    /// occasionally, such as at the end of a benchmark or simulation.
744    ///
745    /// # Example
746    ///
747    /// ```
748    /// use hoomd_spatial::VecCell;
749    /// use log::info;
750    ///
751    /// let vec_cell = VecCell::<usize, 3>::default();
752    ///
753    /// info!("{vec_cell}");
754    /// ```
755    #[allow(
756        clippy::missing_inline_in_public_items,
757        reason = "no need to inline display"
758    )]
759    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
760        let largest_cell_size = self.keys_map.iter().map(Vec::len).fold(0, usize::max);
761
762        writeln!(f, "VecCell<K, {D}>:")?;
763        writeln!(
764            f,
765            "- {} total cells with {} cells on each side.",
766            self.keys_map.len(),
767            self.half_extent * 2 + 1
768        )?;
769        writeln!(f, "- {} points.", self.cell_index.len())?;
770        writeln!(
771            f,
772            "- Nominal, maximum search radii: {}, {}",
773            self.cell_width,
774            self.cell_width.get() * self.stencils.len() as f64
775        )?;
776        write!(f, "- Largest cell length: {largest_cell_size}")
777    }
778}
779
780impl<K, const D: usize> IndexFromPosition<Cartesian<D>> for VecCell<K, D>
781where
782    K: Eq + Hash,
783{
784    type Location = [i64; D];
785
786    #[inline]
787    fn location_from_position(&self, position: &Cartesian<D>) -> Self::Location {
788        self.cell_index_from_position(position)
789    }
790}
791
792#[expect(
793    clippy::used_underscore_binding,
794    reason = "Used for const parameterization."
795)]
796#[cfg(test)]
797mod tests {
798    use assert2::{assert, check};
799    use rand::{
800        RngExt, SeedableRng,
801        distr::{Distribution, Uniform},
802        rngs::StdRng,
803    };
804    use rstest::*;
805
806    use super::*;
807    use hoomd_vector::{Metric, distribution::Ball};
808
809    #[test]
810    fn test_increment_cell_index() {
811        let mut cells = CellIndexIterator::cube(1);
812        check!(cells.next() == Some([-1, -1]));
813        check!(cells.next() == Some([-1, 0]));
814        check!(cells.next() == Some([-1, 1]));
815        check!(cells.next() == Some([0, -1]));
816        check!(cells.next() == Some([0, 0]));
817        check!(cells.next() == Some([0, 1]));
818        check!(cells.next() == Some([1, -1]));
819        check!(cells.next() == Some([1, 0]));
820        check!(cells.next() == Some([1, 1]));
821        check!(cells.next() == None);
822        check!(cells.next() == None);
823        check!(cells.next() == None);
824        check!(cells.next() == None);
825
826        let mut c = CellIndexIterator {
827            cell_index: [1, 2, 2],
828            half_extent: 2,
829        };
830        check!(c.increment_cell_index() == Some([2, -2, -2]));
831        let mut c = CellIndexIterator {
832            cell_index: [0, 1, 2],
833            half_extent: 2,
834        };
835        check!(c.increment_cell_index() == Some([0, 2, -2]));
836        let mut c = CellIndexIterator {
837            cell_index: [0, 0, -2],
838            half_extent: 2,
839        };
840        check!(c.increment_cell_index() == Some([0, 0, -1]));
841        let mut c = CellIndexIterator {
842            cell_index: [2, 2, 2],
843            half_extent: 2,
844        };
845        check!(c.increment_cell_index() == None);
846    }
847
848    #[test]
849    fn test_cell_index() {
850        let cell_list = VecCell::<usize, 3>::builder()
851            .nominal_search_radius(
852                2.0.try_into()
853                    .expect("hard-coded constant should be positive"),
854            )
855            .build();
856        check!(cell_list.cell_index_from_position(&[0.0, 0.0, 0.0].into()) == [0, 0, 0]);
857        check!(cell_list.cell_index_from_position(&[2.0, 0.0, 0.0].into()) == [1, 0, 0]);
858        check!(cell_list.cell_index_from_position(&[0.0, 2.0, 0.0].into()) == [0, 1, 0]);
859        check!(cell_list.cell_index_from_position(&[0.0, 0.0, 2.0].into()) == [0, 0, 1]);
860        check!(cell_list.cell_index_from_position(&[-41.5, 18.5, -0.125].into()) == [-21, 9, -1]);
861    }
862
863    #[test]
864    fn test_insert_one() {
865        let mut cell_list = VecCell::default();
866
867        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
868
869        assert!(cell_list.cell_index.get(&0) == Some(&CellIndex([0, 0])));
870
871        let keys = cell_list.get_keys(&[0, 0]);
872        assert!(keys.len() == 1);
873        assert!(keys.contains(&0));
874    }
875
876    #[test]
877    fn test_insert_many() {
878        let mut cell_list = VecCell::default();
879
880        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
881        cell_list.insert(1, Cartesian::from([0.995, 0.897]));
882        cell_list.insert(2, Cartesian::from([-0.125, 3.25]));
883
884        check!(cell_list.cell_index.get(&0) == Some(&CellIndex([0, 0])));
885        check!(cell_list.cell_index.get(&1) == Some(&CellIndex([0, 0])));
886        check!(cell_list.cell_index.get(&2) == Some(&CellIndex([-1, 3])));
887
888        let keys = cell_list.get_keys(&[0, 0]);
889        assert!(keys.len() == 2);
890        check!(keys.contains(&0));
891        check!(keys.contains(&1));
892
893        let keys = cell_list.get_keys(&[-1, 3]);
894        assert!(keys.len() == 1);
895        check!(keys.contains(&2));
896    }
897
898    #[test]
899    fn test_insert_again_same() {
900        let mut cell_list = VecCell::default();
901
902        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
903        cell_list.insert(0, Cartesian::from([0.25, 0.5]));
904        cell_list.insert(0, Cartesian::from([0.5, 0.75]));
905
906        check!(cell_list.cell_index.get(&0) == Some(&CellIndex([0, 0])));
907
908        let keys = cell_list.get_keys(&[0, 0]);
909        assert!(keys.len() == 1);
910        check!(keys.contains(&0));
911    }
912
913    #[test]
914    fn test_insert_again_different() {
915        let mut cell_list = VecCell::default();
916
917        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
918        cell_list.insert(1, Cartesian::from([0.25, 0.5]));
919        cell_list.insert(1, Cartesian::from([-0.5, -0.75]));
920
921        check!(cell_list.cell_index.get(&0) == Some(&CellIndex([0, 0])));
922        check!(cell_list.cell_index.get(&1) == Some(&CellIndex([-1, -1])));
923
924        let keys = cell_list.get_keys(&[0, 0]);
925        assert!(keys.len() == 1);
926        check!(keys.contains(&0));
927
928        let keys = cell_list.get_keys(&[-1, -1]);
929        assert!(keys.len() == 1);
930        check!(keys.contains(&1));
931    }
932
933    #[test]
934    fn test_remove() {
935        let mut cell_list = VecCell::default();
936
937        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
938        cell_list.insert(1, Cartesian::from([0.995, 0.897]));
939        cell_list.insert(2, Cartesian::from([-0.125, 3.25]));
940
941        cell_list.remove(&1);
942        cell_list.remove(&2);
943
944        check!(cell_list.cell_index.get(&0) == Some(&CellIndex([0, 0])));
945        check!(cell_list.cell_index.get(&1) == None);
946        check!(cell_list.cell_index.get(&2) == None);
947
948        let keys = cell_list.get_keys(&[0, 0]);
949        assert!(keys.len() == 1);
950        check!(keys.contains(&0));
951
952        let keys = cell_list.get_keys(&[-1, 3]);
953        check!(keys.is_empty());
954    }
955
956    #[test]
957    fn test_clear() {
958        let mut cell_list = VecCell::default();
959
960        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
961        cell_list.insert(1, Cartesian::from([0.995, 0.897]));
962        cell_list.insert(2, Cartesian::from([-0.125, 3.25]));
963
964        cell_list.clear();
965
966        check!(cell_list.cell_index.is_empty());
967        for keys in cell_list.keys_map {
968            check!(keys.is_empty());
969        }
970    }
971
972    #[test]
973    fn test_shrink_to_fit() {
974        let mut cell_list = VecCell::default();
975
976        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
977        cell_list.insert(1, Cartesian::from([0.995, 0.897]));
978        cell_list.insert(2, Cartesian::from([-0.125, 3.25]));
979        cell_list.insert(3, Cartesian::from([10.995, -12.897]));
980        cell_list.insert(4, Cartesian::from([15.125, 19.25]));
981
982        cell_list.remove(&1);
983        cell_list.remove(&2);
984        cell_list.remove(&3);
985        cell_list.remove(&4);
986
987        cell_list.shrink_to_fit();
988
989        let filled_cell_index =
990            VecCell::<usize, 2>::map_index_from_cell(cell_list.half_extent, &[0, 0])
991                .expect("hard-coded cell is valid");
992        for (i, keys) in cell_list.keys_map.iter().enumerate() {
993            if i == filled_cell_index {
994                check!(keys.capacity() == 1);
995                check!(keys.len() == 1);
996            } else {
997                check!(keys.capacity() == 0);
998                check!(keys.len() == 0);
999            }
1000        }
1001
1002        let keys = cell_list.get_keys(&[0, 0]);
1003        assert!(keys.len() == 1);
1004        check!(keys.contains(&0));
1005    }
1006
1007    #[rstest]
1008    fn consistency() {
1009        const N_STEPS: usize = 65_536;
1010        let mut rng = StdRng::seed_from_u64(0);
1011        let mut reference = FxHashMap::default();
1012
1013        let cell_width = 0.5;
1014        let mut cell_list = VecCell::builder()
1015            .nominal_search_radius(
1016                cell_width
1017                    .try_into()
1018                    .expect("hard-coded cell with should be positive"),
1019            )
1020            .build();
1021        let position_distribution = Ball {
1022            radius: 20.0.try_into().expect("hardcoded value should be positive"),
1023        };
1024        let key_distribution =
1025            Uniform::new(0, N_STEPS / 4).expect("hardcoded distribution should be valid");
1026
1027        for _ in 0..N_STEPS {
1028            // Add more keys than removing
1029            if rng.random_bool(0.7) {
1030                let position: Cartesian<3> = position_distribution.sample(&mut rng);
1031                let key = key_distribution.sample(&mut rng);
1032
1033                cell_list.insert(key, position);
1034                reference.insert(key, cell_list.cell_index_from_position(&position));
1035            } else {
1036                let key = key_distribution.sample(&mut rng);
1037                cell_list.remove(&key);
1038                reference.remove(&key);
1039            }
1040        }
1041
1042        // Validate that cell_index contains the expected keys and that
1043        // keys_map is consistent.
1044        assert!(cell_list.cell_index.len() == reference.len());
1045        for (reference_key, reference_value) in reference.drain() {
1046            let value = cell_list.cell_index.get(&reference_key);
1047            check!(value == Some(&CellIndex(reference_value)));
1048
1049            let keys = cell_list.get_keys(&reference_value);
1050            check!(keys.contains(&reference_key));
1051        }
1052
1053        // Ensure that there are no extra values in keys_map.
1054        let total = cell_list.keys_map.iter().map(Vec::len).sum();
1055        check!(cell_list.cell_index.len() == total);
1056        check!(total > 2000);
1057    }
1058
1059    #[test]
1060    fn test_outside() {
1061        let mut cell_list = VecCell::default();
1062
1063        cell_list.insert(0, Cartesian::from([0.125, 0.25]));
1064        cell_list.insert(1, Cartesian::from([0.995, 0.897]));
1065        cell_list.insert(2, Cartesian::from([8.125, 0.0]));
1066
1067        check!(cell_list.half_extent == 8);
1068        let potential_neighbors: Vec<_> = cell_list
1069            .points_near_ball(&[9.125, 0.0].into(), 1.0)
1070            .collect();
1071        assert!(potential_neighbors.len() == 1);
1072        check!(potential_neighbors[0] == 2);
1073    }
1074
1075    #[rstest]
1076    #[case::d_2(PhantomData::<VecCell<usize, 2>>)]
1077    #[case::d_3(PhantomData::<VecCell<usize, 3>>)]
1078    fn test_points_near_ball<const D: usize>(
1079        #[case] _d: PhantomData<VecCell<usize, D>>,
1080        #[values(1.0, 0.5, 0.25)] nominal_search_radius: f64,
1081    ) {
1082        let mut rng = StdRng::seed_from_u64(0);
1083        let mut reference = Vec::new();
1084
1085        let mut cell_list = VecCell::builder()
1086            .nominal_search_radius(
1087                nominal_search_radius
1088                    .try_into()
1089                    .expect("hardcoded value should be positive"),
1090            )
1091            .maximum_search_radius(1.0)
1092            .build();
1093        let position_distribution = Ball {
1094            radius: 12.0.try_into().expect("hardcoded value should be positive"),
1095        };
1096
1097        let n = 2048;
1098
1099        for key in 0..n {
1100            let position: Cartesian<D> = position_distribution.sample(&mut rng);
1101
1102            cell_list.insert(key, position);
1103            reference.push(position);
1104        }
1105
1106        let mut n_neighbors = 0;
1107        for p_i in &reference {
1108            let potential_neighbors: Vec<_> = cell_list.points_near_ball(p_i, 1.0).collect();
1109
1110            for (j, p_j) in reference.iter().enumerate() {
1111                if p_i.distance(p_j) <= 1.0 {
1112                    check!(potential_neighbors.contains(&j));
1113                    n_neighbors += 1;
1114                }
1115            }
1116        }
1117        check!(n_neighbors >= n * 2);
1118    }
1119}
hoomd_spatial/vec_cell.rs

hoomd_spatial/
vec_cell.rs
