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 name(&self) -> &str {
73        "density"
74    }
75}
76
77/// Inputs with at least this many finite values use the binned estimator.
78const BINNED_MIN_N: usize = 2048;
79
80/// Exact Gaussian KDE evaluated on the grid `x_min + i·step`.
81fn exact_kde(values: &[f64], bw: f64, x_min: f64, step: f64, n_points: usize) -> Vec<f64> {
82    let n = values.len() as f64;
83    (0..n_points)
84        .map(|i| {
85            let x = x_min + i as f64 * step;
86            values
87                .iter()
88                .map(|xi| gaussian_kernel((x - xi) / bw))
89                .sum::<f64>()
90                / (n * bw)
91        })
92        .collect()
93}
94
95/// Binned Gaussian KDE (Wand 1994): linearly bin the data onto a fine grid
96/// whose spacing is at most `bw / 50` (and divides the output step), then
97/// convolve with a kernel table truncated at ±8 bandwidths. Returns `None` when
98/// the inputs aren't suitable (non-finite data/bandwidth, or a fine grid that
99/// would be unreasonably large) so the caller falls back to the exact sum.
100fn binned_kde(values: &[f64], bw: f64, x_min: f64, step: f64, n_points: usize) -> Option<Vec<f64>> {
101    const MAX_FINE: usize = 1 << 20;
102    if !(bw.is_finite() && bw > 0.0 && step.is_finite() && step > 0.0 && n_points >= 2) {
103        return None;
104    }
105    if values.iter().any(|v| !v.is_finite()) {
106        return None;
107    }
108    // Refinement: r fine cells per output step, fine spacing delta <= bw/50.
109    let r = (step * 50.0 / bw).ceil().max(1.0);
110    let m = (n_points - 1) as f64 * r + 1.0;
111    if !m.is_finite() || m > MAX_FINE as f64 {
112        return None;
113    }
114    let (r, m) = (r as usize, m as usize);
115    let delta = step / r as f64;
116
117    // Linear binning: split each point's unit mass between its two neighbours.
118    let mut w = vec![0.0f64; m];
119    for &x in values {
120        let pos = ((x - x_min) / delta).clamp(0.0, (m - 1) as f64);
121        let j = (pos.floor() as usize).min(m - 1);
122        let f = pos - j as f64;
123        w[j] += 1.0 - f;
124        if f > 0.0 && j + 1 < m {
125            w[j + 1] += f;
126        }
127    }
128
129    // Kernel weights by fine-grid offset, truncated at 8 bandwidths (< 1e-14).
130    let half = ((8.0 * bw / delta).ceil() as usize).min(m);
131    let table: Vec<f64> = (0..=half)
132        .map(|o| gaussian_kernel(o as f64 * delta / bw))
133        .collect();
134
135    let norm = values.len() as f64 * bw;
136    Some(
137        (0..n_points)
138            .map(|i| {
139                let c = i * r;
140                let lo = c.saturating_sub(half);
141                let hi = (c + half).min(m - 1);
142                let mut acc = 0.0;
143                for (j, wj) in w.iter().enumerate().take(hi + 1).skip(lo) {
144                    if *wj != 0.0 {
145                        acc += wj * table[j.abs_diff(c)];
146                    }
147                }
148                acc / norm
149            })
150            .collect(),
151    )
152}
153
154fn gaussian_kernel(x: f64) -> f64 {
155    (-(x * x) / 2.0).exp() / (2.0 * std::f64::consts::PI).sqrt()
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    fn frame(values: &[f64]) -> DataFrame {
163        let mut df = DataFrame::new();
164        df.add_column(
165            "x".into(),
166            values.iter().map(|v| Value::Float(*v)).collect(),
167        );
168        df
169    }
170
171    /// Deterministic, bimodal, heavy-ish tailed sample.
172    fn sample(n: usize) -> Vec<f64> {
173        let mut rng = crate::rng::SplitMix64::new(99);
174        (0..n)
175            .map(|i| {
176                // Box–Muller normal draws, two modes.
177                let u1 = rng.next_f64().max(1e-12);
178                let u2 = rng.next_f64();
179                let z = (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos();
180                if i % 3 == 0 {
181                    10.0 + 0.5 * z
182                } else {
183                    z * 2.0
184                }
185            })
186            .collect()
187    }
188
189    fn setup(values: &[f64]) -> (f64, f64, f64) {
190        let bw = crate::stat::bw_nrd0(values);
191        let lo = values.iter().cloned().fold(f64::INFINITY, f64::min) - 3.0 * bw;
192        let hi = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max) + 3.0 * bw;
193        (bw, lo, (hi - lo) / 511.0)
194    }
195
196    #[test]
197    fn binned_matches_exact_within_tight_tolerance() {
198        for n in [2048, 5000, 20_000] {
199            let v = sample(n);
200            let (bw, lo, step) = setup(&v);
201            let exact = exact_kde(&v, bw, lo, step, 512);
202            let binned = binned_kde(&v, bw, lo, step, 512).expect("binned path");
203            let peak = exact.iter().cloned().fold(0.0, f64::max);
204            let max_err = exact
205                .iter()
206                .zip(&binned)
207                .map(|(a, b)| (a - b).abs())
208                .fold(0.0, f64::max);
209            assert!(
210                max_err / peak < 5e-5,
211                "n={n}: relative error {}",
212                max_err / peak
213            );
214            // Both integrate to ~1 over the grid.
215            let area: f64 = binned.iter().sum::<f64>() * step;
216            assert!((area - 1.0).abs() < 1e-3, "area {area}");
217        }
218    }
219
220    #[test]
221    fn large_input_uses_binned_path_and_stays_close() {
222        let v = sample(10_000);
223        let out = StatDensity::default().compute_group(&frame(&v), &ScaleSet::new());
224        let ys: Vec<f64> = out
225            .column("y")
226            .unwrap()
227            .iter()
228            .map(|v| v.as_f64().unwrap())
229            .collect();
230        let (bw, lo, step) = setup(&v);
231        let exact = exact_kde(&v, bw, lo, step, 512);
232        let peak = exact.iter().cloned().fold(0.0, f64::max);
233        for (a, b) in exact.iter().zip(&ys) {
234            assert!((a - b).abs() / peak < 1e-4);
235        }
236    }
237
238    #[test]
239    fn binned_declines_unsuitable_inputs() {
240        assert!(binned_kde(&[1.0, f64::NAN], 0.5, 0.0, 0.01, 512).is_none());
241        assert!(binned_kde(&[1.0, 2.0], 0.0, 0.0, 0.01, 512).is_none());
242        // A tiny bandwidth over a huge range would need > 2^20 fine cells.
243        assert!(binned_kde(&[0.0, 1e9], 1e-6, 0.0, 1e9 / 511.0, 512).is_none());
244    }
245}
ggplot_rs/stat/density.rs

ggplot_rs/stat/
density.rs
