ggplot_rs/stat/
summary.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Aggregation function type for StatSummary.
8#[derive(Clone)]
9pub enum SummaryFun {
10    Mean,
11    Median,
12    Min,
13    Max,
14    Sum,
15}
16
17impl SummaryFun {
18    pub fn apply(&self, values: &[f64]) -> f64 {
19        if values.is_empty() {
20            return 0.0;
21        }
22        match self {
23            SummaryFun::Mean => values.iter().sum::<f64>() / values.len() as f64,
24            SummaryFun::Median => {
25                let mut sorted = values.to_vec();
26                sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
27                let n = sorted.len();
28                if n.is_multiple_of(2) {
29                    (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0
30                } else {
31                    sorted[n / 2]
32                }
33            }
34            SummaryFun::Min => values.iter().cloned().fold(f64::INFINITY, f64::min),
35            SummaryFun::Max => values.iter().cloned().fold(f64::NEG_INFINITY, f64::max),
36            SummaryFun::Sum => values.iter().sum(),
37        }
38    }
39}
40
41/// Summarize y values for each unique x with a summary function.
42/// Also computes ymin and ymax using configurable functions.
43pub struct StatSummary {
44    pub fun_y: SummaryFun,
45    pub fun_ymin: SummaryFun,
46    pub fun_ymax: SummaryFun,
47}
48
49impl Default for StatSummary {
50    fn default() -> Self {
51        StatSummary {
52            fun_y: SummaryFun::Mean,
53            fun_ymin: SummaryFun::Min,
54            fun_ymax: SummaryFun::Max,
55        }
56    }
57}
58
59impl StatSummary {
60    /// Create with mean_se default (mean +/- standard error).
61    pub fn mean_se() -> Self {
62        StatSummary {
63            fun_y: SummaryFun::Mean,
64            fun_ymin: SummaryFun::Min, // will be overridden in compute_group
65            fun_ymax: SummaryFun::Max,
66        }
67    }
68}
69
70impl Stat for StatSummary {
71    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
72        let x_col = match data.column("x") {
73            Some(c) => c,
74            None => return DataFrame::new(),
75        };
76        let y_col = match data.column("y") {
77            Some(c) => c,
78            None => return DataFrame::new(),
79        };
80
81        // Group y values by x
82        let mut groups: Vec<(String, Value, Vec<f64>)> = Vec::new();
83        for (x, y) in x_col.iter().zip(y_col.iter()) {
84            let key = x.to_group_key();
85            let y_val = y.as_f64().unwrap_or(0.0);
86            if let Some(entry) = groups.iter_mut().find(|(k, _, _)| k == &key) {
87                entry.2.push(y_val);
88            } else {
89                groups.push((key, x.clone(), vec![y_val]));
90            }
91        }
92
93        let n = groups.len();
94        let mut x_vals = Vec::with_capacity(n);
95        let mut y_vals = Vec::with_capacity(n);
96        let mut ymin_vals = Vec::with_capacity(n);
97        let mut ymax_vals = Vec::with_capacity(n);
98
99        for (_, x_val, ys) in &groups {
100            x_vals.push(x_val.clone());
101            y_vals.push(Value::Float(self.fun_y.apply(ys)));
102            ymin_vals.push(Value::Float(self.fun_ymin.apply(ys)));
103            ymax_vals.push(Value::Float(self.fun_ymax.apply(ys)));
104        }
105
106        let mut result = DataFrame::new();
107        result.add_column("x".to_string(), x_vals);
108        result.add_column("y".to_string(), y_vals);
109        result.add_column("ymin".to_string(), ymin_vals);
110        result.add_column("ymax".to_string(), ymax_vals);
111
112        // Carry over grouping columns
113        for col_name in &["color", "fill", "group"] {
114            if let Some(col) = data.column(col_name) {
115                if let Some(first) = col.first() {
116                    result.add_column(col_name.to_string(), vec![first.clone(); n]);
117                }
118            }
119        }
120
121        result
122    }
123
124    fn required_aes(&self) -> Vec<Aesthetic> {
125        vec![Aesthetic::X, Aesthetic::Y]
126    }
127
128    fn name(&self) -> &str {
129        "summary"
130    }
131}
ggplot_rs/stat/summary.rs

ggplot_rs/stat/
summary.rs
