Skip to main content

ggplot_rs/stat/
sum.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Count overlapping (x, y) pairs. Used by geom_count.
8/// Produces x, y, n (count) columns.
9pub struct StatSum;
10
11impl Default for StatSum {
12    fn default() -> Self {
13        StatSum
14    }
15}
16
17impl Stat for StatSum {
18    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
19        let x_col = match data.column("x") {
20            Some(c) => c,
21            None => return DataFrame::new(),
22        };
23        let y_col = match data.column("y") {
24            Some(c) => c,
25            None => return DataFrame::new(),
26        };
27
28        // Group by (x, y) key pairs
29        let mut groups: Vec<(String, Value, Value, usize)> = Vec::new();
30        for (x, y) in x_col.iter().zip(y_col.iter()) {
31            let key = format!("{}|{}", x.to_group_key(), y.to_group_key());
32            if let Some(entry) = groups.iter_mut().find(|(k, _, _, _)| k == &key) {
33                entry.3 += 1;
34            } else {
35                groups.push((key, x.clone(), y.clone(), 1));
36            }
37        }
38
39        let n = groups.len();
40        let mut x_vals = Vec::with_capacity(n);
41        let mut y_vals = Vec::with_capacity(n);
42        let mut n_vals = Vec::with_capacity(n);
43
44        for (_, x_val, y_val, count) in groups {
45            x_vals.push(x_val);
46            y_vals.push(y_val);
47            n_vals.push(Value::Float(count as f64));
48        }
49
50        let mut result = DataFrame::new();
51        result.add_column("x".to_string(), x_vals);
52        result.add_column("y".to_string(), y_vals);
53        result.add_column("n".to_string(), n_vals);
54
55        // Carry over grouping columns
56        for col_name in &["color", "fill", "group"] {
57            if let Some(col) = data.column(col_name) {
58                if let Some(first) = col.first() {
59                    result.add_column(col_name.to_string(), vec![first.clone(); n]);
60                }
61            }
62        }
63
64        result
65    }
66
67    fn required_aes(&self) -> Vec<Aesthetic> {
68        vec![Aesthetic::X, Aesthetic::Y]
69    }
70
71    fn name(&self) -> &str {
72        "sum"
73    }
74}