ggplot_rs/stat/
density.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Gaussian kernel density estimation with Silverman bandwidth.
8pub struct StatDensity {
9    pub n_points: usize,
10}
11
12impl Default for StatDensity {
13    fn default() -> Self {
14        StatDensity { n_points: 512 }
15    }
16}
17
18impl Stat for StatDensity {
19    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
20        let x_col = match data.column("x") {
21            Some(c) => c,
22            None => return DataFrame::new(),
23        };
24
25        let values: Vec<f64> = x_col.iter().filter_map(|v| v.as_f64()).collect();
26        if values.len() < 2 {
27            return DataFrame::new();
28        }
29
30        // R's bw.nrd0 (positive even for zero-spread data).
31        let bandwidth = super::bw_nrd0(&values);
32
33        let x_min = values.iter().cloned().fold(f64::INFINITY, f64::min) - 3.0 * bandwidth;
34        let x_max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max) + 3.0 * bandwidth;
35        let step = (x_max - x_min) / (self.n_points - 1) as f64;
36
37        // Exact O(n × n_points) evaluation for small inputs (bit-for-bit the
38        // historical output); linear binning + a truncated kernel table for
39        // large ones (O(n + n_points × window)), within ~2e-5 of the peak.
40        let dens = if values.len() >= BINNED_MIN_N {
41            binned_kde(&values, bandwidth, x_min, step, self.n_points)
42        } else {
43            None
44        }
45        .unwrap_or_else(|| exact_kde(&values, bandwidth, x_min, step, self.n_points));
46
47        let x_vals: Vec<Value> = (0..self.n_points)
48            .map(|i| Value::Float(x_min + i as f64 * step))
49            .collect();
50        let y_vals: Vec<Value> = dens.into_iter().map(Value::Float).collect();
51
52        let mut result = DataFrame::new();
53        result.add_column("x".to_string(), x_vals);
54        result.add_column("y".to_string(), y_vals);
55
56        // Carry over grouping columns
57        for col_name in &["color", "fill", "group"] {
58            if let Some(col) = data.column(col_name) {
59                if let Some(first) = col.first() {
60                    result.add_column(col_name.to_string(), vec![first.clone(); self.n_points]);
61                }
62            }
63        }
64
65        result
66    }
67
68    fn required_aes(&self) -> Vec<Aesthetic> {
69        vec![Aesthetic::X]
70    }
71
72    fn computed_y(&self) -> bool {
73        true
74    }
75
76    fn name(&self) -> &str {
77        "density"
78    }
79}
80
81/// Inputs with at least this many finite values use the binned estimator.
82const BINNED_MIN_N: usize = 2048;
83
84/// Exact Gaussian KDE evaluated on the grid `x_min + i·step`.
85fn exact_kde(values: &[f64], bw: f64, x_min: f64, step: f64, n_points: usize) -> Vec<f64> {
86    let n = values.len() as f64;
87    (0..n_points)
88        .map(|i| {
89            let x = x_min + i as f64 * step;
90            values
91                .iter()
92                .map(|xi| gaussian_kernel((x - xi) / bw))
93                .sum::<f64>()
94                / (n * bw)
95        })
96        .collect()
97}
98
99/// Binned Gaussian KDE (Wand 1994): linearly bin the data onto a fine grid
100/// whose spacing is at most `bw / 50` (and divides the output step), then
101/// convolve with a kernel table truncated at ±8 bandwidths. Returns `None` when
102/// the inputs aren't suitable (non-finite data/bandwidth, or a fine grid that
103/// would be unreasonably large) so the caller falls back to the exact sum.
104fn binned_kde(values: &[f64], bw: f64, x_min: f64, step: f64, n_points: usize) -> Option<Vec<f64>> {
105    const MAX_FINE: usize = 1 << 20;
106    if !(bw.is_finite() && bw > 0.0 && step.is_finite() && step > 0.0 && n_points >= 2) {
107        return None;
108    }
109    if values.iter().any(|v| !v.is_finite()) {
110        return None;
111    }
112    // Refinement: r fine cells per output step, fine spacing delta <= bw/50.
113    let r = (step * 50.0 / bw).ceil().max(1.0);
114    let m = (n_points - 1) as f64 * r + 1.0;
115    if !m.is_finite() || m > MAX_FINE as f64 {
116        return None;
117    }
118    let (r, m) = (r as usize, m as usize);
119    let delta = step / r as f64;
120
121    // Linear binning: split each point's unit mass between its two neighbours.
122    let mut w = vec![0.0f64; m];
123    for &x in values {
124        let pos = ((x - x_min) / delta).clamp(0.0, (m - 1) as f64);
125        let j = (pos.floor() as usize).min(m - 1);
126        let f = pos - j as f64;
127        w[j] += 1.0 - f;
128        if f > 0.0 && j + 1 < m {
129            w[j + 1] += f;
130        }
131    }
132
133    // Kernel weights by fine-grid offset, truncated at 8 bandwidths (< 1e-14).
134    let half = ((8.0 * bw / delta).ceil() as usize).min(m);
135    let table: Vec<f64> = (0..=half)
136        .map(|o| gaussian_kernel(o as f64 * delta / bw))
137        .collect();
138
139    let norm = values.len() as f64 * bw;
140    Some(
141        (0..n_points)
142            .map(|i| {
143                let c = i * r;
144                let lo = c.saturating_sub(half);
145                let hi = (c + half).min(m - 1);
146                let mut acc = 0.0;
147                for (j, wj) in w.iter().enumerate().take(hi + 1).skip(lo) {
148                    if *wj != 0.0 {
149                        acc += wj * table[j.abs_diff(c)];
150                    }
151                }
152                acc / norm
153            })
154            .collect(),
155    )
156}
157
158fn gaussian_kernel(x: f64) -> f64 {
159    (-(x * x) / 2.0).exp() / (2.0 * std::f64::consts::PI).sqrt()
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    fn frame(values: &[f64]) -> DataFrame {
167        let mut df = DataFrame::new();
168        df.add_column(
169            "x".into(),
170            values.iter().map(|v| Value::Float(*v)).collect(),
171        );
172        df
173    }
174
175    /// Deterministic, bimodal, heavy-ish tailed sample.
176    fn sample(n: usize) -> Vec<f64> {
177        let mut rng = crate::rng::SplitMix64::new(99);
178        (0..n)
179            .map(|i| {
180                // Box–Muller normal draws, two modes.
181                let u1 = rng.next_f64().max(1e-12);
182                let u2 = rng.next_f64();
183                let z = (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos();
184                if i % 3 == 0 {
185                    10.0 + 0.5 * z
186                } else {
187                    z * 2.0
188                }
189            })
190            .collect()
191    }
192
193    fn setup(values: &[f64]) -> (f64, f64, f64) {
194        let bw = crate::stat::bw_nrd0(values);
195        let lo = values.iter().cloned().fold(f64::INFINITY, f64::min) - 3.0 * bw;
196        let hi = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max) + 3.0 * bw;
197        (bw, lo, (hi - lo) / 511.0)
198    }
199
200    #[test]
201    fn binned_matches_exact_within_tight_tolerance() {
202        for n in [2048, 5000, 20_000] {
203            let v = sample(n);
204            let (bw, lo, step) = setup(&v);
205            let exact = exact_kde(&v, bw, lo, step, 512);
206            let binned = binned_kde(&v, bw, lo, step, 512).expect("binned path");
207            let peak = exact.iter().cloned().fold(0.0, f64::max);
208            let max_err = exact
209                .iter()
210                .zip(&binned)
211                .map(|(a, b)| (a - b).abs())
212                .fold(0.0, f64::max);
213            assert!(
214                max_err / peak < 5e-5,
215                "n={n}: relative error {}",
216                max_err / peak
217            );
218            // Both integrate to ~1 over the grid.
219            let area: f64 = binned.iter().sum::<f64>() * step;
220            assert!((area - 1.0).abs() < 1e-3, "area {area}");
221        }
222    }
223
224    #[test]
225    fn large_input_uses_binned_path_and_stays_close() {
226        let v = sample(10_000);
227        let out = StatDensity::default().compute_group(&frame(&v), &ScaleSet::new());
228        let ys: Vec<f64> = out
229            .column("y")
230            .unwrap()
231            .iter()
232            .map(|v| v.as_f64().unwrap())
233            .collect();
234        let (bw, lo, step) = setup(&v);
235        let exact = exact_kde(&v, bw, lo, step, 512);
236        let peak = exact.iter().cloned().fold(0.0, f64::max);
237        for (a, b) in exact.iter().zip(&ys) {
238            assert!((a - b).abs() / peak < 1e-4);
239        }
240    }
241
242    #[test]
243    fn binned_declines_unsuitable_inputs() {
244        assert!(binned_kde(&[1.0, f64::NAN], 0.5, 0.0, 0.01, 512).is_none());
245        assert!(binned_kde(&[1.0, 2.0], 0.0, 0.0, 0.01, 512).is_none());
246        // A tiny bandwidth over a huge range would need > 2^20 fine cells.
247        assert!(binned_kde(&[0.0, 1e9], 1e-6, 0.0, 1e9 / 511.0, 512).is_none());
248    }
249}
ggplot_rs/stat/density.rs

ggplot_rs/stat/
density.rs
