ggplot_rs/stat/
ydensity.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Kernel density estimation on Y per group (for violin plots).
8/// Outputs: x (group value), y (eval points), violinwidth (density normalized to
9/// [0, 1]). The geom mirrors `violinwidth` around the group's x slot.
10///
11/// Also emits `density` (the raw estimate), `n` (the group's count) — which
12/// the geom's `scale = "area" | "count"` need — and, when `quantiles` is set,
13/// a `quantiles` cell with the group's sample quantiles (`"q1,q2,…"`).
14#[derive(Clone, Debug, PartialEq)]
15pub struct StatYDensity {
16    pub n_points: usize,
17    /// Evaluate over the data range only (`trim = TRUE`, default); `false`
18    /// extends ±3 bandwidths into the tails.
19    pub trim: bool,
20    /// Fixed kernel bandwidth; `None` = R's `bw.nrd0`.
21    pub bw: Option<f64>,
22    /// Bandwidth multiplier (`adjust`, default 1).
23    pub adjust: f64,
24    /// Sample quantiles to report (`quantiles`, e.g. `[0.25, 0.5, 0.75]`).
25    pub quantiles: Vec<f64>,
26}
27
28impl Default for StatYDensity {
29    fn default() -> Self {
30        StatYDensity {
31            n_points: 512,
32            trim: true,
33            bw: None,
34            adjust: 1.0,
35            quantiles: Vec::new(),
36        }
37    }
38}
39
40/// R's default (type 7) sample quantile of sorted `v`.
41fn quantile7(v: &[f64], p: f64) -> f64 {
42    let h = (v.len() - 1) as f64 * p.clamp(0.0, 1.0);
43    let (lo, hi) = (h.floor() as usize, h.ceil() as usize);
44    v[lo] + (h - lo as f64) * (v[hi] - v[lo])
45}
46
47impl Stat for StatYDensity {
48    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
49        let x_col = data.column("x");
50        let y_col = match data.column("y") {
51            Some(c) => c,
52            None => return DataFrame::new(),
53        };
54
55        let values: Vec<f64> = y_col.iter().filter_map(|v| v.as_f64()).collect();
56        if values.len() < 2 {
57            return DataFrame::new();
58        }
59
60        // Keep the group's x *value* as-is (e.g. the discrete label "A"). The geom
61        // maps it through the X scale, exactly like boxplot — converting to f64 here
62        // would collapse every discrete group to 0.0.
63        let group_x = x_col
64            .and_then(|c| c.first())
65            .cloned()
66            .unwrap_or(Value::Float(0.0));
67
68        let n = values.len() as f64;
69        // R's bw.nrd0 (positive even for zero-spread data), or a fixed `bw`,
70        // times `adjust`.
71        let adjust = if self.adjust.is_finite() && self.adjust > 0.0 {
72            self.adjust
73        } else {
74            1.0
75        };
76        let bandwidth = self
77            .bw
78            .filter(|b| b.is_finite() && *b > 0.0)
79            .unwrap_or_else(|| super::bw_nrd0(&values))
80            * adjust;
81
82        // ggplot2's geom_violin defaults to trim = TRUE: evaluate the density
83        // over the observed data range; trim = FALSE extends ±3 bandwidths.
84        let tail = if self.trim { 0.0 } else { 3.0 * bandwidth };
85        let y_min = values.iter().cloned().fold(f64::INFINITY, f64::min) - tail;
86        let y_max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max) + tail;
87        let step = (y_max - y_min) / (self.n_points - 1) as f64;
88
89        let mut x_vals = Vec::with_capacity(self.n_points);
90        let mut y_vals = Vec::with_capacity(self.n_points);
91
92        // Compute density at each evaluation point
93        let mut densities = Vec::with_capacity(self.n_points);
94        let mut max_density: f64 = 0.0;
95        for i in 0..self.n_points {
96            let y = y_min + i as f64 * step;
97            let density: f64 = values
98                .iter()
99                .map(|yi| gaussian_kernel((y - yi) / bandwidth))
100                .sum::<f64>()
101                / (n * bandwidth);
102            densities.push((y, density));
103            if density > max_density {
104                max_density = density;
105            }
106        }
107
108        // Normalize density to [0, 1] (peak = 1). The geom scales this by the
109        // per-group slot half-width, so the widest point fills the group's slot.
110        let scale = if max_density > 0.0 {
111            1.0 / max_density
112        } else {
113            1.0
114        };
115
116        let mut width_vals = Vec::with_capacity(self.n_points);
117        let mut density_vals = Vec::with_capacity(self.n_points);
118        for (y, density) in &densities {
119            x_vals.push(group_x.clone());
120            y_vals.push(Value::Float(*y));
121            width_vals.push(Value::Float(density * scale));
122            density_vals.push(Value::Float(*density));
123        }
124
125        let mut result = DataFrame::new();
126        result.add_column("x".to_string(), x_vals);
127        result.add_column("y".to_string(), y_vals);
128        result.add_column("violinwidth".to_string(), width_vals);
129        result.add_column("density".to_string(), density_vals);
130        result.add_column("n".to_string(), vec![Value::Float(n); self.n_points]);
131        if !self.quantiles.is_empty() {
132            let mut sorted = values.clone();
133            sorted.sort_by(f64::total_cmp);
134            let q = self
135                .quantiles
136                .iter()
137                .map(|&p| quantile7(&sorted, p).to_string())
138                .collect::<Vec<_>>()
139                .join(",");
140            result.add_column("quantiles".to_string(), vec![Value::Str(q); self.n_points]);
141        }
142
143        // Carry over grouping columns
144        for col_name in &["color", "fill", "group"] {
145            if let Some(col) = data.column(col_name) {
146                if let Some(first) = col.first() {
147                    result.add_column(col_name.to_string(), vec![first.clone(); self.n_points]);
148                }
149            }
150        }
151
152        result
153    }
154
155    fn required_aes(&self) -> Vec<Aesthetic> {
156        vec![Aesthetic::X, Aesthetic::Y]
157    }
158
159    fn name(&self) -> &str {
160        "ydensity"
161    }
162}
163
164fn gaussian_kernel(x: f64) -> f64 {
165    (-(x * x) / 2.0).exp() / (2.0 * std::f64::consts::PI).sqrt()
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[test]
173    fn test_ydensity_basic() {
174        let mut data = DataFrame::new();
175        data.add_column("x".to_string(), vec![Value::Float(1.0); 50]);
176        let y_vals: Vec<Value> = (0..50).map(|i| Value::Float(i as f64)).collect();
177        data.add_column("y".to_string(), y_vals);
178
179        let stat = StatYDensity::default();
180        let scales = ScaleSet::new();
181        let result = stat.compute_group(&data, &scales);
182
183        assert!(result.nrows() > 0);
184        assert!(result.column("x").is_some());
185        assert!(result.column("y").is_some());
186        assert!(result.column("violinwidth").is_some());
187        // Normalized width peaks at 1.0.
188        let max_w = result
189            .column("violinwidth")
190            .unwrap()
191            .iter()
192            .filter_map(|v| v.as_f64())
193            .fold(0.0_f64, f64::max);
194        assert!(
195            (max_w - 1.0).abs() < 1e-9,
196            "peak width should be 1.0, got {max_w}"
197        );
198    }
199}
ggplot_rs/stat/ydensity.rs

ggplot_rs/stat/
ydensity.rs
