ggplot_rs/stat/
waffle.rs

1//! `StatWaffle`: unit layouts for waffle and parliament charts.
2//!
3//! Each input row contributes `weight` units (1 when unmapped; rounded to
4//! whole units) of its `fill` category; categories keep data order.
5//!
6//! * Waffle: units fill a grid `rows` high, column by column from the
7//!   bottom-left, so each category is a contiguous block.
8//! * Parliament: seats sit on concentric half-circle arcs (inner radius
9//!   `inner`, outer 1); each arc gets seats in proportion to its length and
10//!   categories are assigned left to right by angle.
11//!
12//! Output: one row per unit with its centre `x`/`y`, `fill` (and every other
13//! column of the row it came from), `unit` (half the unit spacing, in data
14//! units — the drawn radius) and constant `xmin`/`xmax`/`ymin`/`ymax`
15//! framing the layout.
16
17use std::f64::consts::PI;
18
19use crate::aes::Aesthetic;
20use crate::data::{DataFrame, Value};
21use crate::scale::ScaleSet;
22
23use super::Stat;
24
25/// Unit layout.
26#[derive(Clone, Copy, Debug, PartialEq)]
27pub enum UnitLayout {
28    /// Square grid `rows` units high.
29    Waffle { rows: usize },
30    /// Half-circle seating; `rows` arcs (0 picks a count from the total) with
31    /// inner radius `inner` (0–1).
32    Parliament { rows: usize, inner: f64 },
33}
34
35/// Unit-chart stat behind [`GeomWaffle`](crate::geom::waffle::GeomWaffle).
36#[derive(Clone, Debug)]
37pub struct StatWaffle {
38    pub layout: UnitLayout,
39}
40
41impl Default for StatWaffle {
42    fn default() -> Self {
43        StatWaffle {
44            layout: UnitLayout::Waffle { rows: 10 },
45        }
46    }
47}
48
49/// Seats per arc for a parliament of `n` seats on `rows` arcs: proportional
50/// to arc radius, rounded so they sum to `n`.
51pub(crate) fn seats_per_row(n: usize, rows: usize, inner: f64) -> Vec<usize> {
52    let radius = |i: usize| {
53        if rows == 1 {
54            1.0
55        } else {
56            inner + (1.0 - inner) * i as f64 / (rows - 1) as f64
57        }
58    };
59    let total_r: f64 = (0..rows).map(radius).sum();
60    let exact: Vec<f64> = (0..rows).map(|i| n as f64 * radius(i) / total_r).collect();
61    let mut seats: Vec<usize> = exact.iter().map(|e| e.floor() as usize).collect();
62    let mut left = n - seats.iter().sum::<usize>();
63    // Largest remainders first.
64    let mut order: Vec<usize> = (0..rows).collect();
65    order.sort_by(|&a, &b| (exact[b] - exact[b].floor()).total_cmp(&(exact[a] - exact[a].floor())));
66    for i in order {
67        if left == 0 {
68            break;
69        }
70        seats[i] += 1;
71        left -= 1;
72    }
73    seats
74}
75
76impl Stat for StatWaffle {
77    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
78        let weight = data.column("weight");
79        // Units in category order: (source row) repeated by its weight.
80        let fill = data.column("fill");
81        let mut cats: indexmap::IndexMap<String, Vec<usize>> = indexmap::IndexMap::new();
82        for i in 0..data.nrows() {
83            let w = match weight {
84                Some(c) => c[i].as_f64().filter(|w| w.is_finite() && *w > 0.0),
85                None => Some(1.0),
86            };
87            let Some(w) = w else { continue };
88            let key = fill
89                .map(|c| c[i].key_str().into_owned())
90                .unwrap_or_default();
91            let units = cats.entry(key).or_default();
92            units.extend(std::iter::repeat_n(i, w.round() as usize));
93        }
94        let units: Vec<usize> = cats.into_values().flatten().collect();
95        let n = units.len();
96        if n == 0 {
97            return DataFrame::new();
98        }
99
100        let (pos, unit, frame): (Vec<(f64, f64)>, f64, [f64; 4]) = match self.layout {
101            UnitLayout::Waffle { rows } => {
102                let rows = rows.max(1);
103                let cols = n.div_ceil(rows);
104                let pos = (0..n)
105                    .map(|k| ((k / rows) as f64 + 1.0, (k % rows) as f64 + 1.0))
106                    .collect();
107                (pos, 0.5, [0.5, cols as f64 + 0.5, 0.5, rows as f64 + 0.5])
108            }
109            UnitLayout::Parliament { rows, inner } => {
110                let inner = inner.clamp(0.05, 0.95);
111                let rows = if rows > 0 {
112                    rows
113                } else {
114                    ((n as f64 / 4.0).sqrt().round() as usize).clamp(1, 40)
115                };
116                let spacing = if rows > 1 {
117                    (1.0 - inner) / (rows - 1) as f64
118                } else {
119                    1.0 - inner
120                };
121                let mut seats: Vec<(f64, f64, f64)> = Vec::with_capacity(n); // (angle, x, y)
122                for (i, k) in seats_per_row(n, rows, inner).into_iter().enumerate() {
123                    let r = if rows == 1 {
124                        1.0
125                    } else {
126                        inner + spacing * i as f64
127                    };
128                    for j in 0..k {
129                        let a = if k == 1 {
130                            PI / 2.0
131                        } else {
132                            PI * (1.0 - j as f64 / (k - 1) as f64)
133                        };
134                        seats.push((a, r * a.cos(), r * a.sin()));
135                    }
136                }
137                // Left to right: decreasing angle, outer rows first on ties.
138                seats.sort_by(|a, b| {
139                    b.0.total_cmp(&a.0)
140                        .then((b.1.hypot(b.2)).total_cmp(&a.1.hypot(a.2)))
141                });
142                // A seat's radius: half the smaller of row spacing and the
143                // outer arc's seat pitch.
144                let unit = 0.45 * spacing.min(PI / (n as f64 / rows as f64).max(1.0));
145                let pos = seats.into_iter().map(|s| (s.1, s.2)).collect();
146                (pos, unit, [-1.0 - unit, 1.0 + unit, -unit, 1.0 + unit])
147            }
148        };
149
150        let mut out = DataFrame::new();
151        out.add_column("x".into(), pos.iter().map(|p| Value::Float(p.0)).collect());
152        out.add_column("y".into(), pos.iter().map(|p| Value::Float(p.1)).collect());
153        out.add_column("unit".into(), vec![Value::Float(unit); n]);
154        for (c, v) in ["xmin", "xmax", "ymin", "ymax"].iter().zip(frame) {
155            out.add_column(c.to_string(), vec![Value::Float(v); n]);
156        }
157        for name in data.column_names() {
158            if out.has_column(name) {
159                continue;
160            }
161            let col = data.column(name).unwrap_or_default();
162            out.add_column(
163                name.to_string(),
164                units.iter().map(|&i| col[i].clone()).collect(),
165            );
166        }
167        out
168    }
169
170    fn required_aes(&self) -> Vec<Aesthetic> {
171        Vec::new()
172    }
173
174    /// All categories share one grid / chamber.
175    fn panelwise(&self) -> bool {
176        true
177    }
178
179    fn name(&self) -> &str {
180        "waffle"
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    fn data() -> DataFrame {
189        let mut d = DataFrame::new();
190        d.add_column(
191            "fill".into(),
192            ["A", "B", "C"]
193                .iter()
194                .map(|s| Value::Str(s.to_string()))
195                .collect(),
196        );
197        d.add_column(
198            "weight".into(),
199            [12.0, 7.0, 4.0].iter().map(|&v| Value::Float(v)).collect(),
200        );
201        d
202    }
203
204    #[test]
205    fn waffle_fills_columns_bottom_up_in_category_order() {
206        let out = StatWaffle {
207            layout: UnitLayout::Waffle { rows: 5 },
208        }
209        .compute_group(&data(), &ScaleSet::new());
210        assert_eq!(out.nrows(), 23);
211        let x: Vec<f64> = out
212            .column("x")
213            .unwrap()
214            .iter()
215            .map(|v| v.as_f64().unwrap())
216            .collect();
217        let y: Vec<f64> = out
218            .column("y")
219            .unwrap()
220            .iter()
221            .map(|v| v.as_f64().unwrap())
222            .collect();
223        assert_eq!((x[0], y[0]), (1.0, 1.0));
224        assert_eq!((x[5], y[5]), (2.0, 1.0));
225        assert_eq!(out.column("fill").unwrap()[12], Value::Str("B".into()));
226        assert_eq!(
227            out.column("xmax").unwrap()[0],
228            Value::Float(5.5),
229            "ceil(23/5) cols"
230        );
231    }
232
233    #[test]
234    fn parliament_seats_sum_and_sit_on_the_half_disk() {
235        assert_eq!(seats_per_row(101, 4, 0.4).iter().sum::<usize>(), 101);
236        let out = StatWaffle {
237            layout: UnitLayout::Parliament {
238                rows: 0,
239                inner: 0.4,
240            },
241        }
242        .compute_group(&data(), &ScaleSet::new());
243        assert_eq!(out.nrows(), 23);
244        for i in 0..23 {
245            let (x, y) = (
246                out.column("x").unwrap()[i].as_f64().unwrap(),
247                out.column("y").unwrap()[i].as_f64().unwrap(),
248            );
249            let r = x.hypot(y);
250            assert!(y >= -1e-9 && (0.39..=1.0 + 1e-9).contains(&r), "{x} {y}");
251        }
252        // Category A takes the leftmost seats.
253        let x0 = out.column("x").unwrap()[0].as_f64().unwrap();
254        assert!(x0 < -0.3, "{x0}");
255    }
256}
ggplot_rs/stat/waffle.rs

ggplot_rs/stat/
waffle.rs
