1use crate::aes::Aesthetic;
7use crate::coord::Coord;
8use crate::data::{DataFrame, Value};
9use crate::position::identity::PositionIdentity;
10use crate::position::Position;
11use crate::render::backend::{DrawBackend, LineStyle, Linetype, RectStyle};
12use crate::render::RenderError;
13use crate::scale::ScaleSet;
14use crate::stat::waterfall::{StatWaterfall, WATERFALL_DECREASE, WATERFALL_TOTAL};
15use crate::stat::Stat;
16use crate::theme::Theme;
17
18use super::{Geom, GeomParams};
19
20pub struct GeomWaterfall {
24 pub up: (u8, u8, u8),
26 pub down: (u8, u8, u8),
28 pub total: (u8, u8, u8),
30 pub color: (u8, u8, u8),
32 pub totals: Vec<String>,
34 pub final_total: Option<String>,
36 pub connectors: bool,
38 pub connector_color: (u8, u8, u8),
39 pub width: f64,
41 pub alpha: f64,
42}
43
44impl Default for GeomWaterfall {
45 fn default() -> Self {
46 GeomWaterfall {
47 up: (0x0c, 0xa6, 0x78),
48 down: (0xe0, 0x31, 0x31),
49 total: (0x4c, 0x6e, 0xa8),
50 color: (50, 50, 50),
51 totals: Vec::new(),
52 final_total: None,
53 connectors: true,
54 connector_color: (110, 110, 110),
55 width: 0.7,
56 alpha: 1.0,
57 }
58 }
59}
60
61impl GeomWaterfall {
62 pub fn with_totals<S: Into<String>>(mut self, levels: impl IntoIterator<Item = S>) -> Self {
64 self.totals = levels.into_iter().map(Into::into).collect();
65 self
66 }
67
68 pub fn with_final_total(mut self, label: impl Into<String>) -> Self {
70 self.final_total = Some(label.into());
71 self
72 }
73
74 pub(crate) fn has_totals(&self) -> bool {
76 !self.totals.is_empty() || self.final_total.is_some()
77 }
78
79 fn direction_color(&self, dir: &Value) -> (u8, u8, u8) {
80 match dir.key_str().as_ref() {
81 WATERFALL_TOTAL => self.total,
82 WATERFALL_DECREASE => self.down,
83 _ => self.up,
84 }
85 }
86}
87
88impl Geom for GeomWaterfall {
89 fn draw(
90 &self,
91 data: &DataFrame,
92 coord: &dyn Coord,
93 scales: &ScaleSet,
94 _theme: &Theme,
95 backend: &mut dyn DrawBackend,
96 ) -> Result<(), RenderError> {
97 let col = |n: &str| {
98 data.column(n)
99 .ok_or_else(|| RenderError::MissingAesthetic(n.into()))
100 };
101 let (x_col, y_col, ymin_col, ymax_col) = (col("x")?, col("y")?, col("ymin")?, col("ymax")?);
102 let delta_col = data.column("delta");
103 let dir_col = data.column("direction");
104 let fill_col = data.column("fill");
105 let area = backend.plot_area();
106 let x_scale = scales.get(&Aesthetic::X);
107 let y_scale = scales.get(&Aesthetic::Y);
108 let my = |v: f64| y_scale.map(|s| s.map(&Value::Float(v))).unwrap_or(0.0);
109
110 let half = if x_scale.map(|s| s.is_discrete()).unwrap_or(false) {
111 let n = x_scale.map(|s| s.breaks().len()).unwrap_or(1).max(1) as f64;
112 let pad = if coord.discrete_expand() { 1.1 } else { 1.0 };
113 self.width / (n * pad) / 2.0
114 } else {
115 super::continuous_bar_half_width(
116 x_col.iter().filter_map(|v| x_scale.map(|s| s.map(v))),
117 self.width,
118 0.02,
119 )
120 };
121
122 let mut prev: Option<(f64, f64)> = None;
124 for i in 0..data.nrows() {
125 let (Some(lo), Some(hi), Some(end)) = (
126 ymin_col[i].as_f64(),
127 ymax_col[i].as_f64(),
128 y_col[i].as_f64(),
129 ) else {
130 continue;
131 };
132 let nx = x_scale.map(|s| s.map(&x_col[i])).unwrap_or(0.5);
133 let (nlo, nhi, nend) = (my(lo), my(hi), my(end));
134
135 if self.connectors {
136 if let Some((px, pend)) = prev {
137 super::clear_mark(backend);
138 let a = coord.transform((px + half, pend), &area);
139 let b = coord.transform((nx - half, pend), &area);
140 backend.draw_line(
141 &[a, b],
142 &LineStyle {
143 color: self.connector_color,
144 alpha: self.alpha,
145 width: 0.8,
146 linetype: Linetype::Dashed,
147 ..Default::default()
148 },
149 )?;
150 }
151 }
152 prev = Some((nx, nend));
153
154 let dir = dir_col.map(|c| c[i].clone()).unwrap_or(Value::Na);
155 let fill = fill_col
156 .and_then(|c| scales.map_color(&Aesthetic::Fill, &c[i]))
157 .unwrap_or_else(|| self.direction_color(&dir));
158 let delta = delta_col.map(|c| &c[i]).unwrap_or(&y_col[i]);
159 let x_txt = super::tip_value(&x_col[i]);
160 let tip = if dir.key_str() == WATERFALL_TOTAL {
161 format!("{x_txt}: {}", super::tip_value(delta))
162 } else {
163 let d = delta.as_f64().unwrap_or(0.0);
164 let sign = if d >= 0.0 { "+" } else { "" };
165 format!(
166 "{x_txt}: {sign}{} (total {})",
167 super::tip_value(delta),
168 super::tip_value(&y_col[i])
169 )
170 };
171 super::set_mark(
172 backend,
173 Some(tip),
174 Some(x_txt),
175 super::raw_value(&dir).or_else(|| super::series_key(data, i)),
176 super::raw_value(delta),
177 );
178 super::set_mark_formatted(backend, scales, Some(&x_col[i]), &Aesthetic::Y, Some(delta));
179 super::set_row_extras(backend, data, i);
180
181 let (x0, y0) = coord.transform((nx - half, nhi), &area);
182 let (x1, y1) = coord.transform((nx + half, nlo), &area);
183 backend.draw_rect(
184 (x0.min(x1), y0.min(y1)),
185 (x0.max(x1), y0.max(y1)),
186 &RectStyle {
187 fill: Some(fill),
188 stroke: Some(self.color),
189 stroke_width: 0.5,
190 alpha: self.alpha,
191 clip: true,
192 corner_radius: 0.0,
193 },
194 )?;
195 }
196 super::clear_mark(backend);
197 super::clear_row_extras(backend);
198 Ok(())
199 }
200
201 fn required_aes(&self) -> Vec<Aesthetic> {
202 vec![Aesthetic::X, Aesthetic::Y]
203 }
204
205 fn default_stat(&self) -> Box<dyn Stat> {
206 Box::new(StatWaterfall {
207 totals: self.totals.clone(),
208 final_total: self.final_total.clone(),
209 })
210 }
211
212 fn default_position(&self) -> Box<dyn Position> {
213 Box::new(PositionIdentity)
214 }
215
216 fn default_params(&self) -> GeomParams {
217 GeomParams::default()
218 }
219
220 fn name(&self) -> &str {
221 "waterfall"
222 }
223
224 fn apply_element_geom(&mut self, el: &crate::theme::elements::ElementGeom) {
225 if let (true, Some(ink)) = (self.color == (50, 50, 50), el.ink) {
226 self.color = ink;
227 }
228 }
229
230 fn consumes_tooltips(&self) -> bool {
231 true
232 }
233
234 fn include_zero_baseline(&self) -> bool {
235 true
236 }
237}