ggplot_rs/position/
jitterdodge.rs1use crate::data::{DataFrame, Value};
2
3use super::{Position, PositionParams};
4
5pub 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 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 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 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 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 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 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}