Skip to main content

hoomd_vector/
distribution.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//! Random distributions of vectors.
5use serde::{Deserialize, Serialize};
6use std::array;
7
8use super::{Cartesian, InnerProduct};
9use hoomd_utility::valid::PositiveReal;
10
11use rand::{Rng, RngExt, distr::Distribution};
12use rand_distr::StandardNormal;
13
14/// A uniform distribution of all points inside or on a sphere with radius `r`.
15///
16/// # Example
17///
18/// ```
19/// use hoomd_vector::{Cartesian, distribution::Ball};
20/// use rand::{Rng, SeedableRng, distr::Distribution, rngs::StdRng};
21///
22/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
23/// let mut rng = StdRng::seed_from_u64(1);
24///
25/// let r = 5.0;
26/// let ball = Ball {
27///     radius: r.try_into()?,
28/// };
29/// let v: Cartesian<3> = ball.sample(&mut rng);
30/// # Ok(())
31/// # }
32/// ```
33#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
34pub struct Ball {
35    /// The radius of the ball *(\[length\])*.
36    pub radius: PositiveReal,
37}
38
39impl<const N: usize> Distribution<Cartesian<N>> for Ball {
40    #[inline]
41    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Cartesian<N> {
42        let r = self.radius.get();
43
44        if N == 2 {
45            // Rejection sampling is fastest in 2D
46            loop {
47                let coordinates_01: [f64; N] = array::from_fn(|_| rng.random());
48
49                let v = Cartesian {
50                    coordinates: array::from_fn(|i| (coordinates_01[i] * 2.0 - 1.0) * r),
51                };
52
53                if v.norm_squared() < r * r {
54                    return v;
55                }
56            }
57        }
58
59        // Muller/Marsaglia 'normalized Gaussians' approach is faster n 3 and larger dimensions.
60        // The implementation is based on:
61        // https://extremelearning.com.au/how-to-generate-uniformly-random-points-on-n-spheres-and-n-balls/
62        let point = Cartesian {
63            coordinates: std::array::from_fn::<_, N, _>(|_| rng.sample(StandardNormal)),
64        };
65        match N {
66            3 => point * (r / point.norm() * rng.random::<f64>().cbrt()),
67            _ => point * (r / point.norm() * rng.random::<f64>().powf(1.0 / N as f64)),
68        }
69    }
70}
71
72#[cfg(test)]
73mod test {
74    use super::*;
75
76    /// Number of random vectors to sample.
77    const N: usize = 1024;
78    use approxim::assert_abs_diff_eq;
79    use rand::{SeedableRng, rngs::StdRng};
80    use rstest::*;
81
82    #[rstest(
83        r => [0.5, 1.0, 12.0])]
84    fn ball(r: f64) {
85        let ball = Ball {
86            radius: r
87                .try_into()
88                .expect("hard-coded constant should be positive"),
89        };
90        let mut rng = StdRng::seed_from_u64(1);
91
92        let mut total = Cartesian::default();
93        for _ in 0..N {
94            let v: Cartesian<3> = ball.sample(&mut rng);
95            assert!(v.norm_squared() < r * r);
96
97            total += v;
98        }
99
100        let average = total / N as f64;
101        for x in average.coordinates {
102            assert_abs_diff_eq!(x, 0.0, epsilon = r * 0.1);
103        }
104    }
105}