Skip to main content

hoomd_mc/
hypercuboid.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 Checkerboard for Hypercuboids
14
15use itertools::izip;
16use rand::{Rng, RngExt};
17use serde::{Deserialize, Serialize};
18use serde_with::serde_as;
19use std::array;
20
21use hoomd_geometry::shape::Hypercuboid;
22use hoomd_microstate::boundary::{Closed, Periodic};
23use hoomd_utility::valid::PositiveReal;
24use hoomd_vector::Cartesian;
25
26use crate::{Checkerboard, Cover};
27
28/// `2^N` color checkerboard with axis-aligned hypercuboidal cells.
29///
30/// A `HypercuboidCheckerboard` is comprised of n x m x ... axis aligned
31/// spaces. Each space has the same shape, but each axis might have a different
32/// edge length. Each axis may be periodic or not.
33///
34/// Along the non-periodic axes, the checkerboard overhangs the boundary so that
35/// the entire domain is always covered (for any origin shift up to 1 cell length).
36/// Along periodic axes, the checkerboard has exactly the same width as the domain.
37/// `HypercuboidCheckerboard` wraps points "outside" the checkerboard into the correct
38/// periodic space.
39///
40/// Obviously, `HypercuboidCheckerboard` is a suitable checkerboard for
41/// [`Hypercuboid`] boundary geometries. It can also be a good choice for
42/// other boundaries. For example: cylindrical boundaries (periodic in one
43/// direction) and closed boundaries of any shape. There may be many overhanging
44/// spaces (some completely outside the boundary) in these cases. However, the
45/// checkerboard is still valid and rayon's dynamic load balancing scheme should
46/// be able to handle the empty cells efficiently.
47#[serde_as]
48#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
49pub struct HypercuboidCheckerboard<const N: usize> {
50    /// Position of the 0,0,0 cell's lower left corner.
51    origin: Cartesian<N>,
52
53    /// Length of each axis aligned space edge.
54    #[serde_as(as = "[_; N]")]
55    space_width: [f64; N],
56
57    /// Number of spaces along each axis.
58    #[serde_as(as = "[_; N]")]
59    shape: [usize; N],
60
61    /// True when an axis is periodic.
62    #[serde_as(as = "[_; N]")]
63    periodic: [bool; N],
64
65    /// The set of all space indices, grouped by color.
66    space_indices_by_color: Vec<Vec<usize>>,
67}
68
69impl<const N: usize> Default for HypercuboidCheckerboard<N> {
70    /// Construct a default `HypercuboidCheckerboard`.
71    ///
72    /// The default is a 2x2x... non-periodic checkerboard with 1.0 x 1.0 x ... spaces.
73    #[inline]
74    fn default() -> Self {
75        Self {
76            origin: Cartesian::default(),
77            space_width: [1.0; N],
78            shape: [2; N],
79            periodic: [false; N],
80            space_indices_by_color: Self::construct_space_indices_by_color([2; N]),
81        }
82    }
83}
84
85impl<const N: usize> Checkerboard<Cartesian<N>> for HypercuboidCheckerboard<N> {
86    #[inline]
87    fn point_to_space_index(&self, point: &Cartesian<N>) -> Option<usize> {
88        let p = *point - self.origin;
89        let mut space_multi_index: [i64; N] =
90            array::from_fn(|i| (p.coordinates[i] / self.space_width[i]).floor() as i64);
91
92        for (index, shape, periodic) in izip!(&mut space_multi_index, self.shape, self.periodic) {
93            // The origin is in the lower left corner of the box and may be up
94            // to one space width to the left of simulation boundary. Therefore,
95            // negative indices are out of bounds (and should never been seen
96            // for wrapped inputs).
97            if *index < 0 {
98                return None;
99            }
100
101            if periodic {
102                // In periodic boundaries, the checkerboard spaces end before
103                // the right side. The space at the rightmost edge is identical
104                // to space 0 to make the checkerboard coloring commensurate
105                // with the periodic boundary conditions.
106                if *index as usize == shape {
107                    *index = 0;
108                }
109                if *index as usize > shape {
110                    return None;
111                }
112            } else if *index as usize >= shape {
113                // In non-periodic boundaries, the checkerboard extends one full
114                // space to the right of the simulation boundary so that when
115                // it is shifted to the left it will still cover the boundary.
116                // Therefore, any points outside the checkerboard shape are
117                // invalid.
118                return None;
119            }
120        }
121
122        Some(Self::multi_index_to_index(
123            array::from_fn(|i| space_multi_index[i] as usize),
124            self.shape,
125        ))
126    }
127
128    #[inline]
129    fn space_indices_by_color(&self) -> &[Vec<usize>] {
130        &self.space_indices_by_color
131    }
132
133    #[inline]
134    fn num_spaces(&self) -> usize {
135        self.shape.iter().product()
136    }
137}
138
139impl<const N: usize> HypercuboidCheckerboard<N> {
140    /// Collapse a multi-dimensional index to a single value in `[0, num_spaces]`.
141    #[inline]
142    fn multi_index_to_index(multi_index: [usize; N], shape: [usize; N]) -> usize {
143        let mut index: usize = 0;
144        let mut width = 1;
145
146        for i in (0..N).rev() {
147            index += multi_index[i] * width;
148            width *= shape[i];
149        }
150
151        index
152    }
153
154    /// Expand a single value in `[0, product(shape)]` back to a multi-dimensional index.
155    #[inline]
156    fn index_to_multi_index(mut index: usize, shape: [usize; N]) -> [usize; N] {
157        let mut multi_index = [0_usize; N];
158        for i in (0..N).rev() {
159            multi_index[i] = index % shape[i];
160            index /= shape[i];
161        }
162        multi_index
163    }
164
165    /// Compute the space width and checkerboard shape.
166    #[inline]
167    fn compute_dimensions(
168        edge_lengths: [PositiveReal; N],
169        interaction_range: PositiveReal,
170        periodic: [bool; N],
171    ) -> ([f64; N], [usize; N]) {
172        let mut shape_inside: [usize; N] =
173            array::from_fn(|i| (edge_lengths[i].get() / interaction_range.get()).floor() as usize);
174
175        assert!(
176            shape_inside.iter().all(|n| *n >= 2),
177            "body interaction range {interaction_range} is too large for the boundary dimensions {edge_lengths:?}"
178        );
179
180        for (width, periodic) in shape_inside.iter_mut().zip(periodic) {
181            // In periodic boundaries, the checkerboard must have an even number
182            // of spaces on a side and fit entirely within the given boundary.
183            // Spaces must be made larger to accommodate.
184            if periodic && !width.is_multiple_of(2) {
185                *width -= 1;
186            }
187
188            // In non-periodic boundaries, the checkerboard must have an even
189            // number of spaces on a side with exactly 1 full space outside
190            // the boundary. Therefore, there must be an odd number of spaces
191            // inside.
192            if !periodic && width.is_multiple_of(2) {
193                *width -= 1;
194            }
195        }
196
197        let space_width = array::from_fn(|i| edge_lengths[i].get() / shape_inside[i] as f64);
198        let shape = array::from_fn(|i| {
199            if periodic[i] {
200                shape_inside[i]
201            } else {
202                shape_inside[i] + 1
203            }
204        });
205
206        (space_width, shape)
207    }
208
209    /// Partition the space indices by color.
210    fn construct_space_indices_by_color(shape: [usize; N]) -> Vec<Vec<usize>> {
211        for width in shape {
212            assert!(width.is_multiple_of(2));
213        }
214        let num_colors = 2usize.pow(N as u32);
215
216        let half: [usize; N] = array::from_fn(|i| shape[i] / 2);
217        let total: usize = half.iter().product();
218
219        let base: Vec<usize> = (0..total)
220            .map(|idx| {
221                let half_multi = Self::index_to_multi_index(idx, half);
222                let multi_index: [usize; N] = array::from_fn(|i| 2 * half_multi[i]);
223                Self::multi_index_to_index(multi_index, shape)
224            })
225            .collect();
226
227        (0..num_colors)
228            .map(|color| {
229                let offset = Self::index_to_multi_index(color, [2; N]);
230                let delta = Self::multi_index_to_index(offset, shape);
231                base.iter().map(|&b| b + delta).collect()
232            })
233            .collect()
234    }
235
236    /// Construct a checkerboard with a given origin (for testing).
237    #[cfg(test)]
238    fn with_fixed_origin(
239        interaction_range: PositiveReal,
240        edge_lengths: [PositiveReal; N],
241        periodic: [bool; N],
242    ) -> Self {
243        let (space_width, shape) =
244            Self::compute_dimensions(edge_lengths, interaction_range, periodic);
245
246        Self {
247            space_width,
248            shape,
249            origin: Cartesian::from(array::from_fn(|i| -edge_lengths[i].get() / 2.0)),
250            space_indices_by_color: Self::construct_space_indices_by_color(shape),
251            periodic,
252        }
253    }
254
255    /// Construct a new `HypercuboidCheckerboard`.
256    ///
257    /// Set `interaction_range` to the largest distance between any two
258    /// interacting bodies. `new` will construct a checkerboard that covers
259    /// the range `[-edge_lengths[i]/2.0, edge_lengths[i]/2.0)` (respecting
260    /// `periodic[i]`) for each dimension `i`.
261    #[inline]
262    pub fn new<R: Rng + ?Sized>(
263        rng: &mut R,
264        interaction_range: PositiveReal,
265        edge_lengths: [PositiveReal; N],
266        periodic: [bool; N],
267    ) -> Self {
268        let (space_width, shape) =
269            Self::compute_dimensions(edge_lengths, interaction_range, periodic);
270
271        let offset: [f64; N] = array::from_fn(|_| rng.random());
272
273        let origin = Cartesian {
274            coordinates: array::from_fn(|i| {
275                -edge_lengths[i].get() / 2.0 - offset[i] * space_width[i]
276            }),
277        };
278
279        Self {
280            space_width,
281            shape,
282            origin,
283            space_indices_by_color: Self::construct_space_indices_by_color(shape),
284            periodic,
285        }
286    }
287
288    /// Update an existing checkerboard.
289    ///
290    /// `update` performs the same steps as `new`, but modifies `self` in place.
291    /// Prefer `update` when possible, as it can reuse the space index partitioning
292    /// when the shape doesn't change.
293    #[inline]
294    pub fn update<R: Rng + ?Sized>(
295        &mut self,
296        rng: &mut R,
297        interaction_range: PositiveReal,
298        edge_lengths: [PositiveReal; N],
299        periodic: [bool; N],
300    ) {
301        let (space_width, shape) =
302            Self::compute_dimensions(edge_lengths, interaction_range, periodic);
303
304        let offset: [f64; N] = array::from_fn(|_| rng.random());
305
306        let origin = Cartesian {
307            coordinates: array::from_fn(|i| {
308                -edge_lengths[i].get() / 2.0 - offset[i] * space_width[i]
309            }),
310        };
311
312        if shape != self.shape || self.space_indices_by_color.is_empty() {
313            self.space_indices_by_color = Self::construct_space_indices_by_color(shape);
314            self.shape = shape;
315        }
316
317        self.space_width = space_width;
318        self.origin = origin;
319        self.periodic = periodic;
320    }
321}
322
323impl<const N: usize> Cover<Cartesian<N>> for Closed<Hypercuboid<N>> {
324    type Checkerboard = HypercuboidCheckerboard<N>;
325
326    #[inline]
327    fn cover<R: Rng + ?Sized>(
328        &self,
329        rng: &mut R,
330        interaction_range: PositiveReal,
331    ) -> Self::Checkerboard {
332        HypercuboidCheckerboard::new(rng, interaction_range, self.0.edge_lengths, [false; N])
333    }
334
335    #[inline]
336    fn cover_into<R: Rng + ?Sized>(
337        &self,
338        checkerboard: &mut Self::Checkerboard,
339        rng: &mut R,
340        interaction_range: PositiveReal,
341    ) {
342        checkerboard.update(rng, interaction_range, self.0.edge_lengths, [false; N]);
343    }
344}
345
346impl<const N: usize> Cover<Cartesian<N>> for Periodic<Hypercuboid<N>> {
347    type Checkerboard = HypercuboidCheckerboard<N>;
348
349    #[inline]
350    fn cover<R: Rng + ?Sized>(
351        &self,
352        rng: &mut R,
353        interaction_range: PositiveReal,
354    ) -> Self::Checkerboard {
355        HypercuboidCheckerboard::new(rng, interaction_range, self.shape().edge_lengths, [true; N])
356    }
357
358    #[inline]
359    fn cover_into<R: Rng + ?Sized>(
360        &self,
361        checkerboard: &mut Self::Checkerboard,
362        rng: &mut R,
363        interaction_range: PositiveReal,
364    ) {
365        checkerboard.update(rng, interaction_range, self.shape().edge_lengths, [true; N]);
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use assert2::{assert, check};
372    use rand::{SeedableRng, rngs::StdRng};
373    use rstest::*;
374
375    use super::*;
376
377    #[test]
378    fn test_multi_index_to_index_2d() {
379        let shape = [12, 4];
380        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([0, 0], shape) == 0);
381        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([0, 1], shape) == 1);
382        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([0, 2], shape) == 2);
383        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([0, 3], shape) == 3);
384        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([1, 0], shape) == 4);
385        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([1, 1], shape) == 5);
386        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([1, 2], shape) == 6);
387        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([1, 3], shape) == 7);
388        check!(HypercuboidCheckerboard::<2>::multi_index_to_index([11, 3], shape) == 47);
389    }
390
391    #[test]
392    fn test_multi_index_to_index_3d() {
393        let shape = [12, 4, 6];
394        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 0, 0], shape) == 0);
395        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 0, 1], shape) == 1);
396        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 0, 2], shape) == 2);
397        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 0, 3], shape) == 3);
398        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 0, 4], shape) == 4);
399        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 0, 5], shape) == 5);
400
401        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 1, 0], shape) == 6);
402        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 2, 0], shape) == 12);
403        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([0, 3, 0], shape) == 18);
404
405        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([1, 0, 0], shape) == 24);
406        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([2, 0, 0], shape) == 48);
407        check!(HypercuboidCheckerboard::<3>::multi_index_to_index([3, 3, 2], shape) == 92);
408    }
409
410    #[test]
411    fn test_multi_index_to_index_4d() {
412        let shape = [12, 4, 6, 8];
413        (0..shape[3]).for_each(|i| {
414            check!(HypercuboidCheckerboard::<4>::multi_index_to_index([0, 0, 0, i], shape) == i);
415        });
416
417        (0..shape[2]).for_each(|i| {
418            check!(
419                HypercuboidCheckerboard::<4>::multi_index_to_index([0, 0, i, 0], shape)
420                    == i * shape[3]
421            );
422        });
423
424        (0..shape[1]).for_each(|i| {
425            check!(
426                HypercuboidCheckerboard::<4>::multi_index_to_index([0, i, 0, 0], shape)
427                    == (i * shape[2]) * shape[3]
428            );
429        });
430
431        check!(HypercuboidCheckerboard::<4>::multi_index_to_index([11, 3, 5, 7], shape) == 2303);
432        check!(
433            HypercuboidCheckerboard::<4>::multi_index_to_index([9, 0, 0, 0], shape)
434                == (9 * 4 * 6) * 8
435        );
436
437        // Same as 3D, multiplied by shape[3]
438        check!(HypercuboidCheckerboard::<4>::multi_index_to_index([1, 0, 0, 0], shape) == 24 * 8);
439        check!(HypercuboidCheckerboard::<4>::multi_index_to_index([2, 0, 0, 0], shape) == 48 * 8);
440        check!(HypercuboidCheckerboard::<4>::multi_index_to_index([3, 3, 2, 0], shape) == 92 * 8);
441    }
442
443    #[test]
444    fn test_compute_dimensions_exact() -> anyhow::Result<()> {
445        let (space_width, shape) = HypercuboidCheckerboard::<2>::compute_dimensions(
446            [16.0.try_into()?, 24.0.try_into()?],
447            2.0.try_into()?,
448            [true; 2],
449        );
450        check!(space_width == [2.0, 2.0]);
451        check!(shape == [8, 12]);
452
453        let (space_width, shape) = HypercuboidCheckerboard::<2>::compute_dimensions(
454            [14.0.try_into()?, 22.0.try_into()?],
455            2.0.try_into()?,
456            [false; 2],
457        );
458        check!(space_width == [2.0, 2.0]);
459        check!(shape == [8, 12]);
460
461        let (space_width, shape) = HypercuboidCheckerboard::<2>::compute_dimensions(
462            [16.0.try_into()?, 22.0.try_into()?],
463            2.0.try_into()?,
464            [true, false],
465        );
466        check!(space_width == [2.0, 2.0]);
467        check!(shape == [8, 12]);
468
469        Ok(())
470    }
471
472    #[test]
473    fn test_compute_dimensions_expand() -> anyhow::Result<()> {
474        let (space_width, shape) = HypercuboidCheckerboard::<2>::compute_dimensions(
475            [15.0.try_into()?, 23.0.try_into()?],
476            1.0.try_into()?,
477            [true; 2],
478        );
479        check!(space_width == [15.0 / 14.0, 23.0 / 22.0]);
480        check!(shape == [14, 22]);
481
482        let (space_width, shape) = HypercuboidCheckerboard::<2>::compute_dimensions(
483            [14.0.try_into()?, 22.0.try_into()?],
484            1.0.try_into()?,
485            [false; 2],
486        );
487        check!(space_width == [14.0 / 13.0, 22.0 / 21.0]);
488        check!(shape == [14, 22]);
489
490        Ok(())
491    }
492
493    #[test]
494    fn test_construct_space_indices_by_color_2d() {
495        let space_indices_by_color =
496            HypercuboidCheckerboard::<2>::construct_space_indices_by_color([6, 4]);
497
498        assert!(space_indices_by_color.len() == 4);
499        assert!(space_indices_by_color[0].len() == 6);
500        check!(space_indices_by_color[0][0] == 0);
501        check!(space_indices_by_color[0][1] == 2);
502        check!(space_indices_by_color[0][2] == 8);
503        check!(space_indices_by_color[0][3] == 10);
504        check!(space_indices_by_color[0][4] == 16);
505        check!(space_indices_by_color[0][5] == 18);
506
507        assert!(space_indices_by_color[1].len() == 6);
508        check!(space_indices_by_color[1][0] == 1);
509        check!(space_indices_by_color[1][1] == 3);
510        check!(space_indices_by_color[1][2] == 9);
511        check!(space_indices_by_color[1][3] == 11);
512        check!(space_indices_by_color[1][4] == 17);
513        check!(space_indices_by_color[1][5] == 19);
514
515        assert!(space_indices_by_color[2].len() == 6);
516        check!(space_indices_by_color[2][0] == 4);
517        check!(space_indices_by_color[2][1] == 6);
518        check!(space_indices_by_color[2][2] == 12);
519        check!(space_indices_by_color[2][3] == 14);
520        check!(space_indices_by_color[2][4] == 20);
521        check!(space_indices_by_color[2][5] == 22);
522
523        assert!(space_indices_by_color[3].len() == 6);
524        check!(space_indices_by_color[3][0] == 5);
525        check!(space_indices_by_color[3][1] == 7);
526        check!(space_indices_by_color[3][2] == 13);
527        check!(space_indices_by_color[3][3] == 15);
528        check!(space_indices_by_color[3][4] == 21);
529        check!(space_indices_by_color[3][5] == 23);
530    }
531
532    #[test]
533    fn test_construct_space_indices_by_color_3d_222() {
534        let space_indices_by_color =
535            HypercuboidCheckerboard::<3>::construct_space_indices_by_color([2, 2, 2]);
536
537        assert!(space_indices_by_color.len() == 8);
538        assert!(space_indices_by_color[0].len() == 1);
539        check!(space_indices_by_color[0][0] == 0);
540
541        assert!(space_indices_by_color[1].len() == 1);
542        check!(space_indices_by_color[1][0] == 1);
543
544        assert!(space_indices_by_color[2].len() == 1);
545        check!(space_indices_by_color[2][0] == 2);
546
547        assert!(space_indices_by_color[3].len() == 1);
548        check!(space_indices_by_color[3][0] == 3);
549
550        assert!(space_indices_by_color[4].len() == 1);
551        check!(space_indices_by_color[4][0] == 4);
552
553        assert!(space_indices_by_color[5].len() == 1);
554        check!(space_indices_by_color[5][0] == 5);
555
556        assert!(space_indices_by_color[6].len() == 1);
557        check!(space_indices_by_color[6][0] == 6);
558
559        assert!(space_indices_by_color[7].len() == 1);
560        check!(space_indices_by_color[7][0] == 7);
561    }
562    #[test]
563    fn test_construct_space_indices_by_color_3d_general() {
564        let space_indices_by_color =
565            HypercuboidCheckerboard::<3>::construct_space_indices_by_color([2, 6, 4]);
566
567        assert!(space_indices_by_color.len() == 8);
568        assert!(space_indices_by_color[0].len() == 6);
569        check!(space_indices_by_color[0][0] == 0);
570        check!(space_indices_by_color[0][1] == 2);
571        check!(space_indices_by_color[0][2] == 8);
572        check!(space_indices_by_color[0][3] == 10);
573        check!(space_indices_by_color[0][4] == 16);
574        check!(space_indices_by_color[0][5] == 18);
575
576        assert!(space_indices_by_color[1].len() == 6);
577        check!(space_indices_by_color[1][0] == 1);
578        check!(space_indices_by_color[1][1] == 3);
579        check!(space_indices_by_color[1][2] == 9);
580        check!(space_indices_by_color[1][3] == 11);
581        check!(space_indices_by_color[1][4] == 17);
582        check!(space_indices_by_color[1][5] == 19);
583
584        assert!(space_indices_by_color[2].len() == 6);
585        check!(space_indices_by_color[2][0] == 4);
586        check!(space_indices_by_color[2][1] == 6);
587        check!(space_indices_by_color[2][2] == 12);
588        check!(space_indices_by_color[2][3] == 14);
589        check!(space_indices_by_color[2][4] == 20);
590        check!(space_indices_by_color[2][5] == 22);
591
592        assert!(space_indices_by_color[3].len() == 6);
593        check!(space_indices_by_color[3][0] == 5);
594        check!(space_indices_by_color[3][1] == 7);
595        check!(space_indices_by_color[3][2] == 13);
596        check!(space_indices_by_color[3][3] == 15);
597        check!(space_indices_by_color[3][4] == 21);
598        check!(space_indices_by_color[3][5] == 23);
599
600        assert!(space_indices_by_color[4].len() == 6);
601        check!(space_indices_by_color[4][0] == 24);
602        check!(space_indices_by_color[4][1] == 24 + 2);
603        check!(space_indices_by_color[4][2] == 24 + 8);
604        check!(space_indices_by_color[4][3] == 24 + 10);
605        check!(space_indices_by_color[4][4] == 24 + 16);
606        check!(space_indices_by_color[4][5] == 24 + 18);
607
608        assert!(space_indices_by_color[5].len() == 6);
609        check!(space_indices_by_color[5][0] == 24 + 1);
610        check!(space_indices_by_color[5][1] == 24 + 3);
611        check!(space_indices_by_color[5][2] == 24 + 9);
612        check!(space_indices_by_color[5][3] == 24 + 11);
613        check!(space_indices_by_color[5][4] == 24 + 17);
614        check!(space_indices_by_color[5][5] == 24 + 19);
615
616        assert!(space_indices_by_color[6].len() == 6);
617        check!(space_indices_by_color[6][0] == 24 + 4);
618        check!(space_indices_by_color[6][1] == 24 + 6);
619        check!(space_indices_by_color[6][2] == 24 + 12);
620        check!(space_indices_by_color[6][3] == 24 + 14);
621        check!(space_indices_by_color[6][4] == 24 + 20);
622        check!(space_indices_by_color[6][5] == 24 + 22);
623
624        assert!(space_indices_by_color[7].len() == 6);
625        check!(space_indices_by_color[7][0] == 24 + 5);
626        check!(space_indices_by_color[7][1] == 24 + 7);
627        check!(space_indices_by_color[7][2] == 24 + 13);
628        check!(space_indices_by_color[7][3] == 24 + 15);
629        check!(space_indices_by_color[7][4] == 24 + 21);
630        check!(space_indices_by_color[7][5] == 24 + 23);
631    }
632
633    #[test]
634    fn test_construct_space_indices_by_color_4d_general() {
635        let space_indices_by_color =
636            HypercuboidCheckerboard::<4>::construct_space_indices_by_color([2, 4, 2, 4]);
637
638        assert!(space_indices_by_color.len() == 16);
639
640        // Colors 0-7: axis 0 offset is 0.
641        assert!(space_indices_by_color[0].len() == 4);
642        check!(space_indices_by_color[0][0] == 0);
643        check!(space_indices_by_color[0][1] == 2);
644        check!(space_indices_by_color[0][2] == 16);
645        check!(space_indices_by_color[0][3] == 18);
646
647        assert!(space_indices_by_color[1].len() == 4);
648        check!(space_indices_by_color[1][0] == 1);
649        check!(space_indices_by_color[1][1] == 3);
650        check!(space_indices_by_color[1][2] == 17);
651        check!(space_indices_by_color[1][3] == 19);
652
653        assert!(space_indices_by_color[2].len() == 4);
654        check!(space_indices_by_color[2][0] == 4);
655        check!(space_indices_by_color[2][1] == 6);
656        check!(space_indices_by_color[2][2] == 20);
657        check!(space_indices_by_color[2][3] == 22);
658
659        assert!(space_indices_by_color[3].len() == 4);
660        check!(space_indices_by_color[3][0] == 5);
661        check!(space_indices_by_color[3][1] == 7);
662        check!(space_indices_by_color[3][2] == 21);
663        check!(space_indices_by_color[3][3] == 23);
664
665        assert!(space_indices_by_color[4].len() == 4);
666        check!(space_indices_by_color[4][0] == 8);
667        check!(space_indices_by_color[4][1] == 10);
668        check!(space_indices_by_color[4][2] == 24);
669        check!(space_indices_by_color[4][3] == 26);
670
671        assert!(space_indices_by_color[5].len() == 4);
672        check!(space_indices_by_color[5][0] == 9);
673        check!(space_indices_by_color[5][1] == 11);
674        check!(space_indices_by_color[5][2] == 25);
675        check!(space_indices_by_color[5][3] == 27);
676
677        assert!(space_indices_by_color[6].len() == 4);
678        check!(space_indices_by_color[6][0] == 12);
679        check!(space_indices_by_color[6][1] == 14);
680        check!(space_indices_by_color[6][2] == 28);
681        check!(space_indices_by_color[6][3] == 30);
682
683        assert!(space_indices_by_color[7].len() == 4);
684        check!(space_indices_by_color[7][0] == 13);
685        check!(space_indices_by_color[7][1] == 15);
686        check!(space_indices_by_color[7][2] == 29);
687        check!(space_indices_by_color[7][3] == 31);
688
689        // Colors 8-15: axis 0 offset is 1 (stride 32 = 4*2*4).
690        assert!(space_indices_by_color[8].len() == 4);
691        check!(space_indices_by_color[8][0] == 32);
692        check!(space_indices_by_color[8][1] == 32 + 2);
693        check!(space_indices_by_color[8][2] == 32 + 16);
694        check!(space_indices_by_color[8][3] == 32 + 18);
695
696        assert!(space_indices_by_color[9].len() == 4);
697        check!(space_indices_by_color[9][0] == 32 + 1);
698        check!(space_indices_by_color[9][1] == 32 + 3);
699        check!(space_indices_by_color[9][2] == 32 + 17);
700        check!(space_indices_by_color[9][3] == 32 + 19);
701
702        assert!(space_indices_by_color[10].len() == 4);
703        check!(space_indices_by_color[10][0] == 32 + 4);
704        check!(space_indices_by_color[10][1] == 32 + 6);
705        check!(space_indices_by_color[10][2] == 32 + 20);
706        check!(space_indices_by_color[10][3] == 32 + 22);
707
708        assert!(space_indices_by_color[11].len() == 4);
709        check!(space_indices_by_color[11][0] == 32 + 5);
710        check!(space_indices_by_color[11][1] == 32 + 7);
711        check!(space_indices_by_color[11][2] == 32 + 21);
712        check!(space_indices_by_color[11][3] == 32 + 23);
713
714        assert!(space_indices_by_color[12].len() == 4);
715        check!(space_indices_by_color[12][0] == 32 + 8);
716        check!(space_indices_by_color[12][1] == 32 + 10);
717        check!(space_indices_by_color[12][2] == 32 + 24);
718        check!(space_indices_by_color[12][3] == 32 + 26);
719
720        assert!(space_indices_by_color[13].len() == 4);
721        check!(space_indices_by_color[13][0] == 32 + 9);
722        check!(space_indices_by_color[13][1] == 32 + 11);
723        check!(space_indices_by_color[13][2] == 32 + 25);
724        check!(space_indices_by_color[13][3] == 32 + 27);
725
726        assert!(space_indices_by_color[14].len() == 4);
727        check!(space_indices_by_color[14][0] == 32 + 12);
728        check!(space_indices_by_color[14][1] == 32 + 14);
729        check!(space_indices_by_color[14][2] == 32 + 28);
730        check!(space_indices_by_color[14][3] == 32 + 30);
731
732        assert!(space_indices_by_color[15].len() == 4);
733        check!(space_indices_by_color[15][0] == 32 + 13);
734        check!(space_indices_by_color[15][1] == 32 + 15);
735        check!(space_indices_by_color[15][2] == 32 + 29);
736        check!(space_indices_by_color[15][3] == 32 + 31);
737    }
738
739    #[test]
740    fn test_point_to_space_index_periodic() -> anyhow::Result<()> {
741        let checkerboard = HypercuboidCheckerboard::with_fixed_origin(
742            2.0.try_into()?,
743            [16.0.try_into()?, 24.0.try_into()?],
744            [true; 2],
745        );
746        check!(checkerboard.space_width == [2.0, 2.0]);
747        check!(checkerboard.shape == [8, 12]);
748
749        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, -12.1])) == None);
750        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, -12.0])) == Some(0));
751        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, -10.0])) == Some(1));
752        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, -8.0])) == Some(2));
753        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, -6.0])) == Some(3));
754        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, 11.9])) == Some(11));
755        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, 12.0])) == Some(0));
756
757        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.1, 0.0])) == None);
758        check!(checkerboard.point_to_space_index(&Cartesian::from([-8.0, 0.0])) == Some(6));
759        check!(checkerboard.point_to_space_index(&Cartesian::from([-6.0, 0.0])) == Some(18));
760        check!(checkerboard.point_to_space_index(&Cartesian::from([-4.0, 0.0])) == Some(30));
761        check!(checkerboard.point_to_space_index(&Cartesian::from([-2.0, 0.0])) == Some(42));
762        check!(checkerboard.point_to_space_index(&Cartesian::from([7.9, 0.0])) == Some(90));
763        check!(checkerboard.point_to_space_index(&Cartesian::from([8.0, 0.0])) == Some(6));
764
765        check!(checkerboard.point_to_space_index(&Cartesian::from([7.9, 11.9])) == Some(95));
766
767        Ok(())
768    }
769
770    #[test]
771    fn test_point_to_space_index_nonperiodic() -> anyhow::Result<()> {
772        let checkerboard = HypercuboidCheckerboard::with_fixed_origin(
773            2.0.try_into()?,
774            [14.0.try_into()?, 22.0.try_into()?],
775            [false; 2],
776        );
777        check!(checkerboard.space_width == [2.0, 2.0]);
778        check!(checkerboard.shape == [8, 12]);
779
780        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, -11.1])) == None);
781        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, -11.0])) == Some(0));
782        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, -9.0])) == Some(1));
783        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, -7.0])) == Some(2));
784        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, -5.0])) == Some(3));
785        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, 10.0])) == Some(10));
786        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, 11.0])) == Some(11));
787        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, 12.9])) == Some(11));
788        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, 13.0])) == None);
789
790        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.1, 1.0])) == None);
791        check!(checkerboard.point_to_space_index(&Cartesian::from([-7.0, 1.0])) == Some(6));
792        check!(checkerboard.point_to_space_index(&Cartesian::from([-5.0, 1.0])) == Some(18));
793        check!(checkerboard.point_to_space_index(&Cartesian::from([-3.0, 1.0])) == Some(30));
794        check!(checkerboard.point_to_space_index(&Cartesian::from([-1.0, 1.0])) == Some(42));
795        check!(checkerboard.point_to_space_index(&Cartesian::from([7.0, 1.0])) == Some(90));
796        check!(checkerboard.point_to_space_index(&Cartesian::from([9.0, 1.0])) == None);
797
798        check!(checkerboard.point_to_space_index(&Cartesian::from([8.9, 12.9])) == Some(95));
799
800        Ok(())
801    }
802
803    #[rstest]
804    fn test_all_points_inside(
805        #[values(true, false)] periodic: bool,
806        #[values(1, 2, 3, 4, 5, 6, 7)] seed: u64,
807    ) -> anyhow::Result<()> {
808        const N_SAMPLES: usize = 512;
809
810        let interaction_range = 1.5.try_into()?;
811        let edge_lengths: [PositiveReal; 2] = [10.0.try_into()?, 7.0.try_into()?];
812        let periodic = [periodic; 2];
813
814        let lower_left =
815            Cartesian::from([-edge_lengths[0].get() / 2.0, -edge_lengths[1].get() / 2.0]);
816        let upper_right = Cartesian::from([
817            edge_lengths[0].get() / 2.0 - 0.1,
818            edge_lengths[1].get() / 2.0 - 0.1,
819        ]);
820
821        let mut rng = StdRng::seed_from_u64(seed);
822
823        for _ in 0..N_SAMPLES {
824            let checkerboard =
825                HypercuboidCheckerboard::new(&mut rng, interaction_range, edge_lengths, periodic);
826
827            let boundary = Hypercuboid { edge_lengths };
828
829            for _ in 0..N_SAMPLES {
830                let v: Cartesian<2> = rng.sample(&boundary);
831
832                check!(checkerboard.point_to_space_index(&v).is_some());
833                check!(checkerboard.point_to_space_index(&lower_left) == Some(0));
834                check!(checkerboard.point_to_space_index(&upper_right).is_some());
835            }
836        }
837
838        Ok(())
839    }
840}