Skip to main content

ggplot_rs/position/
jitterdodge.rs

1use crate::data::{DataFrame, Value};
2
3use super::{Position, PositionParams};
4
5/// Combined jitter and dodge position adjustment (ggplot2's
6/// `position_jitterdodge()`): useful for showing individual points within
7/// dodged groups (e.g. over boxplots).
8///
9/// As in ggplot2, points are first dodged like [`PositionDodge`]
10/// (groups in legend order, counted per x slot), then jittered by
11/// `jitter_width / (ndodge + 2)` either side, where `ndodge` is the most
12/// groups any x slot holds, so each group's points stay inside its own
13/// slot (#429). [`Default`] uses ggplot2's `jitter.width = NULL`: 40 % of
14/// the x resolution (0.4 on a discrete axis).
15///
16/// [`PositionDodge`]: super::dodge::PositionDodge
17pub struct PositionJitterDodge {
18    pub jitter_width: f64,
19    pub jitter_height: f64,
20    pub dodge_width: f64,
21    seed: u64,
22    /// `jitter.width = NULL`: 0.4 × the x resolution, set by [`Default`].
23    auto_width: bool,
24}
25
26impl PositionJitterDodge {
27    /// Explicit `jitter.width` (x data units, before the `ndodge + 2`
28    /// division) and `jitter.height` (the y spread).
29    pub fn new(jitter_width: f64, jitter_height: f64) -> Self {
30        PositionJitterDodge {
31            jitter_width,
32            jitter_height,
33            dodge_width: 0.9,
34            seed: 42,
35            auto_width: false,
36        }
37    }
38
39    pub fn with_dodge_width(mut self, width: f64) -> Self {
40        self.dodge_width = width;
41        self
42    }
43
44    pub fn with_seed(mut self, seed: u64) -> Self {
45        self.seed = seed;
46        self
47    }
48
49    fn pseudo_random(seed: u64, i: usize) -> f64 {
50        // Simple deterministic hash → uniform in [-0.5, 0.5]
51        let h = seed
52            .wrapping_mul(6364136223846793005)
53            .wrapping_add(i as u64)
54            .wrapping_mul(1442695040888963407);
55        let h = h ^ (h >> 33);
56        let h = h.wrapping_mul(0xff51afd7ed558ccd);
57        (h as f64 / u64::MAX as f64) - 0.5
58    }
59}
60
61impl Default for PositionJitterDodge {
62    fn default() -> Self {
63        PositionJitterDodge {
64            auto_width: true,
65            ..PositionJitterDodge::new(0.4, 0.0)
66        }
67    }
68}
69
70/// ggplot2's `resolution(x, zero = FALSE)`: the smallest gap between
71/// distinct x values (1 for a discrete or single-valued x).
72fn resolution(x: &[Value]) -> f64 {
73    if super::is_discrete_x(x) {
74        return 1.0;
75    }
76    let mut v: Vec<f64> = x
77        .iter()
78        .filter_map(Value::as_f64)
79        .filter(|f| f.is_finite())
80        .collect();
81    v.sort_by(f64::total_cmp);
82    v.dedup();
83    v.windows(2)
84        .map(|w| w[1] - w[0])
85        .fold(None, |m: Option<f64>, d| Some(m.map_or(d, |m| m.min(d))))
86        .unwrap_or(1.0)
87}
88
89impl Position for PositionJitterDodge {
90    fn compute(&self, data: &mut DataFrame, params: &PositionParams) {
91        let x_col = match data.column("x") {
92            Some(c) => c.to_vec(),
93            None => return,
94        };
95        // ndodge: the most dodge groups sharing one x slot (1 without groups).
96        let keys = super::dodge::group_keys(data);
97        let mut slots: indexmap::IndexMap<String, indexmap::IndexSet<&str>> =
98            indexmap::IndexMap::new();
99        if let Some(keys) = &keys {
100            for (x, k) in x_col.iter().zip(keys) {
101                if !x.is_na() {
102                    slots.entry(x.to_group_key()).or_default().insert(k);
103                }
104            }
105        }
106        let ndodge = slots.values().map(|s| s.len()).max().unwrap_or(1).max(1) as f64;
107        let width = if self.auto_width {
108            0.4 * resolution(&x_col)
109        } else {
110            self.jitter_width
111        };
112        // ggplot2 jitters ±(width / (ndodge + 2)); `pseudo_random` is ±0.5.
113        let amount = 2.0 * width / (ndodge + 2.0);
114
115        super::dodge::dodge_rows(
116            data,
117            self.dodge_width,
118            None,
119            false,
120            false,
121            params.group_order.as_deref(),
122        );
123
124        let seed = self.seed;
125        if super::is_discrete_x(&x_col) {
126            // Categories can't move: the jitter joins the dodge offset.
127            super::add_x_offsets(data, |i| Self::pseudo_random(seed, i) * amount);
128        } else if let Some(col) = data.column_mut("x") {
129            for (i, v) in col.iter_mut().enumerate() {
130                if let Some(f) = v.as_f64() {
131                    *v = Value::Float(f + Self::pseudo_random(seed, i) * amount);
132                }
133            }
134        }
135
136        // Y jitter
137        if self.jitter_height.abs() > f64::EPSILON {
138            if let Some(col) = data.column_mut("y") {
139                for (i, v) in col.iter_mut().enumerate() {
140                    if let Some(f) = v.as_f64() {
141                        *v = Value::Float(
142                            f + Self::pseudo_random(self.seed.wrapping_add(1), i)
143                                * self.jitter_height,
144                        );
145                    }
146                }
147            }
148        }
149    }
150
151    /// Dodge groups ranked like `PositionDodge` (legend order, #428), so
152    /// the points line up with dodged boxes / violins.
153    fn group_order(&self, data: &DataFrame) -> Option<Vec<String>> {
154        super::dodge::group_keys(data).map(|keys| super::dodge::dedup_order(&keys))
155    }
156
157    fn name(&self) -> &str {
158        "jitterdodge"
159    }
160}