sobol_qmc/
lib.rs

1pub mod params;
2mod type_support;
3use core::{
4    fmt,
5    ops::{AddAssign, BitXorAssign},
6};
7use num_traits::{Bounded, One, PrimInt, Unsigned, Zero};
8pub use statrs;
9use statrs::distribution::Normal;
10use std::iter::repeat_n;
11
12/// A low-discrepancy Sobol sequence generator
13#[derive(Clone)]
14pub struct Sobol<T: SobolType, R: Render<T> = UnitRender> {
15    pub dims: usize,
16    pub resolution: usize,
17    dir_vals: Vec<Vec<T::IT>>,
18    previous: Vec<T::IT>,
19    render: R,
20    pub count: T::IT,
21    pub max_len: T::IT,
22}
23
24#[derive(Debug, Clone, Copy, thiserror::Error)]
25pub enum SobolError {
26    #[error(
27        "Sobol sequence supports a maximum of {max_dims} dimensions, but was configured for {dims}."
28    )]
29    MaxDim { dims: usize, max_dims: usize },
30    #[error("Render supports a {render_dims} dimensions, but Sobol was configured for {dims}.")]
31    RenderDim { dims: usize, render_dims: usize },
32}
33
34impl<T: SobolType> Sobol<T, UnitRender>
35where
36    UnitRender: Render<T>,
37{
38    /// Constructs a new sequence
39    pub fn new<P, Param: SobolParams<P>>(dims: usize, params: &Param) -> Result<Self, SobolError>
40    where
41        T::IT: LossyFrom<P>,
42    {
43        Self::new_with_resolution::<P, Param>(dims, params, None, UnitRender)
44    }
45}
46impl<T: SobolType, R: Render<T>> Sobol<T, R> {
47    /// Constructs a new sequence of given resolution. Resolution is the number of bits used in the
48    /// computation of the sequence and by default is the size of the underlying type. This
49    /// constructor is useful for reducing the number of cycles necessary to generate each point when the
50    /// length of the sequence is not expected to approach it's theoretically maximum (2^res).
51    pub fn new_with_resolution<P, Param: SobolParams<P>>(
52        dims: usize,
53        params: &Param,
54        resolution: Option<usize>,
55        render: R,
56    ) -> Result<Self, SobolError>
57    where
58        T::IT: LossyFrom<P>,
59    {
60        let res = resolution
61            .filter(|res| *res <= T::MAX_RESOLUTION)
62            .unwrap_or(T::MAX_RESOLUTION);
63        let max_dims = params.max_dims();
64        if dims > params.max_dims() {
65            return Err(SobolError::MaxDim { dims, max_dims });
66        }
67        if let Some(render_dims) = render.support_dims() {
68            if dims != render_dims {
69                return Err(SobolError::RenderDim { dims, render_dims });
70            }
71        }
72
73        let dir_values = Self::init_direction_vals::<P, Param>(dims, res, params);
74        // Transpose dir values for better cache locality
75        let dir_values = (0..dir_values[0].len())
76            .map(|i| dir_values.iter().map(|inner| inner[i]).collect::<Vec<_>>())
77            .collect();
78        Ok(Sobol {
79            dims,
80            resolution: res,
81            dir_vals: dir_values,
82            count: T::IT::zero(),
83            max_len: T::IT::max_value() >> (T::IT::BITS - res),
84            previous: Vec::with_capacity(dims),
85            render,
86        })
87    }
88
89    /// Initializes per-dimension direction values given sequence parameters
90    pub fn init_direction_vals<P, Param: SobolParams<P>>(
91        dims: usize,
92        resolution: usize,
93        params: &Param,
94    ) -> Vec<Vec<T::IT>>
95    where
96        T::IT: LossyFrom<P>,
97    {
98        let bits = T::IT::BITS;
99
100        (1..=dims)
101            .map(|dim| match dim {
102                1 => (1..=resolution)
103                    .map(|i| T::IT::one() << (bits - i))
104                    .collect(),
105                _ => {
106                    // Import the parameters needed to prepare this dimension's direction vector
107                    let p = params.get_dim(dim);
108                    let s = if resolution >= p.s() {
109                        p.s()
110                    } else {
111                        resolution
112                    };
113
114                    // Shift initial directions
115                    let mut dirs: Vec<T::IT> = vec![T::IT::zero(); resolution];
116                    for i in 1..=s {
117                        let m = T::IT::lossy_from(p.m(i - 1));
118                        dirs[i - 1] = m << (bits - i);
119                    }
120
121                    // Compute remaining directions
122                    for i in s + 1..=resolution {
123                        dirs[i - 1] = dirs[i - s - 1] ^ (dirs[i - s - 1] >> s);
124
125                        for k in 1..s {
126                            let a = T::IT::lossy_from(p.coefficient(s - k - 1));
127                            let dir = dirs[i - k - 1];
128                            dirs[i - 1] ^= a * dir;
129                        }
130                    }
131
132                    dirs
133                }
134            })
135            .collect()
136    }
137
138    /// Returns zero-based index of the rightmost binary zero. Used for the Gray code optimization
139    #[inline]
140    pub fn rightmost_zero(n: T::IT) -> usize {
141        (n ^ T::IT::max_value()).trailing_zeros() as usize
142    }
143
144    /// Update internal for each `next()`
145    #[inline]
146    pub fn update(&mut self) {
147        if self.previous.is_empty() {
148            self.previous.extend(repeat_n(T::IT::zero(), self.dims));
149        } else {
150            let a = self.count - T::IT::one();
151            let c = Self::rightmost_zero(a);
152            for (p, dir) in self.dir_vals[c].iter().zip(&mut self.previous) {
153                *dir = *p ^ *dir;
154            }
155        }
156        self.count += T::IT::one();
157    }
158    #[inline]
159    pub fn render_next(&self) -> impl Iterator<Item = T> {
160        self.previous
161            .iter()
162            .enumerate()
163            .map(|(dim, val)| self.render.render(dim, *val))
164    }
165}
166
167impl<T: SobolType, R: Render<T>> Iterator for Sobol<T, R> {
168    type Item = Vec<T>;
169    #[inline]
170    fn next(&mut self) -> Option<Self::Item> {
171        if self.count < self.max_len {
172            self.update();
173            Some(self.render_next().collect())
174        } else {
175            None
176        }
177    }
178    #[inline]
179    fn nth(&mut self, n: usize) -> Option<Self::Item> {
180        for _ in 0..n {
181            if self.count < self.max_len {
182                self.update();
183            } else {
184                break;
185            }
186        }
187        self.next()
188    }
189}
190
191#[derive(Debug, Clone, Copy)]
192pub struct UnitRender;
193#[derive(Debug, Clone, Copy)]
194pub struct GaussianRender(pub Normal);
195#[derive(Debug, Clone)]
196pub struct MultiDimGaussianRender(pub Vec<Normal>);
197pub trait Render<T: SobolType>: Clone {
198    /// Converts internal values to those expected by the user. This usually
199    /// involves casting and, for float values, scaling to the range [0,1).
200    fn render(&self, dim: usize, val: <T as SobolType>::IT) -> T;
201
202    fn support_dims(&self) -> Option<usize> {
203        None
204    }
205}
206
207/// The main type parameter for the `Sobol` iterator. This defines the concrete `InternalType`
208/// to be used internally, as well as other properties necessary for sequence generation.
209pub trait SobolType: Sized + fmt::Display {
210    /// The unsigned integer type used internally to compute sequence values.
211    type IT: InternalType;
212
213    /// The maximum number of bits this type can support. By default, this is
214    /// number of bits in the underlying `InternalType`, but it may be less
215    /// in some cases (e.g. floats are limited by the size of their significand).
216    const MAX_RESOLUTION: usize = Self::IT::BITS;
217}
218
219/// Sequences are computed internally using unsigned types with the following capabilities
220pub trait InternalType: PrimInt + Unsigned + One + Zero + AddAssign + BitXorAssign + Copy {
221    const BITS: usize;
222}
223
224/// Primitive polynomial parameters and initial direction values for all sequence dimensions
225pub trait SobolParams<P> {
226    type Dimension: ParamDimension<P>;
227    /// Parameters for a given dimension
228    fn get_dim(&self, dim: usize) -> &Self::Dimension;
229
230    /// Maximum number of dimensions supported by this instance
231    fn max_dims(&self) -> usize;
232}
233
234/// Primitive polynomial parameters and initial direction values for a single dimension
235pub trait ParamDimension<P> {
236    /// The one-based index of this dimension
237    fn d(&self) -> u16;
238
239    /// The degree of the primitive polynomial
240    fn s(&self) -> usize;
241
242    /// The binary coefficient for bit `i`, the zero-based index from the right
243    fn coefficient(&self, i: usize) -> P;
244
245    /// The initial direction value for bit `i`, the zero-based index from the right
246    fn m(&self, i: usize) -> P;
247}
248
249/// A more permissive `From` trait - suitable for cases where lossy casting is
250/// acceptable (i.e. truncation). This is used for casting parameter values to
251/// internal values.
252pub trait LossyFrom<T>: Sized {
253    fn lossy_from(_: T) -> Self;
254}
sobol_qmc/lib.rs

sobol_qmc/
lib.rs
