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