ggplot_rs/stat/
waterfall.rs

1//! `StatWaterfall`: running totals for a waterfall (bridge) chart.
2//!
3//! Walks the panel's rows in data order, accumulating `y`. Each step becomes a
4//! bar floating from the previous running total (`ymin`/`ymax`); rows whose x
5//! level is listed in `totals` instead draw from 0 to the running total (their
6//! own `y` is ignored), and `final_total` appends one closing total bar.
7
8use crate::aes::Aesthetic;
9use crate::data::{DataFrame, Value};
10use crate::scale::ScaleSet;
11
12use super::Stat;
13
14/// Legend / `direction` label of a rising step.
15pub const WATERFALL_INCREASE: &str = "Increase";
16/// Legend / `direction` label of a falling step.
17pub const WATERFALL_DECREASE: &str = "Decrease";
18/// Legend / `direction` label of a total bar.
19pub const WATERFALL_TOTAL: &str = "Total";
20
21/// Cumulative-sum stat behind [`GeomWaterfall`](crate::geom::waterfall::GeomWaterfall).
22///
23/// Output columns: `x`, `y` (running total after the step), `ymin`/`ymax`
24/// (bar extent), `start` (running total before the step), `delta` (the step's
25/// own value; the total for total bars) and `direction`
26/// ([`WATERFALL_INCREASE`] / [`WATERFALL_DECREASE`] / [`WATERFALL_TOTAL`]).
27/// Every other input column is carried through row by row (`Na` on an
28/// appended final total).
29#[derive(Clone, Debug, Default)]
30pub struct StatWaterfall {
31    /// x levels drawn as total bars (0 → running total); their `y` is ignored.
32    pub totals: Vec<String>,
33    /// Append a closing total bar with this x label.
34    pub final_total: Option<String>,
35}
36
37impl Stat for StatWaterfall {
38    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
39        let (Some(x_col), Some(y_col)) = (data.column("x"), data.column("y")) else {
40            return DataFrame::new();
41        };
42        let n = data.nrows();
43        let mut cum = 0.0;
44        let mut out: Vec<(Value, [f64; 5], &'static str)> = Vec::with_capacity(n + 1);
45        let mut kept = Vec::with_capacity(n);
46        for i in 0..n {
47            let is_total = self.totals.iter().any(|t| *t == x_col[i].key_str());
48            if is_total {
49                out.push((x_col[i].clone(), [0.0, cum, 0.0, cum, cum], WATERFALL_TOTAL));
50            } else {
51                let Some(d) = y_col[i].as_f64().filter(|d| d.is_finite()) else {
52                    continue;
53                };
54                let start = cum;
55                cum += d;
56                let dir = if d >= 0.0 {
57                    WATERFALL_INCREASE
58                } else {
59                    WATERFALL_DECREASE
60                };
61                out.push((
62                    x_col[i].clone(),
63                    [start, cum, start.min(cum), start.max(cum), d],
64                    dir,
65                ));
66            }
67            kept.push(Some(i));
68        }
69        if let Some(label) = &self.final_total {
70            out.push((
71                Value::Str(label.clone()),
72                [0.0, cum, 0.0, cum, cum],
73                WATERFALL_TOTAL,
74            ));
75            kept.push(None);
76        }
77
78        let mut result = DataFrame::new();
79        let num = |k: usize| out.iter().map(|r| Value::Float(r.1[k])).collect::<Vec<_>>();
80        result.add_column("x".into(), out.iter().map(|r| r.0.clone()).collect());
81        result.add_column("y".into(), num(1));
82        result.add_column("ymin".into(), num(2));
83        result.add_column("ymax".into(), num(3));
84        result.add_column("start".into(), num(0));
85        result.add_column("delta".into(), num(4));
86        result.add_column(
87            "direction".into(),
88            out.iter().map(|r| Value::Str(r.2.to_string())).collect(),
89        );
90        for name in data.column_names() {
91            if result.has_column(name) {
92                continue;
93            }
94            let col = data.column(name).unwrap_or_default();
95            let vals = kept
96                .iter()
97                .map(|k| k.map(|i| col[i].clone()).unwrap_or(Value::Na))
98                .collect();
99            result.add_column(name.to_string(), vals);
100        }
101        result
102    }
103
104    fn required_aes(&self) -> Vec<Aesthetic> {
105        vec![Aesthetic::X, Aesthetic::Y]
106    }
107
108    /// The running total spans every step of the panel, not one aesthetic
109    /// group (a discrete x would otherwise make each step its own group).
110    fn panelwise(&self) -> bool {
111        true
112    }
113
114    fn name(&self) -> &str {
115        "waterfall"
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    fn df(xs: &[&str], ys: &[f64]) -> DataFrame {
124        let mut d = DataFrame::new();
125        d.add_column(
126            "x".into(),
127            xs.iter().map(|s| Value::Str(s.to_string())).collect(),
128        );
129        d.add_column("y".into(), ys.iter().map(|v| Value::Float(*v)).collect());
130        d
131    }
132
133    fn col(d: &DataFrame, n: &str) -> Vec<f64> {
134        d.column(n)
135            .unwrap()
136            .iter()
137            .map(|v| v.as_f64().unwrap())
138            .collect()
139    }
140
141    #[test]
142    fn steps_float_from_the_running_total() {
143        let stat = StatWaterfall {
144            totals: vec!["Sub".into()],
145            final_total: Some("End".into()),
146        };
147        let out = stat.compute_group(
148            &df(&["A", "B", "Sub", "C"], &[10.0, -4.0, 0.0, 3.0]),
149            &ScaleSet::new(),
150        );
151        assert_eq!(out.nrows(), 5);
152        assert_eq!(col(&out, "ymin"), vec![0.0, 6.0, 0.0, 6.0, 0.0]);
153        assert_eq!(col(&out, "ymax"), vec![10.0, 10.0, 6.0, 9.0, 9.0]);
154        assert_eq!(col(&out, "y"), vec![10.0, 6.0, 6.0, 9.0, 9.0]);
155        assert_eq!(col(&out, "delta"), vec![10.0, -4.0, 6.0, 3.0, 9.0]);
156        let dirs: Vec<String> = out
157            .column("direction")
158            .unwrap()
159            .iter()
160            .map(|v| v.key_str().into_owned())
161            .collect();
162        assert_eq!(dirs, ["Increase", "Decrease", "Total", "Increase", "Total"]);
163        assert_eq!(out.column("x").unwrap()[4], Value::Str("End".into()));
164    }
165}
ggplot_rs/stat/waterfall.rs

ggplot_rs/stat/
waterfall.rs
