ggplot_rs/position/
dodge.rs1use crate::data::{DataFrame, Value};
2
3use super::{Position, PositionParams, DODGE_N_COL, DODGE_OFFSET_COL};
4
5#[derive(Clone, Debug, Default, PartialEq)]
20pub struct PositionDodge {
21 pub width: Option<f64>,
23 pub reverse: bool,
25}
26
27#[allow(non_upper_case_globals)]
30pub const PositionDodge: PositionDodge = PositionDodge {
31 width: None,
32 reverse: false,
33};
34
35impl PositionDodge {
36 pub fn new(width: f64) -> Self {
38 PositionDodge {
39 width: Some(width),
40 reverse: false,
41 }
42 }
43
44 pub fn with_reverse(mut self, reverse: bool) -> Self {
46 self.reverse = reverse;
47 self
48 }
49}
50
51pub 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 fn per_panel(&self) -> bool {
76 true
77 }
78
79 fn group_order(&self, data: &DataFrame) -> Option<Vec<String>> {
86 group_keys(data).map(|keys| dedup_order(&keys))
87 }
88}
89
90pub(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
98pub(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
120pub(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 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 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 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 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 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}