lattice_qcd_rs/statistics/
distribution.rs

1//! Some statistical distribution used by other part of the library
2
3use std::ops::{Add, Div, Mul, Neg, Sub};
4
5use num_traits::{Float, FloatConst, One, Zero};
6use rand::distributions::Uniform;
7use rand_distr::Distribution;
8#[cfg(feature = "serde-serialize")]
9use serde::{Deserialize, Serialize};
10
11use super::super::{su2, su3, CMatrix2, CMatrix3, Real};
12
13/// Distribution given by `x^2 e^{- 2 a x^2}`, `x >= 0` where `x` is the random variable and `a` a parameter of the distribution.
14///
15/// # Example
16/// ```
17/// use lattice_qcd_rs::error::ImplementationError;
18/// use lattice_qcd_rs::statistics::ModifiedNormal;
19/// use rand::{Rng, SeedableRng};
20///
21/// # fn main() -> Result<(), ImplementationError> {
22/// let mut rng = rand::rngs::StdRng::seed_from_u64(0);
23/// let mn = ModifiedNormal::new(0.5_f64).ok_or(ImplementationError::OptionWithUnexpectedNone)?;
24/// let r_number = rng.sample(&mn);
25/// #
26/// # Ok(())
27/// # }
28/// ```
29#[derive(Clone, Debug, Copy, PartialEq, Hash, Eq)]
30#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
31pub struct ModifiedNormal<T>
32where
33    T: One
34        + Div<T, Output = T>
35        + Mul<T, Output = T>
36        + Add<T, Output = T>
37        + Neg<Output = T>
38        + Float
39        + FloatConst
40        + PartialOrd,
41    rand::distributions::OpenClosed01: Distribution<T>,
42{
43    param_exp: T,
44}
45
46impl<T> ModifiedNormal<T>
47where
48    T: One
49        + Div<T, Output = T>
50        + Mul<T, Output = T>
51        + Add<T, Output = T>
52        + Neg<Output = T>
53        + Float
54        + FloatConst
55        + Zero
56        + PartialOrd,
57    rand::distributions::OpenClosed01: Distribution<T>,
58{
59    getter_copy!(
60        /// Returns the parameter `a`.
61        pub const,
62        param_exp,
63        T
64    );
65
66    /// Create the distribution. `param_exp` should be strictly greater than 0 an be finite and a number.
67    /// Otherwise return [`None`].
68    pub fn new(param_exp: T) -> Option<Self> {
69        if param_exp.le(&T::zero()) || param_exp.is_infinite() || param_exp.is_nan() {
70            return None;
71        }
72        Some(Self { param_exp })
73    }
74}
75
76impl<T> Distribution<T> for ModifiedNormal<T>
77where
78    T: One
79        + Div<T, Output = T>
80        + Mul<T, Output = T>
81        + Add<T, Output = T>
82        + Neg<Output = T>
83        + Float
84        + FloatConst
85        + rand_distr::uniform::SampleUniform
86        + PartialOrd,
87    rand::distributions::OpenClosed01: Distribution<T>,
88{
89    fn sample<R>(&self, rng: &mut R) -> T
90    where
91        R: rand::Rng + ?Sized,
92    {
93        let mut r = [T::one(); 3];
94        for element in r.iter_mut() {
95            *element = rng.sample(rand::distributions::OpenClosed01);
96        }
97        let two = T::one() + T::one();
98        (-(r[0].ln() + (two * T::PI() * r[1]).cos().powi(2) * r[2].ln()) / (two * self.param_exp()))
99            .sqrt()
100    }
101}
102
103impl<T> std::fmt::Display for ModifiedNormal<T>
104where
105    T: One
106        + Div<T, Output = T>
107        + Mul<T, Output = T>
108        + Add<T, Output = T>
109        + Neg<Output = T>
110        + Float
111        + FloatConst
112        + PartialOrd
113        + std::fmt::Display,
114    rand::distributions::OpenClosed01: Distribution<T>,
115{
116    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        write!(
118            f,
119            "modified normal distribution with parameter {}",
120            self.param_exp()
121        )
122    }
123}
124
125/// Distribution for the Heat Bath methods with the parameter `param_exp = beta * sqrt(det(A))`.
126///
127/// With distribution `dP(X) = 1/(2 \pi^2) d \cos(\theta) d\phi dx_0 \sqrt(1-x_0^2) e^{param_exp x_0}`.
128///
129/// # Example
130/// ```
131/// use lattice_qcd_rs::error::ImplementationError;
132/// use lattice_qcd_rs::statistics::HeatBathDistribution;
133/// use nalgebra::{Complex, Matrix2};
134/// use rand::{Rng, SeedableRng};
135///
136/// # fn main() -> Result<(), ImplementationError> {
137/// let mut rng = rand::rngs::StdRng::seed_from_u64(0);
138/// let heat_bath =
139///     HeatBathDistribution::new(0.5_f64).ok_or(ImplementationError::OptionWithUnexpectedNone)?;
140/// let r_matrix: Matrix2<Complex<f64>> = rng.sample(&heat_bath);
141/// #
142/// # Ok(())
143/// # }
144/// ```
145#[derive(Clone, Debug, Copy, PartialEq, Hash, Eq)]
146#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
147pub struct HeatBathDistribution<T>
148where
149    T: One
150        + Div<T, Output = T>
151        + Mul<T, Output = T>
152        + Add<T, Output = T>
153        + Sub<T, Output = T>
154        + Neg<Output = T>
155        + Float
156        + FloatConst
157        + Zero
158        + rand_distr::uniform::SampleUniform
159        + PartialOrd,
160    rand::distributions::OpenClosed01: Distribution<T>,
161    Uniform<T>: Distribution<T>,
162{
163    param_exp: T,
164}
165
166impl<T> HeatBathDistribution<T>
167where
168    T: One
169        + Div<T, Output = T>
170        + Mul<T, Output = T>
171        + Add<T, Output = T>
172        + Sub<T, Output = T>
173        + Neg<Output = T>
174        + Float
175        + FloatConst
176        + Zero
177        + rand_distr::uniform::SampleUniform
178        + PartialOrd,
179    rand::distributions::OpenClosed01: Distribution<T>,
180    Uniform<T>: Distribution<T>,
181{
182    getter_copy!(
183        /// Returns the parameter `param_exp`.
184        pub const,
185        param_exp,
186        T
187    );
188
189    /// Create the distribution. `param_exp` should be strictly greater than 0 an be finite and a number.
190    /// Otherwise return [`None`].
191    pub fn new(param_exp: T) -> Option<Self> {
192        if param_exp.le(&T::zero()) || param_exp.is_infinite() || param_exp.is_nan() {
193            return None;
194        }
195        Some(Self { param_exp })
196    }
197}
198
199impl Distribution<CMatrix2> for HeatBathDistribution<f64> {
200    fn sample<R>(&self, rng: &mut R) -> CMatrix2
201    where
202        R: rand::Rng + ?Sized,
203    {
204        // TODO make a function to reduce copy of code with su2::get_random_su2_close_to_unity
205
206        let distr_norm = HeatBathDistributionNorm::new(self.param_exp()).expect("unreachable");
207        // unreachable because self.param_exp() > 0 which Create the distribution
208        let x0: f64 = rng.sample(&distr_norm);
209        let uniform = Uniform::new(-1_f64, 1_f64);
210        let mut x_unorm = na::Vector3::from_fn(|_, _| rng.sample(&uniform));
211        while x_unorm.norm() <= f64::EPSILON {
212            x_unorm = na::Vector3::from_fn(|_, _| rng.sample(&uniform));
213        }
214        let x =
215            x_unorm.try_normalize(f64::EPSILON).expect("unreachable") * (1_f64 - x0 * x0).sqrt();
216        // unreachable because the while loop above guarantee that the norm is bigger than [`f64::EPSILON`]
217        su2::complex_matrix_from_vec(x0, x)
218    }
219}
220
221impl<T> std::fmt::Display for HeatBathDistribution<T>
222where
223    T: One
224        + Div<T, Output = T>
225        + Mul<T, Output = T>
226        + Add<T, Output = T>
227        + Sub<T, Output = T>
228        + Neg<Output = T>
229        + Float
230        + FloatConst
231        + Zero
232        + rand_distr::uniform::SampleUniform
233        + PartialOrd
234        + std::fmt::Display,
235    rand::distributions::OpenClosed01: Distribution<T>,
236    Uniform<T>: Distribution<T>,
237{
238    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
239        write!(
240            f,
241            "heat bath distribution with parameter {}",
242            self.param_exp()
243        )
244    }
245}
246
247/// Distribution for the norm of the SU2 adjoint to generate the [`HeatBathDistribution`] with the parameter
248/// `param_exp = beta * sqrt(det(A))`.
249///
250/// With distribution `dP(x) = dx \sqrt(1-x_0^2) e^{-2 param_exp x^2}`
251/// # Example
252/// ```
253/// use lattice_qcd_rs::error::ImplementationError;
254/// use lattice_qcd_rs::statistics::HeatBathDistributionNorm;
255/// use rand::{Rng, SeedableRng};
256///
257/// # fn main() -> Result<(), ImplementationError> {
258/// let mut rng = rand::rngs::StdRng::seed_from_u64(0);
259/// let heat_bath = HeatBathDistributionNorm::new(0.5_f64)
260///     .ok_or(ImplementationError::OptionWithUnexpectedNone)?;
261/// let r_number = rng.sample(&heat_bath);
262/// #
263/// # Ok(())
264/// # }
265/// ```
266#[derive(Clone, Debug, Copy, PartialEq, Hash, Eq)]
267#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
268pub struct HeatBathDistributionNorm<T>
269where
270    T: One
271        + Div<T, Output = T>
272        + Mul<T, Output = T>
273        + Add<T, Output = T>
274        + Sub<T, Output = T>
275        + Neg<Output = T>
276        + Float
277        + FloatConst
278        + Zero
279        + rand_distr::uniform::SampleUniform
280        + PartialOrd,
281    rand::distributions::OpenClosed01: Distribution<T>,
282    Uniform<T>: Distribution<T>,
283{
284    param_exp: T,
285}
286
287impl<T> HeatBathDistributionNorm<T>
288where
289    T: One
290        + Div<T, Output = T>
291        + Mul<T, Output = T>
292        + Add<T, Output = T>
293        + Neg<Output = T>
294        + Sub<T, Output = T>
295        + Float
296        + FloatConst
297        + Zero
298        + rand_distr::uniform::SampleUniform
299        + PartialOrd,
300    rand::distributions::OpenClosed01: Distribution<T>,
301    Uniform<T>: Distribution<T>,
302{
303    getter_copy!(
304        /// return the parameter `param_exp`.
305        pub const,
306        param_exp,
307        T
308    );
309
310    /// Create the distribution. `param_exp` should be strictly greater than 0 an be finite and a number.
311    /// Otherwise return [`None`].
312    pub fn new(param_exp: T) -> Option<Self> {
313        if param_exp.le(&T::zero()) || param_exp.is_infinite() || param_exp.is_nan() {
314            return None;
315        }
316        Some(Self { param_exp })
317    }
318}
319
320impl<T> Distribution<T> for HeatBathDistributionNorm<T>
321where
322    T: One
323        + Div<T, Output = T>
324        + Mul<T, Output = T>
325        + Add<T, Output = T>
326        + Neg<Output = T>
327        + Sub<T, Output = T>
328        + Float
329        + FloatConst
330        + Zero
331        + rand_distr::uniform::SampleUniform
332        + PartialOrd,
333    rand::distributions::OpenClosed01: Distribution<T>,
334    Uniform<T>: Distribution<T>,
335{
336    fn sample<R>(&self, rng: &mut R) -> T
337    where
338        R: rand::Rng + ?Sized,
339    {
340        let two = T::one() + T::one();
341        // the unwrap is OK because we verify that the param is not zero at creation.
342        let distributions = ModifiedNormal::new(self.param_exp()).unwrap();
343        loop {
344            let r = rng.sample(Uniform::new(T::zero(), T::one()));
345            let lambda = rng.sample(distributions);
346            if r.powi(2) <= T::one() - lambda.powi(2) {
347                return T::one() - two * lambda.powi(2);
348            }
349        }
350    }
351}
352
353impl<T> std::fmt::Display for HeatBathDistributionNorm<T>
354where
355    T: One
356        + Div<T, Output = T>
357        + Mul<T, Output = T>
358        + Add<T, Output = T>
359        + Neg<Output = T>
360        + Sub<T, Output = T>
361        + Float
362        + FloatConst
363        + Zero
364        + rand_distr::uniform::SampleUniform
365        + PartialOrd
366        + std::fmt::Display,
367    rand::distributions::OpenClosed01: Distribution<T>,
368    Uniform<T>: Distribution<T>,
369{
370    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
371        write!(
372            f,
373            "heat bath norm distribution with parameter {}",
374            self.param_exp()
375        )
376    }
377}
378
379/// used to generates matrix close the unit, for su2 close to +/-1, see [`su2::random_su2_close_to_unity`]
380/// and for su(3) `[su3::get_r] (+/- 1) * [su3::get_s] (+/- 1) * [su3::get_t] (+/- 1)`
381///
382/// # Example
383/// ```
384/// use lattice_qcd_rs::error::ImplementationError;
385/// use lattice_qcd_rs::statistics::CloseToUnit;
386/// use nalgebra::{Complex, Matrix2, Matrix3};
387/// use rand::{Rng, SeedableRng};
388///
389/// # fn main() -> Result<(), ImplementationError> {
390/// let mut rng = rand::rngs::StdRng::seed_from_u64(0);
391/// let close_to_unit =
392///     CloseToUnit::new(0.5_f64).ok_or(ImplementationError::OptionWithUnexpectedNone)?;
393/// let r_matrix: Matrix2<Complex<f64>> = rng.sample(&close_to_unit);
394/// let r_matrix: Matrix3<Complex<f64>> = rng.sample(&close_to_unit);
395/// #
396/// # Ok(())
397/// # }
398/// ```
399#[derive(Clone, Debug, Copy, PartialEq)]
400#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
401pub struct CloseToUnit {
402    spread_parameter: Real,
403}
404
405impl CloseToUnit {
406    getter_copy!(
407        /// Get the spread parameter
408        pub const,
409        spread_parameter,
410        f64
411    );
412
413    /// Create a new distribution, spread_parameter should be in `(0,1)` 0 and 1 excluded
414    pub fn new(spread_parameter: Real) -> Option<Self> {
415        if spread_parameter <= 0_f64 || spread_parameter >= 1_f64 || spread_parameter.is_nan() {
416            return None;
417        }
418        Some(Self { spread_parameter })
419    }
420}
421
422impl Distribution<CMatrix2> for CloseToUnit {
423    fn sample<R>(&self, rng: &mut R) -> CMatrix2
424    where
425        R: rand::Rng + ?Sized,
426    {
427        su2::random_su2_close_to_unity(self.spread_parameter, rng)
428    }
429}
430
431impl Distribution<CMatrix3> for CloseToUnit {
432    fn sample<R>(&self, rng: &mut R) -> CMatrix3
433    where
434        R: rand::Rng + ?Sized,
435    {
436        su3::random_su3_close_to_unity(self.spread_parameter, rng)
437    }
438}
439
440impl std::fmt::Display for CloseToUnit {
441    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442        write!(
443            f,
444            "distribution closed to unit with spread parameter {}",
445            self.spread_parameter()
446        )
447    }
448}
449
450#[cfg(test)]
451mod test {
452    use rand::{Rng, SeedableRng};
453
454    use super::*;
455
456    const SEED_RNG: u64 = 0x45_78_93_f4_4a_b0_67_f0;
457
458    #[test]
459    fn modified_normal() {
460        let mut rng = rand::rngs::StdRng::seed_from_u64(SEED_RNG);
461
462        for param in &[0.1_f64, 0.5_f64, 1_f64, 10_f64] {
463            let mn = ModifiedNormal::new(*param).unwrap();
464
465            for _ in 0_u32..1000_u32 {
466                assert!(rng.sample(&mn) >= 0_f64);
467            }
468        }
469    }
470
471    #[test]
472    #[allow(clippy::cognitive_complexity)]
473    fn distribution_creation() {
474        assert!(ModifiedNormal::new(0_f64).is_none());
475        assert!(ModifiedNormal::new(-1_f64).is_none());
476        assert!(ModifiedNormal::new(f64::NAN).is_none());
477        assert!(ModifiedNormal::new(f64::INFINITY).is_none());
478        assert!(ModifiedNormal::new(0.1_f64).is_some());
479        assert!(ModifiedNormal::new(2_f32).is_some());
480
481        let param = 0.5_f64;
482        let mn = ModifiedNormal::new(param).unwrap();
483        assert_eq!(mn.param_exp(), param);
484        assert_eq!(
485            mn.to_string(),
486            "modified normal distribution with parameter 0.5"
487        );
488
489        assert!(HeatBathDistributionNorm::new(0_f64).is_none());
490        assert!(HeatBathDistributionNorm::new(-1_f64).is_none());
491        assert!(HeatBathDistributionNorm::new(f64::NAN).is_none());
492        assert!(HeatBathDistributionNorm::new(f64::INFINITY).is_none());
493        assert!(HeatBathDistributionNorm::new(0.1_f64).is_some());
494        assert!(HeatBathDistributionNorm::new(2_f32).is_some());
495
496        let heat_bath_norm = HeatBathDistributionNorm::new(param).unwrap();
497        assert_eq!(heat_bath_norm.param_exp(), param);
498        assert_eq!(
499            heat_bath_norm.to_string(),
500            "heat bath norm distribution with parameter 0.5"
501        );
502
503        assert!(HeatBathDistribution::new(0_f64).is_none());
504        assert!(HeatBathDistribution::new(-1_f64).is_none());
505        assert!(HeatBathDistribution::new(f64::NAN).is_none());
506        assert!(HeatBathDistribution::new(f64::INFINITY).is_none());
507        assert!(HeatBathDistribution::new(0.1_f64).is_some());
508        assert!(HeatBathDistribution::new(2_f32).is_some());
509
510        let heat_bath = HeatBathDistribution::new(param).unwrap();
511        assert_eq!(heat_bath.param_exp(), param);
512        assert_eq!(
513            heat_bath.to_string(),
514            "heat bath distribution with parameter 0.5"
515        );
516
517        assert!(CloseToUnit::new(0_f64).is_none());
518        assert!(CloseToUnit::new(-1_f64).is_none());
519        assert!(CloseToUnit::new(f64::NAN).is_none());
520        assert!(CloseToUnit::new(f64::INFINITY).is_none());
521        assert!(CloseToUnit::new(2_f64).is_none());
522        assert!(CloseToUnit::new(0.5_f64).is_some());
523
524        let cu = CloseToUnit::new(param).unwrap();
525        assert_eq!(cu.spread_parameter(), param);
526        assert_eq!(
527            cu.to_string(),
528            "distribution closed to unit with spread parameter 0.5"
529        );
530    }
531}
lattice_qcd_rs/statistics/distribution.rs

lattice_qcd_rs/statistics/
distribution.rs
