laddu_generation/
distributions.rs

1use fastrand::Rng;
2use fastrand_contrib::RngExt;
3use laddu_core::{math::Histogram, LadduResult, Vec3, Vec4};
4
5/// Sampler for drawing values from a weighted histogram.
6#[derive(Clone, Debug)]
7pub struct HistogramSampler {
8    pub(crate) hist: Histogram,
9    cdf: Vec<f64>,
10    total: f64,
11}
12
13impl HistogramSampler {
14    /// Construct a histogram sampler.
15    pub fn new(hist: Histogram) -> LadduResult<Self> {
16        hist.validate()?;
17        hist.validate_positive_counts()?;
18        let mut cdf = Vec::with_capacity(hist.counts().len());
19        let mut total = 0.0;
20
21        for &count in hist.counts() {
22            total += count;
23            cdf.push(total);
24        }
25        Ok(Self { hist, cdf, total })
26    }
27
28    /// Sample a value uniformly within a histogram bin selected by bin weight.
29    pub fn sample(&self, rng: &mut Rng) -> f64 {
30        let r = rng.f64() * self.total;
31        let bin = self.cdf.partition_point(|&x| x <= r);
32        let lo = self.hist.bin_edges()[bin];
33        let hi = self.hist.bin_edges()[bin + 1];
34        lo + rng.f64() * (hi - lo)
35    }
36}
37
38#[derive(Clone, Debug)]
39pub enum SimpleDistribution {
40    Fixed(f64),
41    Uniform { min: f64, max: f64 },
42    Histogram(HistogramSampler),
43}
44impl SimpleDistribution {
45    pub fn sample(&self, rng: &mut Rng) -> f64 {
46        match self {
47            Self::Fixed(v) => *v,
48            Self::Uniform { min, max } => rng.uniform(*min, *max),
49            Self::Histogram(sampler) => sampler.sample(rng),
50        }
51    }
52}
53
54#[derive(Clone, Debug)]
55pub enum MandelstamTDistribution {
56    Exponential { slope: f64 },
57    Histogram(HistogramSampler),
58}
59impl MandelstamTDistribution {
60    pub fn sample(&self, rng: &mut Rng, range: Option<(f64, f64)>) -> f64 {
61        match self {
62            Self::Exponential { slope } => {
63                if let Some(range) = range {
64                    let mut result = -rng.truncated_exponential(*slope, range);
65                    while result <= range.0 || result >= range.1 {
66                        result = -rng.truncated_exponential(*slope, range)
67                    }
68                    result
69                } else {
70                    -rng.exponential(*slope)
71                }
72            }
73            Self::Histogram(sampler) => {
74                if let Some(range) = range {
75                    let mut result = sampler.sample(rng);
76                    while result <= range.0 || result >= range.1 {
77                        result = sampler.sample(rng);
78                    }
79                    result
80                } else {
81                    sampler.sample(rng)
82                }
83            }
84        }
85    }
86}
87
88#[derive(Clone, Debug)]
89pub enum Distribution {
90    Fixed(f64),
91    Uniform { min: f64, max: f64 },
92    Normal { mu: f64, sigma: f64 },
93    Exponential { slope: f64 },
94    Histogram(HistogramSampler),
95}
96impl Distribution {
97    pub fn sample(&self, rng: &mut Rng) -> f64 {
98        match self {
99            Self::Fixed(v) => *v,
100            Self::Uniform { min, max } => rng.uniform(*min, *max),
101            Self::Normal { mu, sigma } => rng.normal(*mu, *sigma),
102            Self::Exponential { slope } => rng.exponential(*slope),
103            Self::Histogram(hist) => hist.sample(rng),
104        }
105    }
106}
107
108pub trait LadduGenRngExt {
109    fn uniform(&mut self, min: f64, max: f64) -> f64;
110    fn normal(&mut self, mu: f64, sigma: f64) -> f64;
111    fn exponential(&mut self, slope: f64) -> f64;
112    fn truncated_exponential(&mut self, slope: f64, range: (f64, f64)) -> f64;
113    fn p4(&mut self, mass: f64, energy: f64, direction: Vec3) -> Vec4;
114}
115
116impl LadduGenRngExt for Rng {
117    fn uniform(&mut self, min: f64, max: f64) -> f64 {
118        self.f64_range(min..=max)
119    }
120
121    fn normal(&mut self, mu: f64, sigma: f64) -> f64 {
122        self.f64_normal_approx(mu, sigma)
123    }
124
125    fn exponential(&mut self, slope: f64) -> f64 {
126        -(-self.f64()).ln_1p() / slope
127    }
128
129    fn truncated_exponential(&mut self, slope: f64, range: (f64, f64)) -> f64 {
130        -(1. / slope) * (1.0 - self.f64() * (1.0 - (-slope * (range.1 - range.0)).exp())).ln()
131    }
132
133    fn p4(&mut self, mass: f64, energy: f64, direction: Vec3) -> Vec4 {
134        debug_assert!(
135            energy >= mass,
136            "Mass cannot be greater than energy!\nEnergy: {}\nMass: {}",
137            energy,
138            mass
139        );
140        let momentum = ((energy - mass) * (energy + mass)).max(0.0).sqrt();
141        (momentum * direction).with_mass(mass)
142    }
143}
laddu_generation/distributions.rs

laddu_generation/
distributions.rs
