Skip to main content

ggplot_rs/geom/
waterfall.rs

1//! Waterfall (bridge) chart: `geom_waterfall`. Each step is a bar floating
2//! from the previous running total; total bars rise from 0. Bars are coloured
3//! through the fill scale on the stat's `direction` column (Increase /
4//! Decrease / Total), and thin connectors join each bar's end to the next bar.
5
6use crate::aes::Aesthetic;
7use crate::coord::Coord;
8use crate::data::{DataFrame, Value};
9use crate::position::identity::PositionIdentity;
10use crate::position::Position;
11use crate::render::backend::{DrawBackend, LineStyle, Linetype, RectStyle};
12use crate::render::RenderError;
13use crate::scale::ScaleSet;
14use crate::stat::waterfall::{StatWaterfall, WATERFALL_DECREASE, WATERFALL_TOTAL};
15use crate::stat::Stat;
16use crate::theme::Theme;
17
18use super::{Geom, GeomParams};
19
20/// Waterfall chart. Map `x` (the steps, in data order) and `y` (each step's
21/// change). List subtotal steps in `totals` (their `y` is ignored) and set
22/// `final_total` to append a closing total bar.
23pub struct GeomWaterfall {
24    /// Fill of rising steps (also the default `Increase` legend key).
25    pub up: (u8, u8, u8),
26    /// Fill of falling steps.
27    pub down: (u8, u8, u8),
28    /// Fill of total bars.
29    pub total: (u8, u8, u8),
30    /// Bar outline colour.
31    pub color: (u8, u8, u8),
32    /// x levels drawn as totals (0 → running total).
33    pub totals: Vec<String>,
34    /// Append a closing total bar with this x label.
35    pub final_total: Option<String>,
36    /// Draw dashed connectors from each bar's end to the next bar.
37    pub connectors: bool,
38    pub connector_color: (u8, u8, u8),
39    /// Bar width as a fraction of the category slot.
40    pub width: f64,
41    pub alpha: f64,
42}
43
44impl Default for GeomWaterfall {
45    fn default() -> Self {
46        GeomWaterfall {
47            up: (0x0c, 0xa6, 0x78),
48            down: (0xe0, 0x31, 0x31),
49            total: (0x4c, 0x6e, 0xa8),
50            color: (50, 50, 50),
51            totals: Vec::new(),
52            final_total: None,
53            connectors: true,
54            connector_color: (110, 110, 110),
55            width: 0.7,
56            alpha: 1.0,
57        }
58    }
59}
60
61impl GeomWaterfall {
62    /// Mark these x levels as subtotal bars.
63    pub fn with_totals<S: Into<String>>(mut self, levels: impl IntoIterator<Item = S>) -> Self {
64        self.totals = levels.into_iter().map(Into::into).collect();
65        self
66    }
67
68    /// Append a closing total bar labelled `label`.
69    pub fn with_final_total(mut self, label: impl Into<String>) -> Self {
70        self.final_total = Some(label.into());
71        self
72    }
73
74    /// Whether any total bar will be drawn (so the legend needs a Total key).
75    pub(crate) fn has_totals(&self) -> bool {
76        !self.totals.is_empty() || self.final_total.is_some()
77    }
78
79    fn direction_color(&self, dir: &Value) -> (u8, u8, u8) {
80        match dir.key_str().as_ref() {
81            WATERFALL_TOTAL => self.total,
82            WATERFALL_DECREASE => self.down,
83            _ => self.up,
84        }
85    }
86}
87
88impl Geom for GeomWaterfall {
89    fn draw(
90        &self,
91        data: &DataFrame,
92        coord: &dyn Coord,
93        scales: &ScaleSet,
94        _theme: &Theme,
95        backend: &mut dyn DrawBackend,
96    ) -> Result<(), RenderError> {
97        let col = |n: &str| {
98            data.column(n)
99                .ok_or_else(|| RenderError::MissingAesthetic(n.into()))
100        };
101        let (x_col, y_col, ymin_col, ymax_col) = (col("x")?, col("y")?, col("ymin")?, col("ymax")?);
102        let delta_col = data.column("delta");
103        let dir_col = data.column("direction");
104        let fill_col = data.column("fill");
105        let area = backend.plot_area();
106        let x_scale = scales.get(&Aesthetic::X);
107        let y_scale = scales.get(&Aesthetic::Y);
108        let my = |v: f64| y_scale.map(|s| s.map(&Value::Float(v))).unwrap_or(0.0);
109
110        let half = if x_scale.map(|s| s.is_discrete()).unwrap_or(false) {
111            let n = x_scale.map(|s| s.breaks().len()).unwrap_or(1).max(1) as f64;
112            let pad = if coord.discrete_expand() { 1.1 } else { 1.0 };
113            self.width / (n * pad) / 2.0
114        } else {
115            super::continuous_bar_half_width(
116                x_col.iter().filter_map(|v| x_scale.map(|s| s.map(v))),
117                self.width,
118                0.02,
119            )
120        };
121
122        // (nx, running total after the step) of the previous bar, for connectors.
123        let mut prev: Option<(f64, f64)> = None;
124        for i in 0..data.nrows() {
125            let (Some(lo), Some(hi), Some(end)) = (
126                ymin_col[i].as_f64(),
127                ymax_col[i].as_f64(),
128                y_col[i].as_f64(),
129            ) else {
130                continue;
131            };
132            let nx = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.5);
133            let (nlo, nhi, nend) = (my(lo), my(hi), my(end));
134
135            if self.connectors {
136                if let Some((px, pend)) = prev {
137                    super::clear_mark(backend);
138                    let a = coord.transform((px + half, pend), &area);
139                    let b = coord.transform((nx - half, pend), &area);
140                    backend.draw_line(
141                        &[a, b],
142                        &LineStyle {
143                            color: self.connector_color,
144                            alpha: self.alpha,
145                            width: 0.8,
146                            linetype: Linetype::Dashed,
147                            ..Default::default()
148                        },
149                    )?;
150                }
151            }
152            prev = Some((nx, nend));
153
154            let dir = dir_col.map(|c| c[i].clone()).unwrap_or(Value::Na);
155            let fill = fill_col
156                .and_then(|c| scales.map_color(&Aesthetic::Fill, &c[i]))
157                .unwrap_or_else(|| self.direction_color(&dir));
158            let delta = delta_col.map(|c| &c[i]).unwrap_or(&y_col[i]);
159            let x_txt = super::tip_value(&x_col[i]);
160            let tip = if dir.key_str() == WATERFALL_TOTAL {
161                format!("{x_txt}: {}", super::tip_value(delta))
162            } else {
163                let d = delta.as_f64().unwrap_or(0.0);
164                let sign = if d >= 0.0 { "+" } else { "" };
165                format!(
166                    "{x_txt}: {sign}{} (total {})",
167                    super::tip_value(delta),
168                    super::tip_value(&y_col[i])
169                )
170            };
171            super::set_mark(
172                backend,
173                Some(tip),
174                Some(x_txt),
175                super::raw_value(&dir).or_else(|| super::series_key(data, i)),
176                super::raw_value(delta),
177            );
178            super::set_mark_formatted(backend, scales, Some(&x_col[i]), &Aesthetic::Y, Some(delta));
179            super::set_row_extras(backend, data, i);
180
181            let (x0, y0) = coord.transform((nx - half, nhi), &area);
182            let (x1, y1) = coord.transform((nx + half, nlo), &area);
183            backend.draw_rect(
184                (x0.min(x1), y0.min(y1)),
185                (x0.max(x1), y0.max(y1)),
186                &RectStyle {
187                    fill: Some(fill),
188                    stroke: Some(self.color),
189                    stroke_width: 0.5,
190                    alpha: self.alpha,
191                    clip: true,
192                    corner_radius: 0.0,
193                },
194            )?;
195        }
196        super::clear_mark(backend);
197        super::clear_row_extras(backend);
198        Ok(())
199    }
200
201    fn required_aes(&self) -> Vec<Aesthetic> {
202        vec![Aesthetic::X, Aesthetic::Y]
203    }
204
205    fn default_stat(&self) -> Box<dyn Stat> {
206        Box::new(StatWaterfall {
207            totals: self.totals.clone(),
208            final_total: self.final_total.clone(),
209        })
210    }
211
212    fn default_position(&self) -> Box<dyn Position> {
213        Box::new(PositionIdentity)
214    }
215
216    fn default_params(&self) -> GeomParams {
217        GeomParams::default()
218    }
219
220    fn name(&self) -> &str {
221        "waterfall"
222    }
223
224    fn apply_element_geom(&mut self, el: &crate::theme::elements::ElementGeom) {
225        if let (true, Some(ink)) = (self.color == (50, 50, 50), el.ink) {
226            self.color = ink;
227        }
228    }
229
230    fn consumes_tooltips(&self) -> bool {
231        true
232    }
233
234    fn include_zero_baseline(&self) -> bool {
235        true
236    }
237}