Skip to main content

ggplot_rs/geom/
density_ridges.rs

1//! Ridgeline plot (`geom_density_ridges`, anofox-visualization#917): one
2//! density per group of a discrete y, drawn upward from the group's
3//! baseline and overlapping the next row — distributions over time.
4use std::collections::HashMap;
5
6use crate::aes::Aesthetic;
7use crate::coord::Coord;
8use crate::data::DataFrame;
9use crate::position::identity::PositionIdentity;
10use crate::position::Position;
11use crate::render::backend::{DrawBackend, RectStyle};
12use crate::render::RenderError;
13use crate::scale::ScaleSet;
14use crate::stat::xdensity_ridges::StatXDensityRidges;
15use crate::stat::Stat;
16use crate::theme::Theme;
17
18use super::{Geom, GeomParams};
19
20/// Ridgeline geometry: a filled density per group of a discrete y.
21pub struct GeomDensityRidges {
22    pub fill: (u8, u8, u8),
23    pub color: (u8, u8, u8),
24    pub alpha: f64,
25    pub line_width: f64,
26    /// A ridge's peak height in rows (1.0: it just reaches the next
27    /// baseline; more: ridges overlap). Give the y scale `scale - 0.5` rows
28    /// of room above the last level
29    /// ([`ScaleDiscrete::with_expand`](crate::scale::discrete::ScaleDiscrete::with_expand))
30    /// or the top ridge is clipped.
31    pub scale: f64,
32}
33
34impl Default for GeomDensityRidges {
35    fn default() -> Self {
36        GeomDensityRidges {
37            fill: (97, 156, 255),
38            color: (255, 255, 255),
39            alpha: 0.85,
40            line_width: 0.8,
41            scale: 1.6,
42        }
43    }
44}
45
46impl Geom for GeomDensityRidges {
47    fn draw(
48        &self,
49        data: &DataFrame,
50        coord: &dyn Coord,
51        scales: &ScaleSet,
52        _theme: &Theme,
53        backend: &mut dyn DrawBackend,
54    ) -> Result<(), RenderError> {
55        let x_col = data
56            .column("x")
57            .ok_or(RenderError::MissingAesthetic("x".into()))?;
58        let y_col = data
59            .column("y")
60            .ok_or(RenderError::MissingAesthetic("y".into()))?;
61        let h_col = data
62            .column("height")
63            .ok_or(RenderError::MissingAesthetic("height".into()))?;
64        let fill_col = data.column("fill");
65        let plot_area = backend.plot_area();
66        let x_scale = scales.get(&Aesthetic::X);
67        let y_scale = scales.get(&Aesthetic::Y);
68
69        // One row's height in the y scale's mapped space: the smallest gap
70        // between two neighbouring levels.
71        let mut pos: Vec<f64> = y_scale
72            .map(|s| s.breaks().iter().map(|(p, _)| *p).collect())
73            .unwrap_or_default();
74        pos.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
75        let row = pos
76            .windows(2)
77            .map(|w| (w[1] - w[0]).abs())
78            .filter(|d| *d > 0.0)
79            .fold(f64::INFINITY, f64::min);
80        let row = if row.is_finite() { row } else { 0.5 };
81
82        // The stat stacks one block of rows per group; collect them by y.
83        let mut order: Vec<String> = Vec::new();
84        let mut groups: HashMap<String, Vec<usize>> = HashMap::new();
85        for (i, yv) in y_col.iter().enumerate() {
86            let key = format!("{yv:?}");
87            groups
88                .entry(key.clone())
89                .or_insert_with(|| {
90                    order.push(key.clone());
91                    Vec::new()
92                })
93                .push(i);
94        }
95        // Back to front: the highest baseline first, so a lower ridge is
96        // drawn over the one above it.
97        let base = |idxs: &Vec<usize>| y_scale.map(|s| s.map(&y_col[idxs[0]])).unwrap_or(0.0);
98        order.sort_by(|a, b| {
99            base(&groups[b])
100                .partial_cmp(&base(&groups[a]))
101                .unwrap_or(std::cmp::Ordering::Equal)
102        });
103        for key in &order {
104            let idxs = &groups[key];
105            let first = idxs[0];
106            let ny = base(idxs);
107            let fill = fill_col
108                .and_then(|fc| scales.map_color(&Aesthetic::Fill, &fc[first]))
109                .unwrap_or(self.fill);
110            let group = Some(super::tip_value(&y_col[first])).filter(|s| !s.is_empty());
111            super::set_mark(
112                backend,
113                group.clone(),
114                group,
115                super::series_key(data, first),
116                None,
117            );
118            let mut top: Vec<(f64, f64)> = Vec::with_capacity(idxs.len() + 2);
119            let mut x0 = f64::NAN;
120            let mut x1 = f64::NAN;
121            for &i in idxs {
122                let nx = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.0);
123                let h = h_col[i].as_f64().unwrap_or(0.0).max(0.0);
124                if x0.is_nan() {
125                    x0 = nx;
126                }
127                x1 = nx;
128                top.push(coord.transform((nx, ny + h * self.scale * row), &plot_area));
129            }
130            if top.len() < 2 {
131                continue;
132            }
133            let mut poly = vec![coord.transform((x0, ny), &plot_area)];
134            poly.extend(top);
135            poly.push(coord.transform((x1, ny), &plot_area));
136            backend.draw_polygon(
137                &poly,
138                &RectStyle {
139                    fill: Some(fill),
140                    stroke: Some(self.color),
141                    stroke_width: self.line_width,
142                    alpha: self.alpha,
143                    clip: true,
144                    corner_radius: 0.0,
145                },
146            )?;
147        }
148        super::clear_mark(backend);
149        Ok(())
150    }
151
152    fn required_aes(&self) -> Vec<Aesthetic> {
153        vec![Aesthetic::X, Aesthetic::Y]
154    }
155
156    fn default_stat(&self) -> Box<dyn Stat> {
157        Box::new(StatXDensityRidges::default())
158    }
159
160    fn default_position(&self) -> Box<dyn Position> {
161        Box::new(PositionIdentity)
162    }
163
164    fn default_params(&self) -> GeomParams {
165        GeomParams::default()
166    }
167
168    fn name(&self) -> &str {
169        "density_ridges"
170    }
171
172    fn set_series_color(&mut self, color: (u8, u8, u8)) {
173        self.fill = color;
174    }
175}