Skip to main content

ggplot_rs/stat/
bin2d.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// 2D rectangular binning. Divides x/y ranges into a grid, counts per cell.
8pub struct StatBin2d {
9    pub bins_x: usize,
10    pub bins_y: usize,
11}
12
13impl Default for StatBin2d {
14    fn default() -> Self {
15        StatBin2d {
16            bins_x: 30,
17            bins_y: 30,
18        }
19    }
20}
21
22impl Stat for StatBin2d {
23    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
24        let x_col = match data.column("x") {
25            Some(c) => c,
26            None => return DataFrame::new(),
27        };
28        let y_col = match data.column("y") {
29            Some(c) => c,
30            None => return DataFrame::new(),
31        };
32
33        let xs: Vec<f64> = x_col.iter().filter_map(|v| v.as_f64()).collect();
34        let ys: Vec<f64> = y_col.iter().filter_map(|v| v.as_f64()).collect();
35        let n = xs.len().min(ys.len());
36        if n == 0 {
37            return DataFrame::new();
38        }
39
40        let x_min = xs.iter().cloned().fold(f64::INFINITY, f64::min);
41        let x_max = xs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
42        let y_min = ys.iter().cloned().fold(f64::INFINITY, f64::min);
43        let y_max = ys.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
44
45        let (x_min, x_max) = if (x_max - x_min).abs() < f64::EPSILON {
46            (x_min - 0.5, x_max + 0.5)
47        } else {
48            (x_min, x_max)
49        };
50        let (y_min, y_max) = if (y_max - y_min).abs() < f64::EPSILON {
51            (y_min - 0.5, y_max + 0.5)
52        } else {
53            (y_min, y_max)
54        };
55
56        // Match ggplot2's geom_bin2d: width spans the range in `bins - 1` steps,
57        // with a bin *edge* aligned to 0 (boundary = 0), and right-closed cells.
58        let bw_x = (x_max - x_min) / (self.bins_x.max(2) - 1) as f64;
59        let bw_y = (y_max - y_min) / (self.bins_y.max(2) - 1) as f64;
60        let (x_origin, nbx) = super::bin::aligned_bins_at(x_min, x_max, bw_x, 0.0);
61        let (y_origin, nby) = super::bin::aligned_bins_at(y_min, y_max, bw_y, 0.0);
62
63        let mut counts = vec![vec![0usize; nby]; nbx];
64
65        for i in 0..n {
66            let bx =
67                (((xs[i] - x_origin) / bw_x).ceil() as i64 - 1).clamp(0, nbx as i64 - 1) as usize;
68            let by =
69                (((ys[i] - y_origin) / bw_y).ceil() as i64 - 1).clamp(0, nby as i64 - 1) as usize;
70            counts[bx][by] += 1;
71        }
72
73        let mut xmin_vals = Vec::new();
74        let mut xmax_vals = Vec::new();
75        let mut ymin_vals = Vec::new();
76        let mut ymax_vals = Vec::new();
77        let mut fill_vals = Vec::new();
78
79        for (bx, row) in counts.iter().enumerate() {
80            for (by, &count) in row.iter().enumerate() {
81                if count == 0 {
82                    continue;
83                }
84                let cell_xmin = x_origin + bx as f64 * bw_x;
85                let cell_xmax = cell_xmin + bw_x;
86                let cell_ymin = y_origin + by as f64 * bw_y;
87                let cell_ymax = cell_ymin + bw_y;
88
89                xmin_vals.push(Value::Float(cell_xmin));
90                xmax_vals.push(Value::Float(cell_xmax));
91                ymin_vals.push(Value::Float(cell_ymin));
92                ymax_vals.push(Value::Float(cell_ymax));
93                fill_vals.push(Value::Float(count as f64));
94            }
95        }
96
97        let mut result = DataFrame::new();
98        result.add_column("xmin".to_string(), xmin_vals);
99        result.add_column("xmax".to_string(), xmax_vals);
100        result.add_column("ymin".to_string(), ymin_vals);
101        result.add_column("ymax".to_string(), ymax_vals);
102        result.add_column("fill".to_string(), fill_vals);
103
104        result
105    }
106
107    fn required_aes(&self) -> Vec<Aesthetic> {
108        vec![Aesthetic::X, Aesthetic::Y]
109    }
110
111    fn name(&self) -> &str {
112        "bin2d"
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn test_bin2d_basic() {
122        let mut data = DataFrame::new();
123        let x_vals: Vec<Value> = (0..100).map(|i| Value::Float(i as f64 / 10.0)).collect();
124        let y_vals: Vec<Value> = (0..100).map(|i| Value::Float(i as f64 / 5.0)).collect();
125        data.add_column("x".to_string(), x_vals);
126        data.add_column("y".to_string(), y_vals);
127
128        let stat = StatBin2d {
129            bins_x: 5,
130            bins_y: 5,
131        };
132        let scales = ScaleSet::new();
133        let result = stat.compute_group(&data, &scales);
134
135        assert!(result.nrows() > 0);
136        assert!(result.column("xmin").is_some());
137        assert!(result.column("xmax").is_some());
138        assert!(result.column("ymin").is_some());
139        assert!(result.column("ymax").is_some());
140        assert!(result.column("fill").is_some());
141    }
142}