nereids_io/
rebin.rs

1//! Rebinning: coarsen the TOF axis by an integer factor.
2//!
3//! Counts in adjacent bins are summed, and the merged bins keep every
4//! `factor`-th edge. Rebinning needs bin edges: a merged bin's position
5//! depends on the widths of the bins it merges, which centres do not give.
6
7use ndarray::{Array3, Axis, s};
8
9/// Rebin a 3D counts array along axis 0 by summing groups of `factor` slices.
10///
11/// Remainder slices at the end are summed into the last output bin.
12/// Returns the original array unchanged if `factor < 2`.
13///
14/// **NaN handling**: NaN values propagate through the sum — any NaN in a
15/// group produces NaN in that output bin. This matches ndarray semantics.
16pub fn rebin_counts(data: &Array3<f64>, factor: usize) -> Array3<f64> {
17    if factor < 2 {
18        return data.clone();
19    }
20    let n_old = data.shape()[0];
21    let height = data.shape()[1];
22    let width = data.shape()[2];
23    let n_new = n_old.div_ceil(factor);
24    let mut out = Array3::<f64>::zeros((n_new, height, width));
25    for i in 0..n_new {
26        let start = i * factor;
27        let end = (start + factor).min(n_old);
28        out.slice_mut(s![i, .., ..])
29            .assign(&data.slice(s![start..end, .., ..]).sum_axis(Axis(0)));
30    }
31    out
32}
33
34/// Rebin bin edges by keeping every `factor`-th edge, always including
35/// the final edge.
36///
37/// Input: N+1 edges for N bins.
38/// Output: ceil(N/factor)+1 edges for ceil(N/factor) bins.
39///
40/// Returns the original edges unchanged if `factor < 2`.
41pub fn rebin_edges(edges: &[f64], factor: usize) -> Vec<f64> {
42    if factor < 2 || edges.len() <= 1 {
43        return edges.to_vec();
44    }
45    let mut out: Vec<f64> = edges.iter().step_by(factor).copied().collect();
46    // Always include the final edge
47    if out.last() != edges.last() {
48        out.push(*edges.last().unwrap());
49    }
50    out
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use ndarray::Array3;
57
58    #[test]
59    fn counts_conservation() {
60        let data = Array3::from_shape_fn((6, 2, 3), |(t, y, x)| (t + y + x) as f64);
61        let original_sum: f64 = data.sum();
62        let rebinned = rebin_counts(&data, 2);
63        assert_eq!(rebinned.shape(), &[3, 2, 3]);
64        assert!((rebinned.sum() - original_sum).abs() < 1e-10);
65    }
66
67    #[test]
68    fn counts_conservation_remainder() {
69        // 7 bins / 3 = 3 output bins (2+2+3 or ceil division)
70        let data = Array3::from_shape_fn((7, 2, 2), |(t, _y, _x)| t as f64 + 1.0);
71        let original_sum: f64 = data.sum();
72        let rebinned = rebin_counts(&data, 3);
73        assert_eq!(rebinned.shape(), &[3, 2, 2]);
74        assert!((rebinned.sum() - original_sum).abs() < 1e-10);
75    }
76
77    #[test]
78    fn counts_factor_one_is_identity() {
79        let data = Array3::from_shape_fn((5, 2, 2), |(t, y, x)| (t * 10 + y + x) as f64);
80        let rebinned = rebin_counts(&data, 1);
81        assert_eq!(rebinned, data);
82    }
83
84    #[test]
85    fn counts_factor_ge_nbins_single_bin() {
86        let data = Array3::from_shape_fn((4, 2, 2), |(t, _y, _x)| t as f64);
87        let rebinned = rebin_counts(&data, 100);
88        assert_eq!(rebinned.shape(), &[1, 2, 2]);
89        assert!((rebinned.sum() - data.sum()).abs() < 1e-10);
90    }
91
92    #[test]
93    fn edges_even_division() {
94        let edges = vec![0.0, 1.0, 2.0, 3.0, 4.0]; // 4 bins
95        let rebinned = rebin_edges(&edges, 2);
96        assert_eq!(rebinned, vec![0.0, 2.0, 4.0]); // 2 bins
97    }
98
99    #[test]
100    fn edges_remainder() {
101        let edges = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]; // 5 bins
102        let rebinned = rebin_edges(&edges, 2);
103        // step_by(2): [0, 2, 4] → last edge 5.0 appended
104        assert_eq!(rebinned, vec![0.0, 2.0, 4.0, 5.0]);
105    }
106
107    #[test]
108    fn edges_factor_one() {
109        let edges = vec![1.0, 2.0, 3.0];
110        assert_eq!(rebin_edges(&edges, 1), edges);
111    }
112
113    #[test]
114    fn counts_empty_array() {
115        let data = Array3::<f64>::zeros((0, 2, 2));
116        let rebinned = rebin_counts(&data, 3);
117        assert_eq!(rebinned.shape(), &[0, 2, 2]);
118    }
119
120    #[test]
121    fn counts_factor_zero_returns_clone() {
122        let data = Array3::from_shape_fn((4, 1, 1), |(t, _, _)| t as f64);
123        let rebinned = rebin_counts(&data, 0);
124        assert_eq!(rebinned, data);
125    }
126
127    #[test]
128    fn counts_nan_propagation() {
129        let mut data = Array3::from_shape_fn((4, 1, 1), |(t, _, _)| t as f64);
130        data[[1, 0, 0]] = f64::NAN; // bin 1 is NaN
131        let rebinned = rebin_counts(&data, 2);
132        // Group [0, NaN] should produce NaN
133        assert!(rebinned[[0, 0, 0]].is_nan());
134        // Group [2, 3] should produce 5.0
135        assert!((rebinned[[1, 0, 0]] - 5.0).abs() < 1e-10);
136    }
137
138    #[test]
139    fn edges_factor_zero_returns_clone() {
140        let edges = vec![0.0, 1.0, 2.0, 3.0];
141        assert_eq!(rebin_edges(&edges, 0), edges);
142    }
143}
nereids_io/rebin.rs

nereids_io/
rebin.rs
