ggplot_rs/geom/
sankey.rs

1//! Sankey diagram: `geom_sankey`. Nodes are bars stacked in columns; links
2//! are ribbons whose width is the flow, drawn as smooth S-curves between the
3//! source's and target's bands. Layout by [`StatSankey`].
4
5use crate::aes::Aesthetic;
6use crate::coord::Coord;
7use crate::data::{DataFrame, Value};
8use crate::position::identity::PositionIdentity;
9use crate::position::Position;
10use crate::render::backend::{DrawBackend, FontFace, RectStyle, TextAnchor, TextStyle};
11use crate::render::RenderError;
12use crate::scale::ScaleSet;
13use crate::stat::sankey::StatSankey;
14use crate::stat::Stat;
15use crate::theme::Theme;
16
17use super::{Geom, GeomParams};
18
19/// Sankey diagram. Map `from`, `to` and (optionally) `weight`; nodes are
20/// coloured through the fill scale by node name, links by their source.
21pub struct GeomSankey {
22    /// Node width in column units (columns are 1 apart).
23    pub node_width: f64,
24    /// Gap between nodes as a fraction of the tallest column's flow.
25    pub node_padding: f64,
26    /// Move sink nodes to the last column.
27    pub justify: bool,
28    /// Fill used when no fill scale applies.
29    pub fill: (u8, u8, u8),
30    /// Node outline (`None` for none).
31    pub color: Option<(u8, u8, u8)>,
32    pub node_alpha: f64,
33    pub link_alpha: f64,
34    /// Draw node names beside the nodes.
35    pub labels: bool,
36    pub label_size: f64,
37    pub label_color: (u8, u8, u8),
38    /// Append the node's flow to its label (`"Gas (55)"`).
39    pub label_values: bool,
40    /// Label halo colour (`None` for no halo).
41    pub label_halo: Option<(u8, u8, u8)>,
42}
43
44impl Default for GeomSankey {
45    fn default() -> Self {
46        GeomSankey {
47            node_width: 0.12,
48            node_padding: 0.08,
49            justify: false,
50            fill: (97, 156, 255),
51            color: None,
52            node_alpha: 1.0,
53            link_alpha: 0.4,
54            labels: true,
55            label_size: 10.0,
56            label_color: (40, 40, 40),
57            label_values: false,
58            label_halo: Some((255, 255, 255)),
59        }
60    }
61}
62
63/// Samples per link edge: enough for a smooth S-curve at poster sizes.
64const LINK_STEPS: usize = 32;
65
66impl Geom for GeomSankey {
67    fn draw(
68        &self,
69        data: &DataFrame,
70        coord: &dyn Coord,
71        scales: &ScaleSet,
72        theme: &Theme,
73        backend: &mut dyn DrawBackend,
74    ) -> Result<(), RenderError> {
75        let col = |n: &str| {
76            data.column(n)
77                .ok_or_else(|| RenderError::MissingAesthetic(n.into()))
78        };
79        let (part, xmin, xmax, ymin, ymax) = (
80            col("part")?,
81            col("xmin")?,
82            col("xmax")?,
83            col("ymin")?,
84            col("ymax")?,
85        );
86        let (emin, emax) = (col("yend_min")?, col("yend_max")?);
87        let (node, from, to, value) = (col("node")?, col("from")?, col("to")?, col("value")?);
88        let column = data.column("column");
89        let fill_col = data.column("fill");
90        let area = backend.plot_area();
91        let mx = |v: f64| {
92            scales
93                .get(&Aesthetic::X)
94                .map(|s| s.map(&Value::Float(v)))
95                .unwrap_or(0.0)
96        };
97        let my = |v: f64| {
98            scales
99                .get(&Aesthetic::Y)
100                .map(|s| s.map(&Value::Float(v)))
101                .unwrap_or(0.0)
102        };
103        let f = |c: &[Value], i: usize| c[i].as_f64().unwrap_or(f64::NAN);
104        let fill_of = |i: usize| {
105            fill_col
106                .and_then(|c| scales.map_color(&Aesthetic::Fill, &c[i]))
107                .unwrap_or(self.fill)
108        };
109        let is = |i: usize, p: &str| part[i].key_str() == p;
110
111        // Links first, under the nodes.
112        for i in (0..data.nrows()).filter(|&i| is(i, "link")) {
113            let (x0, x1) = (f(xmin, i), f(xmax, i));
114            let (s0, s1, t0, t1) = (f(ymin, i), f(ymax, i), f(emin, i), f(emax, i));
115            if ![x0, x1, s0, s1, t0, t1].iter().all(|v| v.is_finite()) {
116                continue;
117            }
118            // Smoothstep in y = a cubic Bézier with both controls at mid-x.
119            let ease = |t: f64| t * t * (3.0 - 2.0 * t);
120            let mut pts = Vec::with_capacity(2 * (LINK_STEPS + 1));
121            for k in 0..=LINK_STEPS {
122                let t = k as f64 / LINK_STEPS as f64;
123                let y = s1 + (t1 - s1) * ease(t);
124                pts.push(coord.transform((mx(x0 + (x1 - x0) * t), my(y)), &area));
125            }
126            for k in (0..=LINK_STEPS).rev() {
127                let t = k as f64 / LINK_STEPS as f64;
128                let y = s0 + (t0 - s0) * ease(t);
129                pts.push(coord.transform((mx(x0 + (x1 - x0) * t), my(y)), &area));
130            }
131            let (a, b) = (super::tip_value(&from[i]), super::tip_value(&to[i]));
132            super::set_mark(
133                backend,
134                Some(format!("{a} → {b}: {}", super::tip_value(&value[i]))),
135                Some(a.clone()),
136                Some(a),
137                super::raw_value(&value[i]),
138            );
139            backend.draw_polygon(
140                &pts,
141                &RectStyle {
142                    fill: Some(fill_of(i)),
143                    stroke: None,
144                    stroke_width: 0.0,
145                    alpha: self.link_alpha,
146                    clip: true,
147                    corner_radius: 0.0,
148                },
149            )?;
150        }
151
152        let last_column = column
153            .map(|c| c.iter().filter_map(|v| v.as_f64()).fold(0.0, f64::max))
154            .unwrap_or(0.0);
155        let mut labels = Vec::new();
156        for i in (0..data.nrows()).filter(|&i| is(i, "node")) {
157            let (x0, x1, y0, y1) = (f(xmin, i), f(xmax, i), f(ymin, i), f(ymax, i));
158            let name = super::tip_value(&node[i]);
159            super::set_mark(
160                backend,
161                Some(format!("{name}: {}", super::tip_value(&value[i]))),
162                Some(name.clone()),
163                Some(name.clone()),
164                super::raw_value(&value[i]),
165            );
166            let (px0, py0) = coord.transform((mx(x0), my(y1)), &area);
167            let (px1, py1) = coord.transform((mx(x1), my(y0)), &area);
168            backend.draw_rect(
169                (px0.min(px1), py0.min(py1)),
170                (px0.max(px1), py0.max(py1)),
171                &RectStyle {
172                    fill: Some(fill_of(i)),
173                    stroke: self.color,
174                    stroke_width: if self.color.is_some() { 0.5 } else { 0.0 },
175                    alpha: self.node_alpha,
176                    clip: true,
177                    corner_radius: 0.0,
178                },
179            )?;
180            if self.labels {
181                // Last column: label to the left; elsewhere to the right.
182                let at_end =
183                    column.and_then(|c| c[i].as_f64()) == Some(last_column) && last_column > 0.0;
184                let ymid = my((y0 + y1) / 2.0);
185                let (pos, anchor) = if at_end {
186                    (coord.transform((mx(x0), ymid), &area), TextAnchor::End)
187                } else {
188                    (coord.transform((mx(x1), ymid), &area), TextAnchor::Start)
189                };
190                let dx = if at_end { -4.0 } else { 4.0 };
191                let text = if self.label_values {
192                    format!("{name} ({})", super::tip_value(&value[i]))
193                } else {
194                    name
195                };
196                labels.push((text, (pos.0 + dx, pos.1), anchor));
197            }
198        }
199        super::clear_mark(backend);
200        let family = super::resolve_font_family("", theme);
201        for (text, pos, anchor) in labels {
202            backend.draw_text(
203                &text,
204                pos,
205                &TextStyle {
206                    color: self.label_color,
207                    size: self.label_size,
208                    anchor,
209                    angle: 0.0,
210                    family: family.clone(),
211                    face: FontFace::Plain,
212                    bg_colour: self.label_halo,
213                    bg_r: 1.5,
214                },
215            )?;
216        }
217        Ok(())
218    }
219
220    fn required_aes(&self) -> Vec<Aesthetic> {
221        vec![Aesthetic::From, Aesthetic::To]
222    }
223
224    fn default_stat(&self) -> Box<dyn Stat> {
225        Box::new(StatSankey {
226            node_width: self.node_width,
227            node_padding: self.node_padding,
228            justify: self.justify,
229            ..Default::default()
230        })
231    }
232
233    fn default_position(&self) -> Box<dyn Position> {
234        Box::new(PositionIdentity)
235    }
236
237    fn default_params(&self) -> GeomParams {
238        GeomParams::default()
239    }
240
241    fn name(&self) -> &str {
242        "sankey"
243    }
244
245    fn apply_element_geom(&mut self, el: &crate::theme::elements::ElementGeom) {
246        if let Some(ink) = el.ink {
247            self.label_color = ink;
248        }
249        if let (Some(paper), Some(_)) = (el.paper, self.label_halo) {
250            self.label_halo = Some(paper);
251        }
252    }
253}
ggplot_rs/geom/sankey.rs

ggplot_rs/geom/
sankey.rs
