Skip to main content

ggplot_rs/geom/
step.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, LineStyle, Linetype};
7use crate::render::RenderError;
8use crate::scale::ScaleSet;
9use crate::stat::identity::StatIdentity;
10use crate::stat::Stat;
11use crate::theme::Theme;
12
13use super::{Geom, GeomParams};
14
15/// Step direction: horizontal-then-vertical or vertical-then-horizontal.
16pub enum StepDirection {
17    /// Draw horizontal first, then vertical (default).
18    Hv,
19    /// Draw vertical first, then horizontal.
20    Vh,
21    /// Step half-way between adjacent x values (ggplot2's `"mid"`).
22    Mid,
23}
24
25/// Step function line geometry.
26pub struct GeomStep {
27    pub color: (u8, u8, u8),
28    pub width: f64,
29    pub alpha: f64,
30    pub direction: StepDirection,
31}
32
33impl Default for GeomStep {
34    fn default() -> Self {
35        GeomStep {
36            color: (0, 0, 0),
37            width: 1.5,
38            alpha: 1.0,
39            direction: StepDirection::Hv,
40        }
41    }
42}
43
44impl Geom for GeomStep {
45    fn draw(
46        &self,
47        data: &DataFrame,
48        coord: &dyn Coord,
49        scales: &ScaleSet,
50        _theme: &Theme,
51        backend: &mut dyn DrawBackend,
52    ) -> Result<(), RenderError> {
53        let x_col = data
54            .column("x")
55            .ok_or(RenderError::MissingAesthetic("x".into()))?;
56        let y_col = data
57            .column("y")
58            .ok_or(RenderError::MissingAesthetic("y".into()))?;
59        let plot_area = backend.plot_area();
60        let x_scale = scales.get(&Aesthetic::X);
61        let y_scale = scales.get(&Aesthetic::Y);
62
63        // One step line per group (colour / group / linetype level), so e.g.
64        // a Kaplan–Meier curve per stratum or an ECDF per group.
65        for rows in super::support::row_groups(data) {
66            let mut raw: Vec<(f64, f64)> = rows
67                .iter()
68                .map(|&i| {
69                    let nx = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.0);
70                    let ny = y_scale.map(|s| s.map(&y_col[i])).unwrap_or(0.0);
71                    (nx, ny)
72                })
73                .collect();
74            // Sort by x (stable, so tied x keep their row order).
75            raw.sort_by(|a, b| a.0.total_cmp(&b.0));
76            let step_points = super::ribbon::stepped(&raw, &self.direction);
77
78            let points: Vec<(f64, f64)> = step_points
79                .iter()
80                .map(|&(nx, ny)| {
81                    let (px, py) = coord.transform((nx, ny), &plot_area);
82                    // Clamp non-finite coords (e.g. stat_ecdf's ±Inf padding, which
83                    // extends the step to the panel edge) to the panel border so the
84                    // flat segments draw to the edge instead of off-canvas.
85                    let px = if px.is_finite() {
86                        px
87                    } else if nx < 0.0 {
88                        plot_area.x
89                    } else {
90                        plot_area.x + plot_area.width
91                    };
92                    let py = if py.is_finite() {
93                        py
94                    } else if ny < 0.0 {
95                        plot_area.y + plot_area.height
96                    } else {
97                        plot_area.y
98                    };
99                    (px, py)
100                })
101                .collect();
102
103            let first = rows[0];
104            let line_color = data
105                .column("color")
106                .and_then(|cc| scales.map_color(&Aesthetic::Color, &cc[first]))
107                .unwrap_or(self.color);
108            let linetype = data
109                .column("linetype")
110                .and_then(|c| scales.map_linetype(&c[first]))
111                .unwrap_or(Linetype::Solid);
112            let series = super::series_key(data, first);
113            super::set_mark(backend, series.clone(), None, series, None);
114
115            if points.len() >= 2 {
116                backend.draw_line(
117                    &points,
118                    &LineStyle {
119                        color: line_color,
120                        alpha: self.alpha,
121                        width: self.width,
122                        linetype,
123                    },
124                )?;
125            }
126        }
127        super::clear_mark(backend);
128
129        Ok(())
130    }
131
132    fn required_aes(&self) -> Vec<Aesthetic> {
133        vec![Aesthetic::X, Aesthetic::Y]
134    }
135
136    fn default_stat(&self) -> Box<dyn Stat> {
137        Box::new(StatIdentity)
138    }
139
140    fn default_position(&self) -> Box<dyn Position> {
141        Box::new(PositionIdentity)
142    }
143
144    fn default_params(&self) -> GeomParams {
145        GeomParams::default()
146    }
147
148    fn name(&self) -> &str {
149        "step"
150    }
151
152    /// stat_ecdf pads the step with ±Inf (ggplot2), drawn to the panel edge.
153    fn allows_infinite(&self) -> bool {
154        true
155    }
156
157    fn set_series_color(&mut self, color: (u8, u8, u8)) {
158        self.color = color;
159    }
160}