1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7pub 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 let bw_x = (x_max - x_min) / self.bins_x as f64;
57 let bw_y = (y_max - y_min) / self.bins_y as f64;
58
59 let mut counts = vec![vec![0usize; self.bins_y]; self.bins_x];
60
61 for i in 0..n {
62 let bx = ((xs[i] - x_min) / bw_x).floor() as usize;
63 let by = ((ys[i] - y_min) / bw_y).floor() as usize;
64 let bx = bx.min(self.bins_x - 1);
65 let by = by.min(self.bins_y - 1);
66 counts[bx][by] += 1;
67 }
68
69 let mut xmin_vals = Vec::new();
70 let mut xmax_vals = Vec::new();
71 let mut ymin_vals = Vec::new();
72 let mut ymax_vals = Vec::new();
73 let mut fill_vals = Vec::new();
74
75 for (bx, row) in counts.iter().enumerate() {
76 for (by, &count) in row.iter().enumerate() {
77 if count == 0 {
78 continue;
79 }
80 let cell_xmin = x_min + bx as f64 * bw_x;
81 let cell_xmax = cell_xmin + bw_x;
82 let cell_ymin = y_min + by as f64 * bw_y;
83 let cell_ymax = cell_ymin + bw_y;
84
85 xmin_vals.push(Value::Float(cell_xmin));
86 xmax_vals.push(Value::Float(cell_xmax));
87 ymin_vals.push(Value::Float(cell_ymin));
88 ymax_vals.push(Value::Float(cell_ymax));
89 fill_vals.push(Value::Float(count as f64));
90 }
91 }
92
93 let mut result = DataFrame::new();
94 result.add_column("xmin".to_string(), xmin_vals);
95 result.add_column("xmax".to_string(), xmax_vals);
96 result.add_column("ymin".to_string(), ymin_vals);
97 result.add_column("ymax".to_string(), ymax_vals);
98 result.add_column("fill".to_string(), fill_vals);
99
100 result
101 }
102
103 fn required_aes(&self) -> Vec<Aesthetic> {
104 vec![Aesthetic::X, Aesthetic::Y]
105 }
106
107 fn name(&self) -> &str {
108 "bin2d"
109 }
110}
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115
116 #[test]
117 fn test_bin2d_basic() {
118 let mut data = DataFrame::new();
119 let x_vals: Vec<Value> = (0..100).map(|i| Value::Float(i as f64 / 10.0)).collect();
120 let y_vals: Vec<Value> = (0..100).map(|i| Value::Float(i as f64 / 5.0)).collect();
121 data.add_column("x".to_string(), x_vals);
122 data.add_column("y".to_string(), y_vals);
123
124 let stat = StatBin2d {
125 bins_x: 5,
126 bins_y: 5,
127 };
128 let scales = ScaleSet::new();
129 let result = stat.compute_group(&data, &scales);
130
131 assert!(result.nrows() > 0);
132 assert!(result.column("xmin").is_some());
133 assert!(result.column("xmax").is_some());
134 assert!(result.column("ymin").is_some());
135 assert!(result.column("ymax").is_some());
136 assert!(result.column("fill").is_some());
137 }
138}