laddu_core/utils/
mod.rs

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

laddu_core/utils/
mod.rs
