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        // ggplot2 stacks the first group at the TOP (so the stack order top-to-
65        // bottom matches the legend), so accumulate downward from each x's total
66        // rather than upward from 0.
67        // Per-x accumulators keyed by the borrowed x key: O(1) per row instead
68        // of a linear scan over the distinct x values.
69        let mut totals: HashMap<Cow<'_, str>, f64> = HashMap::new();
70        for (x, y) in x_col.iter().zip(y_col.iter()) {
71            *totals.entry(x.key_str()).or_insert(0.0) += y.as_f64().unwrap_or(0.0);
72        }
73
74        let mut consumed: HashMap<Cow<'_, str>, f64> = HashMap::new();
75        let mut tops = Vec::with_capacity(y_col.len());
76        let mut bottoms = Vec::with_capacity(y_col.len());
77
78        for (x, y) in x_col.iter().zip(y_col.iter()) {
79            let x_key = x.key_str();
80            let y_val = y.as_f64().unwrap_or(0.0);
81            let total = totals.get(&x_key).copied().unwrap_or(0.0);
82            let run = consumed.get(&x_key).copied().unwrap_or(0.0);
83
84            // This group occupies [total - run - y, total - run] (top-down).
85            tops.push(total - run);
86            bottoms.push(total - run - y_val);
87
88            *consumed.entry(x_key).or_insert(0.0) += y_val;
89        }
90
91        // A non-finite vjust is validated (and warned on) separately; here
92        // it falls back to 1 (the byte-identical top-anchored path).
93        let vjust = if self.vjust.is_finite() {
94            self.vjust
95        } else {
96            1.0
97        };
98
99        let write_ymax = (vjust - 1.0).abs() > f64::EPSILON;
100        let new_y: Vec<Value> = if write_ymax {
101            bottoms
102                .iter()
103                .zip(tops.iter())
104                .map(|(&b, &t)| Value::Float((1.0 - vjust) * b + vjust * t))
105                .collect()
106        } else {
107            tops.iter().map(|&t| Value::Float(t)).collect()
108        };
109        let ymin_vals: Vec<Value> = bottoms.into_iter().map(Value::Float).collect();
110
111        if let Some(col) = data.column_mut("y") {
112            *col = new_y;
113        }
114        if !data.has_column("ymin") {
115            data.add_column("ymin".to_string(), ymin_vals);
116        } else if let Some(col) = data.column_mut("ymin") {
117            *col = ymin_vals;
118        }
119        if write_ymax {
120            let ymax_vals: Vec<Value> = tops.into_iter().map(Value::Float).collect();
121            if !data.has_column("ymax") {
122                data.add_column("ymax".to_string(), ymax_vals);
123            } else if let Some(col) = data.column_mut("ymax") {
124                *col = ymax_vals;
125            }
126        }
127    }
128
129    fn name(&self) -> &str {
130        "stack"
131    }
132
133    fn per_panel(&self) -> bool {
134        true
135    }
136
137    fn validate(&self, warnings: &mut Vec<String>) {
138        if !self.vjust.is_finite() {
139            warnings.push(format!(
140                "position_stack: vjust({}) must be finite; falling back to 1",
141                self.vjust
142            ));
143        }
144    }
145}