ggplot_rs/stat/
density2d.rs

1use crate::aes::Aesthetic;
2use crate::data::DataFrame;
3use crate::scale::ScaleSet;
4
5use super::marching_squares::extract_contours;
6use super::Stat;
7
8/// Compute 2D density contour lines from (x, y) point data.
9/// Uses Gaussian 2D KDE with Silverman bandwidth selection.
10pub struct StatDensity2d {
11    pub n_grid: usize,
12    pub n_levels: usize,
13}
14
15impl Default for StatDensity2d {
16    fn default() -> Self {
17        StatDensity2d {
18            n_grid: 50,
19            n_levels: 10,
20        }
21    }
22}
23
24impl StatDensity2d {
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    pub fn with_grid(mut self, n: usize) -> Self {
30        self.n_grid = n;
31        self
32    }
33
34    pub fn with_levels(mut self, n: usize) -> Self {
35        self.n_levels = n;
36        self
37    }
38}
39
40impl Stat for StatDensity2d {
41    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
42        let x_col = match data.column("x") {
43            Some(c) => c,
44            None => return DataFrame::new(),
45        };
46        let y_col = match data.column("y") {
47            Some(c) => c,
48            None => return DataFrame::new(),
49        };
50
51        // Extract finite (x, y) points
52        let points: Vec<(f64, f64)> = x_col
53            .iter()
54            .zip(y_col.iter())
55            .filter_map(|(x, y)| match (x.as_f64(), y.as_f64()) {
56                (Some(xv), Some(yv)) if xv.is_finite() && yv.is_finite() => Some((xv, yv)),
57                _ => None,
58            })
59            .collect();
60
61        if points.len() < 2 {
62            return DataFrame::new();
63        }
64
65        let n = points.len() as f64;
66
67        // Compute means and standard deviations
68        let x_mean = points.iter().map(|p| p.0).sum::<f64>() / n;
69        let y_mean = points.iter().map(|p| p.1).sum::<f64>() / n;
70        let x_var = points.iter().map(|p| (p.0 - x_mean).powi(2)).sum::<f64>() / (n - 1.0);
71        let y_var = points.iter().map(|p| (p.1 - y_mean).powi(2)).sum::<f64>() / (n - 1.0);
72        let x_sd = x_var.sqrt();
73        let y_sd = y_var.sqrt();
74
75        if x_sd < f64::EPSILON || y_sd < f64::EPSILON {
76            return DataFrame::new();
77        }
78
79        // 2D Silverman bandwidth: bw = sd * n^(-1/6)
80        let bw_x = x_sd * n.powf(-1.0 / 6.0);
81        let bw_y = y_sd * n.powf(-1.0 / 6.0);
82
83        // Compute grid bounds (extend by 3*bw beyond data range)
84        let x_min = points.iter().map(|p| p.0).fold(f64::INFINITY, f64::min) - 3.0 * bw_x;
85        let x_max = points.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max) + 3.0 * bw_x;
86        let y_min = points.iter().map(|p| p.1).fold(f64::INFINITY, f64::min) - 3.0 * bw_y;
87        let y_max = points.iter().map(|p| p.1).fold(f64::NEG_INFINITY, f64::max) + 3.0 * bw_y;
88
89        let nx = self.n_grid;
90        let ny = self.n_grid;
91        let dx = (x_max - x_min) / nx as f64;
92        let dy = (y_max - y_min) / ny as f64;
93
94        // Build density grid via 2D Gaussian KDE
95        let mut grid = vec![0.0f64; (nx + 1) * (ny + 1)];
96        let inv_2bwx2 = 1.0 / (2.0 * bw_x * bw_x);
97        let inv_2bwy2 = 1.0 / (2.0 * bw_y * bw_y);
98        let norm = 1.0 / (2.0 * std::f64::consts::PI * bw_x * bw_y * n);
99
100        for iy in 0..=ny {
101            let gy = y_min + iy as f64 * dy;
102            for ix in 0..=nx {
103                let gx = x_min + ix as f64 * dx;
104                let mut density = 0.0;
105                for &(px, py) in &points {
106                    let dx2 = (gx - px) * (gx - px) * inv_2bwx2;
107                    let dy2 = (gy - py) * (gy - py) * inv_2bwy2;
108                    density += (-dx2 - dy2).exp();
109                }
110                grid[iy * (nx + 1) + ix] = density * norm;
111            }
112        }
113
114        // Determine density range for contour levels
115        let d_min = grid.iter().copied().fold(f64::INFINITY, f64::min);
116        let d_max = grid.iter().copied().fold(f64::NEG_INFINITY, f64::max);
117        let d_range = d_max - d_min;
118
119        if d_range.abs() < f64::EPSILON {
120            return DataFrame::new();
121        }
122
123        let n_levels = self.n_levels;
124        let levels: Vec<f64> = (1..=n_levels)
125            .map(|li| d_min + li as f64 * d_range / (n_levels + 1) as f64)
126            .collect();
127
128        let (x_vals, y_vals, level_vals, group_vals) =
129            extract_contours(&grid, nx, ny, x_min, y_min, dx, dy, &levels);
130
131        let mut result = DataFrame::new();
132        result.add_column("x".to_string(), x_vals);
133        result.add_column("y".to_string(), y_vals);
134        result.add_column("level".to_string(), level_vals);
135        result.add_column("group".to_string(), group_vals);
136        result
137    }
138
139    fn required_aes(&self) -> Vec<Aesthetic> {
140        vec![Aesthetic::X, Aesthetic::Y]
141    }
142
143    fn name(&self) -> &str {
144        "density_2d"
145    }
146}
ggplot_rs/stat/density2d.rs

ggplot_rs/stat/
density2d.rs
