Skip to main content

ggplot_rs/stat/
summary_bin.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::summary::SummaryFun;
6use super::Stat;
7
8/// Bin x values and apply summary function to y within each bin.
9/// Produces x (bin center), y, ymin, ymax columns.
10pub struct StatSummaryBin {
11    pub bins: usize,
12    pub fun_y: SummaryFun,
13    pub fun_ymin: SummaryFun,
14    pub fun_ymax: SummaryFun,
15}
16
17impl Default for StatSummaryBin {
18    fn default() -> Self {
19        StatSummaryBin {
20            bins: 30,
21            fun_y: SummaryFun::Mean,
22            fun_ymin: SummaryFun::Min,
23            fun_ymax: SummaryFun::Max,
24        }
25    }
26}
27
28impl StatSummaryBin {
29    pub fn with_bins(mut self, bins: usize) -> Self {
30        self.bins = bins;
31        self
32    }
33
34    pub fn with_fun(mut self, fun_y: SummaryFun) -> Self {
35        self.fun_y = fun_y;
36        self
37    }
38
39    pub fn with_fun_range(mut self, fun_ymin: SummaryFun, fun_ymax: SummaryFun) -> Self {
40        self.fun_ymin = fun_ymin;
41        self.fun_ymax = fun_ymax;
42        self
43    }
44}
45
46impl Stat for StatSummaryBin {
47    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
48        let x_col = match data.column("x") {
49            Some(c) => c,
50            None => return DataFrame::new(),
51        };
52        let y_col = match data.column("y") {
53            Some(c) => c,
54            None => return DataFrame::new(),
55        };
56
57        // Extract numeric pairs
58        let mut pairs: Vec<(f64, f64)> = Vec::new();
59        for (x, y) in x_col.iter().zip(y_col.iter()) {
60            if let (Some(xv), Some(yv)) = (x.as_f64(), y.as_f64()) {
61                if xv.is_finite() && yv.is_finite() {
62                    pairs.push((xv, yv));
63                }
64            }
65        }
66
67        if pairs.is_empty() {
68            return DataFrame::new();
69        }
70
71        let x_min = pairs.iter().map(|p| p.0).fold(f64::INFINITY, f64::min);
72        let x_max = pairs.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
73
74        let (x_min, x_max) = if (x_max - x_min).abs() < f64::EPSILON {
75            (x_min - 0.5, x_max + 0.5)
76        } else {
77            (x_min, x_max)
78        };
79
80        let bin_width = (x_max - x_min) / self.bins as f64;
81        let n_bins = self.bins;
82
83        // Collect y values per bin
84        let mut bin_ys: Vec<Vec<f64>> = vec![Vec::new(); n_bins];
85        for &(x, y) in &pairs {
86            let bin = ((x - x_min) / bin_width).floor() as usize;
87            let bin = bin.min(n_bins - 1);
88            bin_ys[bin].push(y);
89        }
90
91        let mut x_vals = Vec::new();
92        let mut y_vals = Vec::new();
93        let mut ymin_vals = Vec::new();
94        let mut ymax_vals = Vec::new();
95
96        for (i, ys) in bin_ys.iter().enumerate() {
97            if ys.is_empty() {
98                continue;
99            }
100            let bin_center = x_min + (i as f64 + 0.5) * bin_width;
101            x_vals.push(Value::Float(bin_center));
102            y_vals.push(Value::Float(self.fun_y.apply(ys)));
103            ymin_vals.push(Value::Float(self.fun_ymin.apply(ys)));
104            ymax_vals.push(Value::Float(self.fun_ymax.apply(ys)));
105        }
106
107        let mut result = DataFrame::new();
108        result.add_column("x".to_string(), x_vals);
109        result.add_column("y".to_string(), y_vals);
110        result.add_column("ymin".to_string(), ymin_vals);
111        result.add_column("ymax".to_string(), ymax_vals);
112
113        // Carry over grouping columns
114        let n = result.nrows();
115        for col_name in &["color", "fill", "group"] {
116            if let Some(col) = data.column(col_name) {
117                if let Some(first) = col.first() {
118                    result.add_column(col_name.to_string(), vec![first.clone(); n]);
119                }
120            }
121        }
122
123        result
124    }
125
126    fn required_aes(&self) -> Vec<Aesthetic> {
127        vec![Aesthetic::X, Aesthetic::Y]
128    }
129
130    fn name(&self) -> &str {
131        "summary_bin"
132    }
133}