ggplot_rs/stat/
xdensity_ridges.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Kernel density estimation on X per group, for ridgeline plots
8/// ([`GeomDensityRidges`](crate::geom::density_ridges::GeomDensityRidges)).
9/// Outputs: y (the group's value, e.g. a discrete label, kept as-is), x
10/// (evaluation points over the group's range ± 3 bandwidths, so each ridge
11/// falls to zero at its ends) and `height` (the density normalised to a peak
12/// of 1 per group).
13pub struct StatXDensityRidges {
14    pub n_points: usize,
15}
16
17impl Default for StatXDensityRidges {
18    fn default() -> Self {
19        StatXDensityRidges { n_points: 256 }
20    }
21}
22
23impl Stat for StatXDensityRidges {
24    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
25        let Some(x_col) = data.column("x") else {
26            return DataFrame::new();
27        };
28        let values: Vec<f64> = x_col
29            .iter()
30            .filter_map(|v| v.as_f64())
31            .filter(|f| f.is_finite())
32            .collect();
33        if values.len() < 2 {
34            return DataFrame::new();
35        }
36        let group_y = data
37            .column("y")
38            .and_then(|c| c.first())
39            .cloned()
40            .unwrap_or(Value::Float(0.0));
41        let n = values.len() as f64;
42        let bw = super::bw_nrd0(&values);
43        let lo = values.iter().cloned().fold(f64::INFINITY, f64::min) - 3.0 * bw;
44        let hi = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max) + 3.0 * bw;
45        let m = self.n_points.max(2);
46        let step = (hi - lo) / (m - 1) as f64;
47        let dens: Vec<(f64, f64)> = (0..m)
48            .map(|i| {
49                let x = lo + i as f64 * step;
50                let d = values
51                    .iter()
52                    .map(|xi| {
53                        let u = (x - xi) / bw;
54                        (-(u * u) / 2.0).exp()
55                    })
56                    .sum::<f64>()
57                    / (n * bw * (2.0 * std::f64::consts::PI).sqrt());
58                (x, d)
59            })
60            .collect();
61        let peak = dens.iter().map(|(_, d)| *d).fold(0.0_f64, f64::max);
62        let k = if peak > 0.0 { 1.0 / peak } else { 1.0 };
63        let mut result = DataFrame::new();
64        result.add_column(
65            "x".into(),
66            dens.iter().map(|(x, _)| Value::Float(*x)).collect(),
67        );
68        result.add_column("y".into(), vec![group_y; m]);
69        result.add_column(
70            "height".into(),
71            dens.iter().map(|(_, d)| Value::Float(d * k)).collect(),
72        );
73        for col_name in &["color", "fill", "group"] {
74            if let Some(first) = data.column(col_name).and_then(|c| c.first()) {
75                result.add_column(col_name.to_string(), vec![first.clone(); m]);
76            }
77        }
78        result
79    }
80
81    fn required_aes(&self) -> Vec<Aesthetic> {
82        vec![Aesthetic::X, Aesthetic::Y]
83    }
84
85    fn name(&self) -> &str {
86        "xdensity_ridges"
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn one_ridge_per_group_peaks_at_one_and_ends_near_zero() {
96        let mut data = DataFrame::new();
97        data.add_column(
98            "x".into(),
99            (0..40).map(|i| Value::Float((i % 10) as f64)).collect(),
100        );
101        data.add_column("y".into(), vec![Value::Str("A".into()); 40]);
102        let r = StatXDensityRidges::default().compute_group(&data, &ScaleSet::new());
103        let h: Vec<f64> = r
104            .column("height")
105            .unwrap()
106            .iter()
107            .filter_map(|v| v.as_f64())
108            .collect();
109        let peak = h.iter().cloned().fold(0.0_f64, f64::max);
110        assert!((peak - 1.0).abs() < 1e-9, "{peak}");
111        assert!(
112            h[0] < 0.05 && h[h.len() - 1] < 0.05,
113            "{} {}",
114            h[0],
115            h[h.len() - 1]
116        );
117        assert!(r
118            .column("y")
119            .unwrap()
120            .iter()
121            .all(|v| *v == Value::Str("A".into())));
122    }
123
124    #[test]
125    fn too_few_values_give_nothing() {
126        let mut data = DataFrame::new();
127        data.add_column("x".into(), vec![Value::Float(1.0)]);
128        data.add_column("y".into(), vec![Value::Str("A".into())]);
129        assert_eq!(
130            StatXDensityRidges::default()
131                .compute_group(&data, &ScaleSet::new())
132                .nrows(),
133            0
134        );
135    }
136}
ggplot_rs/stat/xdensity_ridges.rs

ggplot_rs/stat/
xdensity_ridges.rs
