lindhard/tally/
hist.rs

1//! Fixed-binning histograms with integer counts.
2//!
3//! Counts are `u64`, so merging two histograms is an exact, associative and
4//! commutative sum: the merged result does not depend on how the samples were
5//! split between workers.
6
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9
10/// Errors from building a [`Binning`].
11#[derive(Debug, Clone, PartialEq, Error)]
12pub enum BinningError {
13    /// The edges are not finite with `hi > lo`, or there are no bins.
14    #[error(
15        "invalid binning: lo = {lo}, hi = {hi}, bins = {bins} (need finite lo < hi and bins > 0)"
16    )]
17    Invalid {
18        /// Lower edge.
19        lo: f64,
20        /// Upper edge.
21        hi: f64,
22        /// Number of bins.
23        bins: usize,
24    },
25}
26
27/// Uniform binning of `[lo, hi)` into `bins` equal bins.
28///
29/// The unit of `lo` and `hi` is the unit of the quantity binned (metres for
30/// depths, eV for energies, radians for angles); each use documents it.
31#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
32pub struct Binning {
33    /// Lower edge of the first bin.
34    pub lo: f64,
35    /// Upper edge of the last bin.
36    pub hi: f64,
37    /// Number of bins.
38    pub bins: usize,
39}
40
41impl Binning {
42    /// Validated binning of `[lo, hi)` into `bins` bins.
43    pub fn new(lo: f64, hi: f64, bins: usize) -> Result<Self, BinningError> {
44        if lo.is_finite() && hi.is_finite() && hi > lo && bins > 0 {
45            Ok(Self { lo, hi, bins })
46        } else {
47            Err(BinningError::Invalid { lo, hi, bins })
48        }
49    }
50
51    /// Width of one bin.
52    pub fn width(&self) -> f64 {
53        (self.hi - self.lo) / self.bins as f64
54    }
55
56    /// Lower edge of bin `i`.
57    pub fn edge(&self, i: usize) -> f64 {
58        self.lo + (self.hi - self.lo) * (i as f64 / self.bins as f64)
59    }
60
61    /// Centre of bin `i`.
62    pub fn center(&self, i: usize) -> f64 {
63        0.5 * (self.edge(i) + self.edge(i + 1))
64    }
65
66    /// Where `x` falls: `Ok(bin)`, or `Err(Below)` / `Err(Above)`. NaN counts
67    /// as above.
68    pub fn locate(&self, x: f64) -> Result<usize, OutOfRange> {
69        if x < self.lo {
70            return Err(OutOfRange::Below);
71        }
72        if x.is_nan() || x >= self.hi {
73            return Err(OutOfRange::Above);
74        }
75        let i = ((x - self.lo) / (self.hi - self.lo) * self.bins as f64) as usize;
76        Ok(i.min(self.bins - 1))
77    }
78}
79
80/// Which side of a [`Binning`] a value fell on.
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum OutOfRange {
83    /// `x < lo`.
84    Below,
85    /// `x >= hi`, or NaN.
86    Above,
87}
88
89/// Counts over a [`Binning`], with underflow and overflow.
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
91pub struct Histogram {
92    /// The binning.
93    pub binning: Binning,
94    /// Count per bin.
95    pub counts: Vec<u64>,
96    /// Samples below `binning.lo`.
97    pub underflow: u64,
98    /// Samples at or above `binning.hi` (and NaN samples).
99    pub overflow: u64,
100}
101
102impl Histogram {
103    /// Empty histogram.
104    pub fn new(binning: Binning) -> Self {
105        Self {
106            binning,
107            counts: vec![0; binning.bins],
108            underflow: 0,
109            overflow: 0,
110        }
111    }
112
113    /// Count one sample.
114    pub fn fill(&mut self, x: f64) {
115        match self.binning.locate(x) {
116            Ok(i) => self.counts[i] += 1,
117            Err(OutOfRange::Below) => self.underflow += 1,
118            Err(OutOfRange::Above) => self.overflow += 1,
119        }
120    }
121
122    /// All samples, in range or not.
123    pub fn total(&self) -> u64 {
124        self.counts.iter().sum::<u64>() + self.underflow + self.overflow
125    }
126
127    /// Add `other`, which must have the same binning.
128    ///
129    /// # Panics
130    /// If the binnings differ (merging tallies built with different settings
131    /// is a programming error).
132    pub fn merge(&mut self, other: &Histogram) {
133        assert_eq!(
134            self.binning, other.binning,
135            "merging histograms with different binnings"
136        );
137        for (a, b) in self.counts.iter_mut().zip(&other.counts) {
138            *a += b;
139        }
140        self.underflow += other.underflow;
141        self.overflow += other.overflow;
142    }
143
144    /// Probability density per bin: `count / (total * width)`, so that the
145    /// in-range densities integrate to the in-range fraction. Unit: inverse
146    /// of the binned quantity's unit. All zeros if the histogram is empty.
147    pub fn density(&self) -> Vec<f64> {
148        let n = self.total();
149        if n == 0 {
150            return vec![0.0; self.counts.len()];
151        }
152        let norm = 1.0 / (n as f64 * self.binning.width());
153        self.counts.iter().map(|&c| c as f64 * norm).collect()
154    }
155}
156
157/// Weighted sums over a [`Binning`], with underflow and overflow.
158///
159/// Unlike [`Histogram`] the bin contents are `f64` sums, so merging is
160/// deterministic only for a fixed merge order (chunk order, as for the other
161/// float accumulators of the ion tally).
162#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
163pub struct WeightedHistogram {
164    /// The binning.
165    pub binning: Binning,
166    /// Summed weight per bin.
167    pub sums: Vec<f64>,
168    /// Weight of samples below `binning.lo`.
169    pub underflow: f64,
170    /// Weight of samples at or above `binning.hi` (and NaN samples).
171    pub overflow: f64,
172}
173
174impl WeightedHistogram {
175    /// Empty histogram.
176    pub fn new(binning: Binning) -> Self {
177        Self {
178            binning,
179            sums: vec![0.0; binning.bins],
180            underflow: 0.0,
181            overflow: 0.0,
182        }
183    }
184
185    /// Add weight `w` at `x`.
186    pub fn fill(&mut self, x: f64, w: f64) {
187        match self.binning.locate(x) {
188            Ok(i) => self.sums[i] += w,
189            Err(OutOfRange::Below) => self.underflow += w,
190            Err(OutOfRange::Above) => self.overflow += w,
191        }
192    }
193
194    /// All weight, in range or not.
195    pub fn total(&self) -> f64 {
196        self.sums.iter().sum::<f64>() + self.underflow + self.overflow
197    }
198
199    /// Add `other`, which must have the same binning.
200    ///
201    /// # Panics
202    /// If the binnings differ.
203    pub fn merge(&mut self, other: &WeightedHistogram) {
204        assert_eq!(
205            self.binning, other.binning,
206            "merging histograms with different binnings"
207        );
208        for (a, b) in self.sums.iter_mut().zip(&other.sums) {
209            *a += b;
210        }
211        self.underflow += other.underflow;
212        self.overflow += other.overflow;
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn binning_rejects_bad_input() {
222        assert!(Binning::new(0.0, 0.0, 1).is_err());
223        assert!(Binning::new(1.0, 0.0, 1).is_err());
224        assert!(Binning::new(0.0, 1.0, 0).is_err());
225        assert!(Binning::new(f64::NAN, 1.0, 3).is_err());
226        assert!(Binning::new(0.0, f64::INFINITY, 3).is_err());
227    }
228
229    #[test]
230    fn fill_and_edges() {
231        let mut h = Histogram::new(Binning::new(0.0, 10.0, 5).unwrap());
232        for x in [-1.0, 0.0, 1.99, 2.0, 9.999, 10.0, f64::NAN] {
233            h.fill(x);
234        }
235        assert_eq!(h.counts, vec![2, 1, 0, 0, 1]);
236        assert_eq!((h.underflow, h.overflow, h.total()), (1, 2, 7));
237        assert_eq!(h.binning.center(0), 1.0);
238        let mut g = h.clone();
239        g.merge(&h);
240        assert_eq!(g.counts, vec![4, 2, 0, 0, 2]);
241        assert_eq!(g.total(), 14);
242    }
243}
lindhard/tally/hist.rs

lindhard/tally/
hist.rs
