Skip to main content

ggplot_rs/stat/
funnel.rs

1//! `StatFunnel`: stage polygons for funnel and pyramid charts.
2//!
3//! Input: `x` (the stage; stages keep data order, repeated stages are summed)
4//! and `y` (the stage's value). Output: polygon vertices — several rows per
5//! stage sharing a `stage` key, in drawing order — as `fx`/`fy` in a layout
6//! where the chart spans `fx ∈ [-0.5, 0.5]` (centred) and `fy ∈ [0, 1]`
7//! (top = first stage). The layout lives outside `x`/`y` because the stage
8//! mapping makes the x scale discrete; `x`/`y` keep the stage and value.
9//! Per-row columns: `fx`, `fy`, `x`, `y`, `stage`, `value`, `pct_first`
10//! (value as a share of the first stage), `pct_total` (share of the sum) and
11//! the label anchor `label_y`.
12
13use crate::aes::Aesthetic;
14use crate::data::{DataFrame, Value};
15use crate::scale::ScaleSet;
16
17use super::Stat;
18
19/// How stage widths and heights encode the values.
20#[derive(Clone, Copy, Debug, PartialEq)]
21pub enum FunnelShape {
22    /// Equal-height stages whose width is proportional to the value; each
23    /// stage tapers to the next stage's width (a conversion funnel).
24    Proportional,
25    /// A fixed silhouette — full width at the top narrowing to `neck_width`
26    /// at `neck_height` above the bottom — cut into stages whose *height* is
27    /// proportional to the value (Highcharts' funnel).
28    Silhouette { neck_width: f64, neck_height: f64 },
29    /// A triangle with its apex on top, cut into stages whose height is
30    /// proportional to the value (Highcharts' pyramid).
31    Pyramid,
32}
33
34/// Funnel stat behind [`GeomFunnel`](crate::geom::funnel::GeomFunnel).
35#[derive(Clone, Debug)]
36pub struct StatFunnel {
37    pub shape: FunnelShape,
38    /// Vertical gap between stages as a fraction of the chart height.
39    pub gap: f64,
40}
41
42impl Default for StatFunnel {
43    fn default() -> Self {
44        StatFunnel {
45            shape: FunnelShape::Proportional,
46            gap: 0.01,
47        }
48    }
49}
50
51impl StatFunnel {
52    /// Half-width at height `y` (0 = bottom, 1 = top) of a silhouette shape.
53    fn half_width(&self, y: f64) -> f64 {
54        match self.shape {
55            FunnelShape::Silhouette {
56                neck_width,
57                neck_height,
58            } => {
59                let nw = neck_width.clamp(0.0, 1.0);
60                let nh = neck_height.clamp(0.0, 0.999);
61                if y <= nh {
62                    nw / 2.0
63                } else {
64                    (nw + (1.0 - nw) * (y - nh) / (1.0 - nh)) / 2.0
65                }
66            }
67            FunnelShape::Pyramid => (1.0 - y) / 2.0,
68            FunnelShape::Proportional => 0.5,
69        }
70    }
71}
72
73impl Stat for StatFunnel {
74    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
75        let (Some(xc), Some(yc)) = (data.column("x"), data.column("y")) else {
76            return DataFrame::new();
77        };
78        let mut stages: indexmap::IndexMap<String, (Value, f64)> = indexmap::IndexMap::new();
79        for (x, y) in xc.iter().zip(yc) {
80            let Some(v) = y.as_f64().filter(|v| v.is_finite() && *v >= 0.0) else {
81                continue;
82            };
83            stages
84                .entry(x.key_str().into_owned())
85                .or_insert_with(|| (x.clone(), 0.0))
86                .1 += v;
87        }
88        let stages: Vec<(Value, f64)> = stages.into_values().collect();
89        let n = stages.len();
90        let total: f64 = stages.iter().map(|s| s.1).sum();
91        let max = stages.iter().map(|s| s.1).fold(0.0, f64::max);
92        if n == 0 || total <= 0.0 {
93            return DataFrame::new();
94        }
95        let first = stages[0].1;
96        let gap = self.gap.clamp(0.0, 0.5 / n as f64);
97
98        let mut cols: [Vec<f64>; 3] = Default::default(); // fx, fy, label_y
99        let (mut stage, mut value, mut pct_first, mut pct_total) = (vec![], vec![], vec![], vec![]);
100        let mut top = 1.0;
101        for (i, (key, v)) in stages.iter().enumerate() {
102            let h = match self.shape {
103                FunnelShape::Proportional => 1.0 / n as f64,
104                _ => v / total,
105            };
106            let (y1, y0) = (top, top - h);
107            top = y0;
108            let (y1g, y0g) = (y1 - gap / 2.0, (y0 + gap / 2.0).min(y1 - gap / 2.0));
109            let poly: Vec<(f64, f64)> = match self.shape {
110                FunnelShape::Proportional => {
111                    let w_top = if max > 0.0 { v / max / 2.0 } else { 0.0 };
112                    let w_bot = stages
113                        .get(i + 1)
114                        .map(|s| if max > 0.0 { s.1 / max / 2.0 } else { 0.0 })
115                        .unwrap_or(w_top);
116                    vec![(-w_top, y1g), (w_top, y1g), (w_bot, y0g), (-w_bot, y0g)]
117                }
118                _ => {
119                    // Right edge top → bottom (through the neck corner when
120                    // the stage spans it), then the left edge back up.
121                    let mut ys = vec![y1g];
122                    if let FunnelShape::Silhouette { neck_height, .. } = self.shape {
123                        if neck_height < y1g && neck_height > y0g {
124                            ys.push(neck_height);
125                        }
126                    }
127                    ys.push(y0g);
128                    let mut p: Vec<(f64, f64)> =
129                        ys.iter().map(|&y| (self.half_width(y), y)).collect();
130                    p.extend(ys.iter().rev().map(|&y| (-self.half_width(y), y)));
131                    p
132                }
133            };
134            for &(px, py) in &poly {
135                cols[0].push(px);
136                cols[1].push(py);
137                cols[2].push((y1g + y0g) / 2.0);
138                stage.push(key.clone());
139                value.push(Value::Float(*v));
140                pct_first.push(if first > 0.0 {
141                    Value::Float(v / first)
142                } else {
143                    Value::Na
144                });
145                pct_total.push(Value::Float(v / total));
146            }
147        }
148        let f = |v: &Vec<f64>| v.iter().map(|&x| Value::Float(x)).collect::<Vec<_>>();
149        let mut out = DataFrame::new();
150        out.add_column("fx".into(), f(&cols[0]));
151        out.add_column("fy".into(), f(&cols[1]));
152        out.add_column("x".into(), stage.clone());
153        out.add_column("y".into(), value.clone());
154        out.add_column("stage".into(), stage);
155        out.add_column("value".into(), value);
156        out.add_column("pct_first".into(), pct_first);
157        out.add_column("pct_total".into(), pct_total);
158        out.add_column("label_y".into(), f(&cols[2]));
159        out
160    }
161
162    fn required_aes(&self) -> Vec<Aesthetic> {
163        vec![Aesthetic::X, Aesthetic::Y]
164    }
165
166    /// Stages share one silhouette: every row of the panel at once.
167    fn panelwise(&self) -> bool {
168        true
169    }
170
171    fn name(&self) -> &str {
172        "funnel"
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    fn data() -> DataFrame {
181        let mut d = DataFrame::new();
182        d.add_column(
183            "x".into(),
184            ["Visit", "Cart", "Buy", "Cart"]
185                .iter()
186                .map(|s| Value::Str(s.to_string()))
187                .collect(),
188        );
189        d.add_column(
190            "y".into(),
191            [1000.0, 300.0, 120.0, 100.0]
192                .iter()
193                .map(|&v| Value::Float(v))
194                .collect(),
195        );
196        d
197    }
198
199    fn col(d: &DataFrame, c: &str) -> Vec<f64> {
200        d.column(c)
201            .unwrap()
202            .iter()
203            .map(|v| v.as_f64().unwrap())
204            .collect()
205    }
206
207    #[test]
208    fn proportional_widths_and_shares() {
209        let out = StatFunnel {
210            gap: 0.0,
211            ..Default::default()
212        }
213        .compute_group(&data(), &ScaleSet::new());
214        assert_eq!(out.nrows(), 3 * 4);
215        let x = col(&out, "fx");
216        assert_eq!(
217            &x[..4],
218            &[-0.5, 0.5, 0.2, -0.2],
219            "Visit tapers to Cart (400)"
220        );
221        assert_eq!(
222            &x[8..12],
223            &[-0.06, 0.06, 0.06, -0.06],
224            "last stage is a bar"
225        );
226        let pf = col(&out, "pct_first");
227        assert_eq!((pf[0], pf[4], pf[8]), (1.0, 0.4, 0.12));
228    }
229
230    #[test]
231    fn silhouette_heights_are_proportional_and_pass_the_neck() {
232        let stat = StatFunnel {
233            shape: FunnelShape::Silhouette {
234                neck_width: 0.3,
235                neck_height: 0.25,
236            },
237            gap: 0.0,
238        };
239        let out = stat.compute_group(&data(), &ScaleSet::new());
240        let y = col(&out, "fy");
241        // Visit spans 1000/1520 of the height from the top.
242        assert!((y[0] - 1.0).abs() < 1e-12);
243        assert!((y[1] - (1.0 - 1000.0 / 1520.0)).abs() < 1e-12);
244        // Buy (bottom 120/1520) sits inside the neck: constant width.
245        let n = out.nrows();
246        let x = col(&out, "fx");
247        assert!((x[n - 1] + 0.15).abs() < 1e-12);
248        assert!((stat.half_width(1.0) - 0.5).abs() < 1e-12);
249    }
250}