laddu_core/utils/
mod.rs

1/// Useful enumerations for various frames and variables common in particle physics analyses.
2pub mod enums;
3/// Standard special functions like spherical harmonics and momentum definitions.
4pub mod functions;
5/// Traits and structs which can be used to extract complex information from
6/// [`EventData`](crate::data::EventData)s.
7pub mod variables;
8/// Traits to give additional functionality to [`nalgebra::Vector3`] and [`nalgebra::Vector4`] (in
9/// particular, to treat the latter as a four-momentum).
10pub mod vectors;
11
12/// A helper method to get histogram edges from evenly-spaced `bins` over a given `range`
13///
14/// # See Also
15/// [`Histogram`]
16/// [`get_bin_index`]
17///
18/// # Examples
19/// ```rust
20/// use laddu_core::utils::get_bin_edges;
21///
22/// assert_eq!(get_bin_edges(3, (0.0, 3.0)), vec![0.0, 1.0, 2.0, 3.0]);
23/// ```
24pub fn get_bin_edges(bins: usize, range: (f64, f64)) -> Vec<f64> {
25    let bin_width = (range.1 - range.0) / (bins as f64);
26    (0..=bins)
27        .map(|i| range.0 + (i as f64 * bin_width))
28        .collect()
29}
30
31/// A helper method to obtain the index of a bin where a value should go in a histogram with evenly
32/// spaced `bins` over a given `range`
33///
34/// # See Also
35/// [`Histogram`]
36/// [`get_bin_edges`]
37///
38/// # Examples
39/// ```rust
40/// use laddu_core::utils::get_bin_index;
41///
42/// assert_eq!(get_bin_index(0.25, 4, (0.0, 1.0)), Some(1));
43/// assert_eq!(get_bin_index(1.5, 4, (0.0, 1.0)), None);
44/// ```
45pub fn get_bin_index(value: f64, bins: usize, range: (f64, f64)) -> Option<usize> {
46    if value >= range.0 && value < range.1 {
47        let bin_width = (range.1 - range.0) / bins as f64;
48        let bin_index = ((value - range.0) / bin_width).floor() as usize;
49        Some(bin_index.min(bins - 1))
50    } else {
51        None
52    }
53}
54
55/// A simple struct which represents a histogram
56pub struct Histogram {
57    /// The number of counts in each bin (can be [`f64`]s since these might be weighted counts)
58    pub counts: Vec<f64>,
59    /// The edges of each bin (length is one greater than `counts`)
60    pub bin_edges: Vec<f64>,
61}
62
63/// A method which creates a histogram from some data by binning it with evenly spaced `bins` within
64/// the given `range`
65///
66/// # Examples
67/// ```rust
68/// use laddu_core::utils::histogram;
69///
70/// let values = vec![0.1, 0.4, 0.8];
71/// let weights: Option<&[f64]> = None;
72/// let hist = histogram(values.as_slice(), 2, (0.0, 1.0), weights);
73/// assert_eq!(hist.counts, vec![2.0, 1.0]);
74/// assert_eq!(hist.bin_edges, vec![0.0, 0.5, 1.0]);
75/// ```
76pub fn histogram<T: AsRef<[f64]>>(
77    values: T,
78    bins: usize,
79    range: (f64, f64),
80    weights: Option<T>,
81) -> Histogram {
82    assert!(bins > 0, "Number of bins must be greater than zero!");
83    assert!(
84        range.1 > range.0,
85        "The lower edge of the range must be smaller than the upper edge!"
86    );
87    if let Some(w) = &weights {
88        assert_eq!(
89            values.as_ref().len(),
90            w.as_ref().len(),
91            "`values` and `weights` must have the same length!"
92        );
93    }
94    let mut counts = vec![0.0; bins];
95    for (i, &value) in values.as_ref().iter().enumerate() {
96        if let Some(bin_index) = get_bin_index(value, bins, range) {
97            let weight = weights.as_ref().map_or(1.0, |w| w.as_ref()[i]);
98            counts[bin_index] += weight;
99        }
100    }
101    Histogram {
102        counts,
103        bin_edges: get_bin_edges(bins, range),
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use std::sync::Arc;
110
111    use crate::{
112        data::test_dataset,
113        traits::Variable,
114        utils::{get_bin_index, histogram},
115        Mass,
116    };
117
118    #[test]
119    fn test_binning() {
120        let mut v = Mass::new(["kshort1"]);
121        let dataset = Arc::new(test_dataset());
122        v.bind(dataset.metadata()).unwrap();
123        let values = v.value_on(&dataset).unwrap();
124        let bin_index = get_bin_index(values[0], 3, (0.0, 1.0));
125        assert_eq!(bin_index, Some(1));
126        let bin_index = get_bin_index(0.0, 3, (0.0, 1.0));
127        assert_eq!(bin_index, Some(0));
128        let bin_index = get_bin_index(0.1, 3, (0.0, 1.0));
129        assert_eq!(bin_index, Some(0));
130        let bin_index = get_bin_index(0.9, 3, (0.0, 1.0));
131        assert_eq!(bin_index, Some(2));
132        let bin_index = get_bin_index(1.0, 3, (0.0, 1.0));
133        assert_eq!(bin_index, None);
134        let bin_index = get_bin_index(2.0, 3, (0.0, 1.0));
135        assert_eq!(bin_index, None);
136        let weights = dataset.weights();
137        let histogram = histogram(&values, 3, (0.0, 1.0), Some(&weights));
138        assert_eq!(histogram.counts, vec![0.0, 0.48, 0.0]);
139        assert_eq!(histogram.bin_edges, vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0])
140    }
141}
laddu_core/utils/mod.rs

laddu_core/utils/
mod.rs
