Skip to main content

ggplot_rs/position/
dodge.rs

1use crate::data::{DataFrame, Value};
2
3use super::{Position, PositionParams, DODGE_N_COL, DODGE_OFFSET_COL};
4
5/// Place groups side-by-side (ggplot2's `position_dodge(width)`).
6///
7/// Within each x position the groups present there (group / fill / colour
8/// level, in first-appearance order — the legend order — or reversed with
9/// [`with_reverse`](Self::with_reverse)) share `width` (in x units; one
10/// category on a discrete axis) and are centred on x:
11/// `offset = (k − (n − 1)/2) · width / n` for the k-th of n groups. On a
12/// continuous axis x itself moves; on a discrete axis the offset is stored in
13/// the `.x_dodge_offset` column and applied by the geoms that support it
14/// (point, pointrange, errorbar, linerange, ribbon, censor marks), so several
15/// models' intervals per term sit side by side — also under `coord_flip`.
16///
17/// `PositionDodge` is also usable as a value (the default, `width` from the
18/// layer, 0.9), as before 0.17.
19#[derive(Clone, Debug, Default, PartialEq)]
20pub struct PositionDodge {
21    /// Total dodge width; `None` = the layer default (0.9).
22    pub width: Option<f64>,
23    /// Reverse the group order (ggplot2 ≥ 3.5 `reverse = TRUE`).
24    pub reverse: bool,
25}
26
27/// The default dodge (`width` 0.9, groups in legend order): lets
28/// `.position(PositionDodge)` keep compiling now that the type has fields.
29#[allow(non_upper_case_globals)]
30pub const PositionDodge: PositionDodge = PositionDodge {
31    width: None,
32    reverse: false,
33};
34
35impl PositionDodge {
36    /// Dodge with an explicit total `width` (ggplot2's `position_dodge(width)`).
37    pub fn new(width: f64) -> Self {
38        PositionDodge {
39            width: Some(width),
40            reverse: false,
41        }
42    }
43
44    /// Reverse the order of the dodged groups.
45    pub fn with_reverse(mut self, reverse: bool) -> Self {
46        self.reverse = reverse;
47        self
48    }
49}
50
51/// `position_dodge(width)` — see [`PositionDodge`](struct@PositionDodge).
52pub fn position_dodge(width: f64) -> PositionDodge {
53    PositionDodge::new(width)
54}
55
56impl Position for PositionDodge {
57    fn compute(&self, data: &mut DataFrame, params: &PositionParams) {
58        dodge_rows(
59            data,
60            self.width.unwrap_or(params.width),
61            None,
62            self.reverse,
63            params.group_order.as_deref(),
64        );
65    }
66
67    fn name(&self) -> &str {
68        "dodge"
69    }
70
71    /// ggplot2 computes `position_dodge`'s group count/offset per facet
72    /// panel (`compute_panel`), not pooled across panels (#153) — verified
73    /// against ggplot2 4.0.3: a group missing from one panel's x level
74    /// doesn't shift the groups that panel does have.
75    fn per_panel(&self) -> bool {
76        true
77    }
78
79    /// Review #167 finding 2: the group *rank* still comes from the whole
80    /// layer (first-appearance order across every panel), even though the
81    /// per-x group *count* stays per-panel. Without this, a panel whose
82    /// rows happen to list groups in a different order than another panel
83    /// drew that group on the opposite side — a regression vs 0.22.2 and
84    /// ggplot2 4.0.3, which both rank by the global `group` id.
85    fn group_order(&self, data: &DataFrame) -> Option<Vec<String>> {
86        group_keys(data).map(|keys| dedup_order(&keys))
87    }
88}
89
90/// First-appearance-order, de-duplicated group keys — the "legend order"
91/// `group_order` (and, absent an external order, `dodge_rows` itself)
92/// ranks groups by.
93pub(crate) fn dedup_order(keys: &[String]) -> Vec<String> {
94    let set: indexmap::IndexSet<&str> = keys.iter().map(|k| k.as_str()).collect();
95    set.into_iter().map(str::to_string).collect()
96}
97
98/// The dodge key of each row: its group / fill / colour levels combined
99/// (`None` when the frame has none of these columns).
100pub(crate) fn group_keys(data: &DataFrame) -> Option<Vec<String>> {
101    let cols: Vec<&[Value]> = ["group", "fill", "color"]
102        .iter()
103        .filter_map(|c| data.column(c))
104        .collect();
105    if cols.is_empty() {
106        return None;
107    }
108    Some(
109        (0..data.nrows())
110            .map(|i| {
111                cols.iter()
112                    .map(|c| c[i].to_group_key())
113                    .collect::<Vec<_>>()
114                    .join("\u{1f}")
115            })
116            .collect(),
117    )
118}
119
120/// Shared dodge / dodge2 implementation. `padding` (dodge2) also rewrites
121/// `xmin`/`xmax` on a continuous axis to the shrunken element extent.
122///
123/// `external_order` (review #167 finding 2) is the whole layer's group
124/// rank order, computed by `build::compute_position_per_panel` from
125/// [`Position::group_order`] before the layer's rows were split into
126/// panels. When given, it — not this call's own (possibly per-panel)
127/// `data` — decides group rank; `None` (the non-per-panel call path, and
128/// every pre-#153 caller) falls back to the previous, locally-computed
129/// order.
130pub(crate) fn dodge_rows(
131    data: &mut DataFrame,
132    width: f64,
133    padding: Option<f64>,
134    reverse: bool,
135    external_order: Option<&[String]>,
136) {
137    let Some(x_col) = data.column("x").map(|c| c.to_vec()) else {
138        return;
139    };
140    let Some(keys) = group_keys(data) else {
141        return;
142    };
143    let width = if width.is_finite() && width > 0.0 {
144        width
145    } else {
146        0.9
147    };
148    // Group order (legend order), optionally reversed: the whole layer's
149    // order when given (#167 finding 2), else this call's own data —
150    // identical to the whole layer's order whenever there's no per-panel
151    // split (the non-per-panel path, and direct `compute()` calls in unit
152    // tests below).
153    let mut order: indexmap::IndexSet<&str> = match external_order {
154        Some(ext) => ext.iter().map(|s| s.as_str()).collect(),
155        None => keys.iter().map(|k| k.as_str()).collect(),
156    };
157    if order.len() <= 1 {
158        return;
159    }
160    if reverse {
161        order.reverse();
162    }
163    let rank = |k: &str| order.get_index_of(k).unwrap_or(0);
164
165    // Groups present at each x position, sorted by the global order.
166    let mut slots: indexmap::IndexMap<String, Vec<usize>> = indexmap::IndexMap::new();
167    for (i, x) in x_col.iter().enumerate() {
168        if x.is_na() {
169            continue;
170        }
171        let present = slots.entry(x.to_group_key()).or_default();
172        let r = rank(&keys[i]);
173        if !present.contains(&r) {
174            present.push(r);
175        }
176    }
177    for present in slots.values_mut() {
178        present.sort_unstable();
179    }
180
181    let discrete = x_col
182        .iter()
183        .any(|v| matches!(v, Value::Str(_) | Value::Bool(_)));
184    let mut offsets = vec![Value::Float(0.0); x_col.len()];
185    let mut counts = vec![Value::Float(1.0); x_col.len()];
186    let mut new_x = x_col.clone();
187    let mut new_xmin = data.column("xmin").map(|c| c.to_vec());
188    let mut new_xmax = data.column("xmax").map(|c| c.to_vec());
189    for (i, x) in x_col.iter().enumerate() {
190        let Some(present) = slots.get(&x.to_group_key()).filter(|_| !x.is_na()) else {
191            continue;
192        };
193        let n = present.len() as f64;
194        let k = present
195            .iter()
196            .position(|&r| r == rank(&keys[i]))
197            .unwrap_or(0) as f64;
198        let group_width = width / n;
199        let offset = (k - (n - 1.0) / 2.0) * group_width;
200        if discrete {
201            offsets[i] = Value::Float(offset);
202            counts[i] = Value::Float(n);
203            continue;
204        }
205        let Some(xv) = x.as_f64() else {
206            continue;
207        };
208        let center = xv + offset;
209        new_x[i] = Value::Float(center);
210        if let Some(pad) = padding {
211            let half = group_width * (1.0 - pad) / 2.0;
212            if let Some(c) = new_xmin.as_mut() {
213                c[i] = Value::Float(center - half);
214            }
215            if let Some(c) = new_xmax.as_mut() {
216                c[i] = Value::Float(center + half);
217            }
218        }
219    }
220
221    if discrete {
222        if let Some(col) = data.column_mut(DODGE_OFFSET_COL) {
223            *col = offsets;
224        } else {
225            data.add_column(DODGE_OFFSET_COL.to_string(), offsets);
226        }
227        if let Some(col) = data.column_mut(DODGE_N_COL) {
228            *col = counts;
229        } else {
230            data.add_column(DODGE_N_COL.to_string(), counts);
231        }
232        return;
233    }
234    if let Some(col) = data.column_mut("x") {
235        *col = new_x;
236    }
237    if padding.is_some() {
238        for (name, vals) in [("xmin", new_xmin), ("xmax", new_xmax)] {
239            if let (Some(vals), Some(col)) = (vals, data.column_mut(name)) {
240                *col = vals;
241            }
242        }
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    fn frame(x: Vec<Value>, g: &[&str]) -> DataFrame {
251        let mut df = DataFrame::new();
252        df.add_column("x".into(), x);
253        df.add_column(
254            "color".into(),
255            g.iter().map(|s| Value::Str(s.to_string())).collect(),
256        );
257        df
258    }
259
260    fn col(df: &DataFrame, c: &str) -> Vec<f64> {
261        df.column(c)
262            .unwrap()
263            .iter()
264            .map(|v| v.as_f64().unwrap())
265            .collect()
266    }
267
268    #[test]
269    fn continuous_dodge_matches_ggplot2() {
270        // ggplot2: position_dodge(width = 0.6) of 3 groups at x = 1 →
271        // 0.8, 1.0, 1.2.
272        let mut df = frame(vec![Value::Float(1.0); 3], &["a", "b", "c"]);
273        PositionDodge::new(0.6).compute(&mut df, &PositionParams::default());
274        let x = col(&df, "x");
275        for (got, want) in x.iter().zip([0.8, 1.0, 1.2]) {
276            assert!((got - want).abs() < 1e-12);
277        }
278        let mut df = frame(vec![Value::Float(1.0); 3], &["a", "b", "c"]);
279        PositionDodge::new(0.6)
280            .with_reverse(true)
281            .compute(&mut df, &PositionParams::default());
282        let x = col(&df, "x");
283        assert!((x[0] - 1.2).abs() < 1e-12 && (x[2] - 0.8).abs() < 1e-12);
284    }
285
286    #[test]
287    fn discrete_dodge_stores_offsets_and_centres_partial_slots() {
288        let x: Vec<Value> = ["t1", "t1", "t2"]
289            .iter()
290            .map(|s| Value::Str(s.to_string()))
291            .collect();
292        let mut df = frame(x, &["m1", "m2", "m2"]);
293        PositionDodge::new(0.5).compute(&mut df, &PositionParams::default());
294        // x itself untouched (a category), offsets in category units.
295        assert_eq!(df.column("x").unwrap()[0].as_str(), Some("t1"));
296        let off = col(&df, DODGE_OFFSET_COL);
297        assert!((off[0] + 0.125).abs() < 1e-12);
298        assert!((off[1] - 0.125).abs() < 1e-12);
299        // t2 has only m2 → centred (ggplot2's preserve = "total").
300        assert_eq!(off[2], 0.0);
301    }
302
303    #[test]
304    fn unit_value_still_works() {
305        let mut df = frame(vec![Value::Float(1.0); 2], &["a", "b"]);
306        PositionDodge.compute(&mut df, &PositionParams::default());
307        let x = col(&df, "x");
308        assert!((x[0] - 0.775).abs() < 1e-12 && (x[1] - 1.225).abs() < 1e-12);
309    }
310}