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.
6/// Useful for showing individual points within dodged groups (e.g., over boxplots).
7pub struct PositionJitterDodge {
8    pub jitter_width: f64,
9    pub jitter_height: f64,
10    pub dodge_width: f64,
11    seed: u64,
12}
13
14impl PositionJitterDodge {
15    pub fn new(jitter_width: f64, jitter_height: f64) -> Self {
16        PositionJitterDodge {
17            jitter_width,
18            jitter_height,
19            dodge_width: 0.9,
20            seed: 42,
21        }
22    }
23
24    pub fn with_dodge_width(mut self, width: f64) -> Self {
25        self.dodge_width = width;
26        self
27    }
28
29    pub fn with_seed(mut self, seed: u64) -> Self {
30        self.seed = seed;
31        self
32    }
33
34    fn pseudo_random(seed: u64, i: usize) -> f64 {
35        // Simple deterministic hash → uniform in [-0.5, 0.5]
36        let h = seed
37            .wrapping_mul(6364136223846793005)
38            .wrapping_add(i as u64)
39            .wrapping_mul(1442695040888963407);
40        let h = h ^ (h >> 33);
41        let h = h.wrapping_mul(0xff51afd7ed558ccd);
42        (h as f64 / u64::MAX as f64) - 0.5
43    }
44}
45
46impl Default for PositionJitterDodge {
47    fn default() -> Self {
48        PositionJitterDodge::new(0.4, 0.0)
49    }
50}
51
52impl Position for PositionJitterDodge {
53    fn compute(&self, data: &mut DataFrame, _params: &PositionParams) {
54        let x_col = match data.column("x") {
55            Some(c) => c.to_vec(),
56            None => return,
57        };
58
59        let group_col = data
60            .column("fill")
61            .or_else(|| data.column("color"))
62            .or_else(|| data.column("group"));
63
64        let group_keys: Vec<String> = match group_col {
65            Some(col) => col.iter().map(|v| v.to_group_key()).collect(),
66            None => {
67                // No groups — just jitter
68                if super::is_discrete_x(&x_col) {
69                    let (seed, w) = (self.seed, self.jitter_width);
70                    super::add_x_offsets(data, |i| Self::pseudo_random(seed, i) * w);
71                    return;
72                }
73                let mut new_x = x_col;
74                for (i, v) in new_x.iter_mut().enumerate() {
75                    if let Some(f) = v.as_f64() {
76                        *v =
77                            Value::Float(f + Self::pseudo_random(self.seed, i) * self.jitter_width);
78                    }
79                }
80                if let Some(col) = data.column_mut("x") {
81                    *col = new_x;
82                }
83                return;
84            }
85        };
86
87        let mut unique_groups: Vec<String> = Vec::new();
88        for g in &group_keys {
89            if !unique_groups.contains(g) {
90                unique_groups.push(g.clone());
91            }
92        }
93
94        let n_groups = unique_groups.len() as f64;
95        let group_width = if n_groups > 1.0 {
96            self.dodge_width / n_groups
97        } else {
98            0.0
99        };
100
101        if super::is_discrete_x(&x_col) {
102            // Categories can't move: dodge + jitter go to the offset column.
103            let offsets: Vec<f64> = group_keys
104                .iter()
105                .enumerate()
106                .map(|(i, group)| {
107                    let dodge = if n_groups > 1.0 {
108                        let k = unique_groups.iter().position(|g| g == group).unwrap() as f64;
109                        (k - (n_groups - 1.0) / 2.0) * group_width
110                    } else {
111                        0.0
112                    };
113                    dodge + Self::pseudo_random(self.seed, i) * self.jitter_width
114                })
115                .collect();
116            super::add_x_offsets(data, |i| offsets[i]);
117        }
118
119        let mut new_x = x_col.clone();
120        for (i, (x, group)) in x_col.iter().zip(group_keys.iter()).enumerate() {
121            if let Some(x_val) = x.as_f64() {
122                // Dodge offset
123                let dodge_offset = if n_groups > 1.0 {
124                    let group_idx = unique_groups.iter().position(|g| g == group).unwrap() as f64;
125                    (group_idx - (n_groups - 1.0) / 2.0) * group_width
126                } else {
127                    0.0
128                };
129
130                // Jitter offset
131                let jitter_x = Self::pseudo_random(self.seed, i) * self.jitter_width;
132
133                new_x[i] = Value::Float(x_val + dodge_offset + jitter_x);
134            }
135        }
136
137        if let Some(col) = data.column_mut("x") {
138            *col = new_x;
139        }
140
141        // Y jitter
142        if self.jitter_height.abs() > f64::EPSILON {
143            if let Some(col) = data.column_mut("y") {
144                for (i, v) in col.iter_mut().enumerate() {
145                    if let Some(f) = v.as_f64() {
146                        *v = Value::Float(
147                            f + Self::pseudo_random(self.seed.wrapping_add(1), i)
148                                * self.jitter_height,
149                        );
150                    }
151                }
152            }
153        }
154    }
155
156    fn name(&self) -> &str {
157        "jitterdodge"
158    }
159}