Skip to main content

ggplot_rs/position/
stack.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3
4use crate::data::{DataFrame, Value};
5
6use super::{Position, PositionParams};
7
8/// Stack bars/areas on top of each other (ggplot2's `position_stack(vjust
9/// = )`).
10///
11/// `vjust` places the row's `y` within its stacked segment: `1` (the
12/// default) puts it at the segment's top — the pre-B3-4 behaviour, and
13/// byte-identical (no `ymax` column is written). Anything else writes
14/// `y = (1 − vjust)·ymin + vjust·ymax` and an `ymax` column holding the
15/// segment's top, so e.g. `vjust(0.5)` centres a label in its segment.
16///
17/// `PositionStack` is also usable as a bare value (`vjust` 1, as before
18/// the struct gained fields) via the [`PositionStack`](const@PositionStack)
19/// const.
20#[derive(Clone, Debug, PartialEq)]
21pub struct PositionStack {
22    pub vjust: f64,
23}
24
25impl Default for PositionStack {
26    fn default() -> Self {
27        PositionStack { vjust: 1.0 }
28    }
29}
30
31/// The default stack (`vjust` 1): lets `.position(PositionStack)` keep
32/// compiling now that the type has fields.
33#[allow(non_upper_case_globals)]
34pub const PositionStack: PositionStack = PositionStack { vjust: 1.0 };
35
36impl PositionStack {
37    /// Place `y` within each stacked segment (`0` = bottom, `0.5` =
38    /// middle, `1` = top, the default).
39    pub fn with_vjust(mut self, vjust: f64) -> Self {
40        self.vjust = vjust;
41        self
42    }
43}
44
45/// `position_stack()` / `position_stack().with_vjust(v)` — see
46/// [`PositionStack`](struct@PositionStack).
47pub fn position_stack() -> PositionStack {
48    PositionStack::default()
49}
50
51impl Position for PositionStack {
52    fn compute(&self, data: &mut DataFrame, _params: &PositionParams) {
53        // Group by x, accumulate y values
54        let x_col = match data.column("x") {
55            Some(c) => c.to_vec(),
56            None => return,
57        };
58        let y_col = match data.column("y") {
59            Some(c) => c.to_vec(),
60            None => return,
61        };
62        super::preserve_raw_y(data, &y_col);
63
64        let n = y_col.len();
65        let mut tops = vec![0.0_f64; n];
66        let mut bottoms = vec![0.0_f64; n];
67
68        // ggplot2 stacks a group's positive and negative values
69        // independently within each x (#154, ggplot2's internal
70        // `pos`/`neg` split in `collide`/`pos_stack`): run the same top-
71        // down, legend-order accumulation twice, once over the rows
72        // with y >= 0 and once over y < 0, so e.g. a bar chart mixing a
73        // group's positive and negative segments at the same x gets two
74        // independent stacks (one above zero, one below) instead of a
75        // single pooled running total that lets a negative row cancel
76        // into a positive one's segment. ggplot2 stacks the first group
77        // at the top of its bucket (so top-to-bottom matches the
78        // legend), so each bucket still accumulates downward from its
79        // own total rather than upward from 0 — "top"/"bottom" here are
80        // in that accumulation order, not value order (a negative
81        // bucket's "top", nearer 0, has the smaller `tops[i]`), so
82        // `ymin`/`ymax`/`y` below always take the min/max of the pair
83        // rather than assume `tops[i] >= bottoms[i]`.
84        // Per-x accumulators keyed by the borrowed x key: O(1) per row
85        // instead of a linear scan over the distinct x values.
86        for positive in [true, false] {
87            let mut totals: HashMap<Cow<'_, str>, f64> = HashMap::new();
88            for (x, y) in x_col.iter().zip(y_col.iter()) {
89                let y_val = y.as_f64().unwrap_or(0.0);
90                if (y_val >= 0.0) != positive {
91                    continue;
92                }
93                *totals.entry(x.key_str()).or_insert(0.0) += y_val;
94            }
95
96            let mut consumed: HashMap<Cow<'_, str>, f64> = HashMap::new();
97            for (i, (x, y)) in x_col.iter().zip(y_col.iter()).enumerate() {
98                let y_val = y.as_f64().unwrap_or(0.0);
99                if (y_val >= 0.0) != positive {
100                    continue;
101                }
102                let x_key = x.key_str();
103                let total = totals.get(&x_key).copied().unwrap_or(0.0);
104                let run = consumed.get(&x_key).copied().unwrap_or(0.0);
105
106                // This group occupies [total - run - y, total - run]
107                // within its own (same-sign) bucket.
108                tops[i] = total - run;
109                bottoms[i] = total - run - y_val;
110
111                *consumed.entry(x_key).or_insert(0.0) += y_val;
112            }
113        }
114
115        // A non-finite vjust is validated (and warned on) separately; here
116        // it falls back to 1 (the byte-identical top-anchored path).
117        let vjust = if self.vjust.is_finite() {
118            self.vjust
119        } else {
120            1.0
121        };
122
123        let write_ymax = (vjust - 1.0).abs() > f64::EPSILON;
124        let lo_hi: Vec<(f64, f64)> = tops
125            .iter()
126            .zip(bottoms.iter())
127            .map(|(&t, &b)| (t.min(b), t.max(b)))
128            .collect();
129        let new_y: Vec<Value> = if write_ymax {
130            lo_hi
131                .iter()
132                .map(|&(lo, hi)| Value::Float((1.0 - vjust) * lo + vjust * hi))
133                .collect()
134        } else {
135            lo_hi.iter().map(|&(_, hi)| Value::Float(hi)).collect()
136        };
137        let ymin_vals: Vec<Value> = lo_hi.iter().map(|&(lo, _)| Value::Float(lo)).collect();
138
139        if let Some(col) = data.column_mut("y") {
140            *col = new_y;
141        }
142        if !data.has_column("ymin") {
143            data.add_column("ymin".to_string(), ymin_vals);
144        } else if let Some(col) = data.column_mut("ymin") {
145            *col = ymin_vals;
146        }
147        if write_ymax {
148            let ymax_vals: Vec<Value> = lo_hi.iter().map(|&(_, hi)| Value::Float(hi)).collect();
149            if !data.has_column("ymax") {
150                data.add_column("ymax".to_string(), ymax_vals);
151            } else if let Some(col) = data.column_mut("ymax") {
152                *col = ymax_vals;
153            }
154        }
155    }
156
157    fn name(&self) -> &str {
158        "stack"
159    }
160
161    fn per_panel(&self) -> bool {
162        true
163    }
164
165    fn validate(&self, warnings: &mut Vec<String>) {
166        if !self.vjust.is_finite() {
167            warnings.push(format!(
168                "position_stack: vjust({}) must be finite; falling back to 1",
169                self.vjust
170            ));
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    fn frame(x: &[f64], g: &[&str], y: &[f64]) -> DataFrame {
180        let mut df = DataFrame::new();
181        df.add_column("x".into(), x.iter().map(|v| Value::Float(*v)).collect());
182        df.add_column(
183            "fill".into(),
184            g.iter().map(|s| Value::Str(s.to_string())).collect(),
185        );
186        df.add_column("y".into(), y.iter().map(|v| Value::Float(*v)).collect());
187        df
188    }
189
190    fn col(df: &DataFrame, c: &str) -> Vec<f64> {
191        df.column(c)
192            .unwrap()
193            .iter()
194            .map(|v| v.as_f64().unwrap())
195            .collect()
196    }
197
198    /// R parity (ggplot2 4.0.3, `layer_data()`): x=1, fill groups a/b/c/d
199    /// with y = 3, -2, 4, -1 and `position_stack(vjust = 0.5)` gives
200    /// y/ymin/ymax = (5.5, 4, 7), (-2, -3, -1), (2, 0, 4), (-0.5, -1, 0)
201    /// (#154) — positives (a, c) stack above 0 using only their own sum
202    /// (7) as the total; negatives (b, d) stack below 0 using only their
203    /// own sum (-3), independently.
204    #[test]
205    fn negative_and_positive_values_stack_independently() {
206        let mut df = frame(&[1.0; 4], &["a", "b", "c", "d"], &[3.0, -2.0, 4.0, -1.0]);
207        PositionStack::default()
208            .with_vjust(0.5)
209            .compute(&mut df, &PositionParams::default());
210        let y = col(&df, "y");
211        let ymin = col(&df, "ymin");
212        let ymax = col(&df, "ymax");
213        for (i, (want_y, want_min, want_max)) in [
214            (5.5, 4.0, 7.0),
215            (-2.0, -3.0, -1.0),
216            (2.0, 0.0, 4.0),
217            (-0.5, -1.0, 0.0),
218        ]
219        .into_iter()
220        .enumerate()
221        {
222            assert!((y[i] - want_y).abs() < 1e-9, "y[{i}]={}", y[i]);
223            assert!((ymin[i] - want_min).abs() < 1e-9, "ymin[{i}]={}", ymin[i]);
224            assert!((ymax[i] - want_max).abs() < 1e-9, "ymax[{i}]={}", ymax[i]);
225        }
226    }
227
228    /// All-positive data (the common case) is byte-identical to the
229    /// pre-#154 pooled algorithm: the negative bucket is simply empty.
230    #[test]
231    fn all_positive_values_are_unaffected() {
232        let mut df = frame(&[1.0, 1.0], &["a", "b"], &[3.0, 5.0]);
233        PositionStack.compute(&mut df, &PositionParams::default());
234        assert_eq!(col(&df, "y"), vec![8.0, 5.0]);
235        assert_eq!(col(&df, "ymin"), vec![5.0, 0.0]);
236    }
237}