Skip to main content

ggplot_rs/geom/
ribbon.rs

1use crate::aes::Aesthetic;
2use crate::coord::Coord;
3use crate::data::DataFrame;
4use crate::position::identity::PositionIdentity;
5use crate::position::Position;
6use crate::render::backend::{DrawBackend, RectStyle};
7use crate::render::RenderError;
8use crate::scale::ScaleSet;
9use crate::stat::identity::StatIdentity;
10use crate::stat::Stat;
11use crate::theme::Theme;
12
13use super::step::StepDirection;
14use super::{Geom, GeomParams};
15
16/// How a band's edges connect consecutive `(x, ymin, ymax)` rows.
17pub(crate) enum BandShape<'a> {
18    /// Straight segments (`geom_ribbon`).
19    Linear,
20    /// A step function (`geom_stepribbon`), like `geom_step`'s direction.
21    Step(&'a StepDirection),
22    /// A smooth curve through the points (`areasplinerange`).
23    Smooth(super::support::Interpolate),
24}
25
26/// Insert step corners into a run of normalized `(x, y)` points.
27pub(crate) fn stepped(raw: &[(f64, f64)], direction: &StepDirection) -> Vec<(f64, f64)> {
28    let mut out = Vec::with_capacity(raw.len() * 2);
29    for (j, &(x, y)) in raw.iter().enumerate() {
30        if j > 0 {
31            let (px, py) = raw[j - 1];
32            match direction {
33                StepDirection::Hv => out.push((x, py)),
34                StepDirection::Vh => out.push((px, y)),
35                StepDirection::Mid => {
36                    let m = 0.5 * (px + x);
37                    out.push((m, py));
38                    out.push((m, y));
39                }
40            }
41        }
42        out.push((x, y));
43    }
44    out
45}
46
47/// Draw one filled band per group (group / colour / fill level) between
48/// `ymin` and `ymax`, rows ordered by x. Fill comes from a mapped `fill`,
49/// else `fill`. Shared by `geom_ribbon`, `geom_stepribbon`, `geom_qq_band`.
50pub(crate) fn draw_bands(
51    data: &DataFrame,
52    coord: &dyn Coord,
53    scales: &ScaleSet,
54    backend: &mut dyn DrawBackend,
55    fill: (u8, u8, u8),
56    alpha: f64,
57    shape: BandShape<'_>,
58) -> Result<(), RenderError> {
59    let x_col = data
60        .column("x")
61        .ok_or(RenderError::MissingAesthetic("x".into()))?;
62    let ymin_col = data
63        .column("ymin")
64        .ok_or(RenderError::MissingAesthetic("ymin".into()))?;
65    let ymax_col = data
66        .column("ymax")
67        .ok_or(RenderError::MissingAesthetic("ymax".into()))?;
68    let plot_area = backend.plot_area();
69    let y_scale = scales.get(&Aesthetic::Y);
70    let xm = super::support::XMapper::new(data, scales);
71
72    for rows in super::support::row_groups(data) {
73        let mut pts: Vec<(f64, f64, f64)> = rows
74            .iter()
75            .map(|&i| {
76                (
77                    xm.map(&x_col[i], i),
78                    y_scale.map(|s| s.map(&ymin_col[i])).unwrap_or(0.0),
79                    y_scale.map(|s| s.map(&ymax_col[i])).unwrap_or(0.0),
80                )
81            })
82            .filter(|(x, a, b)| x.is_finite() && a.is_finite() && b.is_finite())
83            .collect();
84        if pts.len() < 2 {
85            continue;
86        }
87        pts.sort_by(|a, b| a.0.total_cmp(&b.0));
88        let upper: Vec<(f64, f64)> = pts.iter().map(|p| (p.0, p.2)).collect();
89        let lower: Vec<(f64, f64)> = pts.iter().map(|p| (p.0, p.1)).collect();
90        let (upper, mut lower) = match shape {
91            BandShape::Linear => (upper, lower),
92            BandShape::Step(dir) => (stepped(&upper, dir), stepped(&lower, dir)),
93            BandShape::Smooth(m) => (
94                super::support::interpolate(&upper, m).0,
95                super::support::interpolate(&lower, m).0,
96            ),
97        };
98        lower.reverse();
99        let polygon: Vec<(f64, f64)> = upper
100            .into_iter()
101            .chain(lower)
102            .map(|p| coord.transform(p, &plot_area))
103            .collect();
104
105        let first = rows[0];
106        let band_fill = data
107            .column("fill")
108            .and_then(|c| scales.map_color(&Aesthetic::Fill, &c[first]))
109            .unwrap_or(fill);
110        let series = super::series_key(data, first);
111        super::set_mark(backend, series.clone(), None, series, None);
112        super::set_mark_formatted(backend, scales, None, &Aesthetic::Y, None);
113        backend.draw_polygon(
114            &polygon,
115            &RectStyle {
116                fill: Some(band_fill),
117                stroke: None,
118                stroke_width: 0.0,
119                alpha,
120                clip: true,
121                corner_radius: 0.0,
122            },
123        )?;
124    }
125    super::clear_mark(backend);
126    Ok(())
127}
128
129/// Ribbon geometry — filled band between ymin and ymax (one per group).
130pub struct GeomRibbon {
131    pub fill: (u8, u8, u8),
132    pub alpha: f64,
133    /// How consecutive points are joined (straight by default; a smooth
134    /// `areasplinerange` with [`Interpolate::MonotoneCubic`](super::support::Interpolate::MonotoneCubic)).
135    pub interpolate: super::support::Interpolate,
136}
137
138impl Default for GeomRibbon {
139    fn default() -> Self {
140        GeomRibbon {
141            fill: (97, 156, 255),
142            alpha: 0.3,
143            interpolate: super::support::Interpolate::Linear,
144        }
145    }
146}
147
148impl Geom for GeomRibbon {
149    fn key_fill(&self) -> Option<((u8, u8, u8), f64)> {
150        Some((self.fill, self.alpha))
151    }
152
153    fn draw(
154        &self,
155        data: &DataFrame,
156        coord: &dyn Coord,
157        scales: &ScaleSet,
158        _theme: &Theme,
159        backend: &mut dyn DrawBackend,
160    ) -> Result<(), RenderError> {
161        draw_bands(
162            data,
163            coord,
164            scales,
165            backend,
166            self.fill,
167            self.alpha,
168            match self.interpolate {
169                super::support::Interpolate::Linear => BandShape::Linear,
170                m => BandShape::Smooth(m),
171            },
172        )
173    }
174
175    fn required_aes(&self) -> Vec<Aesthetic> {
176        vec![Aesthetic::X, Aesthetic::Ymin, Aesthetic::Ymax]
177    }
178
179    fn default_stat(&self) -> Box<dyn Stat> {
180        Box::new(StatIdentity)
181    }
182    fn default_position(&self) -> Box<dyn Position> {
183        Box::new(PositionIdentity)
184    }
185    fn default_params(&self) -> GeomParams {
186        GeomParams::default()
187    }
188    fn name(&self) -> &str {
189        "ribbon"
190    }
191
192    fn set_series_color(&mut self, color: (u8, u8, u8)) {
193        self.fill = color;
194    }
195}
196
197/// Step ribbon (pammtools' `geom_stepribbon`): a band whose edges are step
198/// functions — the confidence band of a Kaplan–Meier curve or an ECDF.
199/// `direction` matches `geom_step` (`Hv` default: the value holds until the
200/// next x). `±Inf` x (e.g. `stat_ecdf`'s padding) extends to the panel edge.
201pub struct GeomStepribbon {
202    pub fill: (u8, u8, u8),
203    pub alpha: f64,
204    pub direction: StepDirection,
205}
206
207impl Default for GeomStepribbon {
208    fn default() -> Self {
209        GeomStepribbon {
210            fill: (97, 156, 255),
211            alpha: 0.3,
212            direction: StepDirection::Hv,
213        }
214    }
215}
216
217impl GeomStepribbon {
218    /// A step ribbon with the given direction.
219    pub fn new(direction: StepDirection) -> Self {
220        GeomStepribbon {
221            direction,
222            ..Default::default()
223        }
224    }
225}
226
227impl Geom for GeomStepribbon {
228    fn key_fill(&self) -> Option<((u8, u8, u8), f64)> {
229        Some((self.fill, self.alpha))
230    }
231
232    fn draw(
233        &self,
234        data: &DataFrame,
235        coord: &dyn Coord,
236        scales: &ScaleSet,
237        _theme: &Theme,
238        backend: &mut dyn DrawBackend,
239    ) -> Result<(), RenderError> {
240        draw_bands(
241            data,
242            coord,
243            scales,
244            backend,
245            self.fill,
246            self.alpha,
247            BandShape::Step(&self.direction),
248        )
249    }
250
251    fn required_aes(&self) -> Vec<Aesthetic> {
252        vec![Aesthetic::X, Aesthetic::Ymin, Aesthetic::Ymax]
253    }
254
255    fn default_stat(&self) -> Box<dyn Stat> {
256        Box::new(StatIdentity)
257    }
258    fn default_position(&self) -> Box<dyn Position> {
259        Box::new(PositionIdentity)
260    }
261    fn default_params(&self) -> GeomParams {
262        GeomParams::default()
263    }
264    fn name(&self) -> &str {
265        "stepribbon"
266    }
267
268    fn allows_infinite(&self) -> bool {
269        true
270    }
271
272    fn set_series_color(&mut self, color: (u8, u8, u8)) {
273        self.fill = color;
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280
281    #[test]
282    fn step_corners() {
283        let raw = [(0.0, 1.0), (1.0, 2.0), (2.0, 3.0)];
284        assert_eq!(
285            stepped(&raw, &StepDirection::Hv),
286            vec![(0.0, 1.0), (1.0, 1.0), (1.0, 2.0), (2.0, 2.0), (2.0, 3.0)]
287        );
288        assert_eq!(
289            stepped(&raw, &StepDirection::Vh),
290            vec![(0.0, 1.0), (0.0, 2.0), (1.0, 2.0), (1.0, 3.0), (2.0, 3.0)]
291        );
292        assert_eq!(
293            stepped(&raw[..2], &StepDirection::Mid),
294            vec![(0.0, 1.0), (0.5, 1.0), (0.5, 2.0), (1.0, 2.0)]
295        );
296    }
297}