1use crate::aes::Aesthetic;
14use crate::data::{DataFrame, Value};
15use crate::scale::ScaleSet;
16
17use super::Stat;
18
19#[derive(Clone, Copy, Debug, PartialEq)]
21pub enum FunnelShape {
22 Proportional,
25 Silhouette { neck_width: f64, neck_height: f64 },
29 Pyramid,
32}
33
34#[derive(Clone, Debug)]
36pub struct StatFunnel {
37 pub shape: FunnelShape,
38 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 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(); 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 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 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 assert!((y[0] - 1.0).abs() < 1e-12);
243 assert!((y[1] - (1.0 - 1000.0 / 1520.0)).abs() < 1e-12);
244 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}