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(data, self.width.unwrap_or(params.width), None, self.reverse);
59    }
60
61    fn name(&self) -> &str {
62        "dodge"
63    }
64}
65
66/// The dodge key of each row: its group / fill / colour levels combined
67/// (`None` when the frame has none of these columns).
68fn group_keys(data: &DataFrame) -> Option<Vec<String>> {
69    let cols: Vec<&[Value]> = ["group", "fill", "color"]
70        .iter()
71        .filter_map(|c| data.column(c))
72        .collect();
73    if cols.is_empty() {
74        return None;
75    }
76    Some(
77        (0..data.nrows())
78            .map(|i| {
79                cols.iter()
80                    .map(|c| c[i].to_group_key())
81                    .collect::<Vec<_>>()
82                    .join("\u{1f}")
83            })
84            .collect(),
85    )
86}
87
88/// Shared dodge / dodge2 implementation. `padding` (dodge2) also rewrites
89/// `xmin`/`xmax` on a continuous axis to the shrunken element extent.
90pub(crate) fn dodge_rows(data: &mut DataFrame, width: f64, padding: Option<f64>, reverse: bool) {
91    let Some(x_col) = data.column("x").map(|c| c.to_vec()) else {
92        return;
93    };
94    let Some(keys) = group_keys(data) else {
95        return;
96    };
97    let width = if width.is_finite() && width > 0.0 {
98        width
99    } else {
100        0.9
101    };
102    // Global group order (legend order), optionally reversed.
103    let mut order: indexmap::IndexSet<&str> = keys.iter().map(|k| k.as_str()).collect();
104    if order.len() <= 1 {
105        return;
106    }
107    if reverse {
108        order.reverse();
109    }
110    let rank = |k: &str| order.get_index_of(k).unwrap_or(0);
111
112    // Groups present at each x position, sorted by the global order.
113    let mut slots: indexmap::IndexMap<String, Vec<usize>> = indexmap::IndexMap::new();
114    for (i, x) in x_col.iter().enumerate() {
115        if x.is_na() {
116            continue;
117        }
118        let present = slots.entry(x.to_group_key()).or_default();
119        let r = rank(&keys[i]);
120        if !present.contains(&r) {
121            present.push(r);
122        }
123    }
124    for present in slots.values_mut() {
125        present.sort_unstable();
126    }
127
128    let discrete = x_col
129        .iter()
130        .any(|v| matches!(v, Value::Str(_) | Value::Bool(_)));
131    let mut offsets = vec![Value::Float(0.0); x_col.len()];
132    let mut counts = vec![Value::Float(1.0); x_col.len()];
133    let mut new_x = x_col.clone();
134    let mut new_xmin = data.column("xmin").map(|c| c.to_vec());
135    let mut new_xmax = data.column("xmax").map(|c| c.to_vec());
136    for (i, x) in x_col.iter().enumerate() {
137        let Some(present) = slots.get(&x.to_group_key()).filter(|_| !x.is_na()) else {
138            continue;
139        };
140        let n = present.len() as f64;
141        let k = present
142            .iter()
143            .position(|&r| r == rank(&keys[i]))
144            .unwrap_or(0) as f64;
145        let group_width = width / n;
146        let offset = (k - (n - 1.0) / 2.0) * group_width;
147        if discrete {
148            offsets[i] = Value::Float(offset);
149            counts[i] = Value::Float(n);
150            continue;
151        }
152        let Some(xv) = x.as_f64() else {
153            continue;
154        };
155        let center = xv + offset;
156        new_x[i] = Value::Float(center);
157        if let Some(pad) = padding {
158            let half = group_width * (1.0 - pad) / 2.0;
159            if let Some(c) = new_xmin.as_mut() {
160                c[i] = Value::Float(center - half);
161            }
162            if let Some(c) = new_xmax.as_mut() {
163                c[i] = Value::Float(center + half);
164            }
165        }
166    }
167
168    if discrete {
169        if let Some(col) = data.column_mut(DODGE_OFFSET_COL) {
170            *col = offsets;
171        } else {
172            data.add_column(DODGE_OFFSET_COL.to_string(), offsets);
173        }
174        if let Some(col) = data.column_mut(DODGE_N_COL) {
175            *col = counts;
176        } else {
177            data.add_column(DODGE_N_COL.to_string(), counts);
178        }
179        return;
180    }
181    if let Some(col) = data.column_mut("x") {
182        *col = new_x;
183    }
184    if padding.is_some() {
185        for (name, vals) in [("xmin", new_xmin), ("xmax", new_xmax)] {
186            if let (Some(vals), Some(col)) = (vals, data.column_mut(name)) {
187                *col = vals;
188            }
189        }
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    fn frame(x: Vec<Value>, g: &[&str]) -> DataFrame {
198        let mut df = DataFrame::new();
199        df.add_column("x".into(), x);
200        df.add_column(
201            "color".into(),
202            g.iter().map(|s| Value::Str(s.to_string())).collect(),
203        );
204        df
205    }
206
207    fn col(df: &DataFrame, c: &str) -> Vec<f64> {
208        df.column(c)
209            .unwrap()
210            .iter()
211            .map(|v| v.as_f64().unwrap())
212            .collect()
213    }
214
215    #[test]
216    fn continuous_dodge_matches_ggplot2() {
217        // ggplot2: position_dodge(width = 0.6) of 3 groups at x = 1 →
218        // 0.8, 1.0, 1.2.
219        let mut df = frame(vec![Value::Float(1.0); 3], &["a", "b", "c"]);
220        PositionDodge::new(0.6).compute(&mut df, &PositionParams::default());
221        let x = col(&df, "x");
222        for (got, want) in x.iter().zip([0.8, 1.0, 1.2]) {
223            assert!((got - want).abs() < 1e-12);
224        }
225        let mut df = frame(vec![Value::Float(1.0); 3], &["a", "b", "c"]);
226        PositionDodge::new(0.6)
227            .with_reverse(true)
228            .compute(&mut df, &PositionParams::default());
229        let x = col(&df, "x");
230        assert!((x[0] - 1.2).abs() < 1e-12 && (x[2] - 0.8).abs() < 1e-12);
231    }
232
233    #[test]
234    fn discrete_dodge_stores_offsets_and_centres_partial_slots() {
235        let x: Vec<Value> = ["t1", "t1", "t2"]
236            .iter()
237            .map(|s| Value::Str(s.to_string()))
238            .collect();
239        let mut df = frame(x, &["m1", "m2", "m2"]);
240        PositionDodge::new(0.5).compute(&mut df, &PositionParams::default());
241        // x itself untouched (a category), offsets in category units.
242        assert_eq!(df.column("x").unwrap()[0].as_str(), Some("t1"));
243        let off = col(&df, DODGE_OFFSET_COL);
244        assert!((off[0] + 0.125).abs() < 1e-12);
245        assert!((off[1] - 0.125).abs() < 1e-12);
246        // t2 has only m2 → centred (ggplot2's preserve = "total").
247        assert_eq!(off[2], 0.0);
248    }
249
250    #[test]
251    fn unit_value_still_works() {
252        let mut df = frame(vec![Value::Float(1.0); 2], &["a", "b"]);
253        PositionDodge.compute(&mut df, &PositionParams::default());
254        let x = col(&df, "x");
255        assert!((x[0] - 0.775).abs() < 1e-12 && (x[1] - 1.225).abs() < 1e-12);
256    }
257}
ggplot_rs/position/dodge.rs

ggplot_rs/position/
dodge.rs
