laddu-generation 0.19.6

Monte Carlo event generation tools for laddu
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use fastrand::Rng;
use fastrand_contrib::RngExt;
use laddu_core::{math::Histogram, LadduResult, Vec3, Vec4};

/// Sampler for drawing values from a weighted histogram.
#[derive(Clone, Debug)]
pub struct HistogramSampler {
    pub(crate) hist: Histogram,
    cdf: Vec<f64>,
    total: f64,
}

impl HistogramSampler {
    /// Construct a histogram sampler.
    pub fn new(hist: Histogram) -> LadduResult<Self> {
        hist.validate()?;
        hist.validate_positive_counts()?;
        let mut cdf = Vec::with_capacity(hist.counts().len());
        let mut total = 0.0;

        for &count in hist.counts() {
            total += count;
            cdf.push(total);
        }
        Ok(Self { hist, cdf, total })
    }

    /// Sample a value uniformly within a histogram bin selected by bin weight.
    pub fn sample(&self, rng: &mut Rng) -> f64 {
        let r = rng.f64() * self.total;
        let bin = self.cdf.partition_point(|&x| x <= r);
        let lo = self.hist.bin_edges()[bin];
        let hi = self.hist.bin_edges()[bin + 1];
        lo + rng.f64() * (hi - lo)
    }
}

#[derive(Clone, Debug)]
pub enum SimpleDistribution {
    Fixed(f64),
    Uniform { min: f64, max: f64 },
    Histogram(HistogramSampler),
}
impl SimpleDistribution {
    pub fn sample(&self, rng: &mut Rng) -> f64 {
        match self {
            Self::Fixed(v) => *v,
            Self::Uniform { min, max } => rng.uniform(*min, *max),
            Self::Histogram(sampler) => sampler.sample(rng),
        }
    }
}

#[derive(Clone, Debug)]
pub enum MandelstamTDistribution {
    Exponential { slope: f64 },
    Histogram(HistogramSampler),
}
impl MandelstamTDistribution {
    pub fn sample(&self, rng: &mut Rng, range: Option<(f64, f64)>) -> f64 {
        match self {
            Self::Exponential { slope } => {
                if let Some(range) = range {
                    let mut result = -rng.truncated_exponential(*slope, range);
                    while result <= range.0 || result >= range.1 {
                        result = -rng.truncated_exponential(*slope, range)
                    }
                    result
                } else {
                    -rng.exponential(*slope)
                }
            }
            Self::Histogram(sampler) => {
                if let Some(range) = range {
                    let mut result = sampler.sample(rng);
                    while result <= range.0 || result >= range.1 {
                        result = sampler.sample(rng);
                    }
                    result
                } else {
                    sampler.sample(rng)
                }
            }
        }
    }
}

#[derive(Clone, Debug)]
pub enum Distribution {
    Fixed(f64),
    Uniform { min: f64, max: f64 },
    Normal { mu: f64, sigma: f64 },
    Exponential { slope: f64 },
    Histogram(HistogramSampler),
}
impl Distribution {
    pub fn sample(&self, rng: &mut Rng) -> f64 {
        match self {
            Self::Fixed(v) => *v,
            Self::Uniform { min, max } => rng.uniform(*min, *max),
            Self::Normal { mu, sigma } => rng.normal(*mu, *sigma),
            Self::Exponential { slope } => rng.exponential(*slope),
            Self::Histogram(hist) => hist.sample(rng),
        }
    }
}

pub trait LadduGenRngExt {
    fn uniform(&mut self, min: f64, max: f64) -> f64;
    fn normal(&mut self, mu: f64, sigma: f64) -> f64;
    fn exponential(&mut self, slope: f64) -> f64;
    fn truncated_exponential(&mut self, slope: f64, range: (f64, f64)) -> f64;
    fn p4(&mut self, mass: f64, energy: f64, direction: Vec3) -> Vec4;
}

impl LadduGenRngExt for Rng {
    fn uniform(&mut self, min: f64, max: f64) -> f64 {
        self.f64_range(min..=max)
    }

    fn normal(&mut self, mu: f64, sigma: f64) -> f64 {
        self.f64_normal_approx(mu, sigma)
    }

    fn exponential(&mut self, slope: f64) -> f64 {
        -(-self.f64()).ln_1p() / slope
    }

    fn truncated_exponential(&mut self, slope: f64, range: (f64, f64)) -> f64 {
        -(1. / slope) * (1.0 - self.f64() * (1.0 - (-slope * (range.1 - range.0)).exp())).ln()
    }

    fn p4(&mut self, mass: f64, energy: f64, direction: Vec3) -> Vec4 {
        debug_assert!(
            energy >= mass,
            "Mass cannot be greater than energy!\nEnergy: {}\nMass: {}",
            energy,
            mass
        );
        let momentum = ((energy - mass) * (energy + mass)).max(0.0).sqrt();
        (momentum * direction).with_mass(mass)
    }
}