ggplot_rs/stat/
tree.rs

1//! `StatTree`: hierarchical node-link layouts (tidy tree / dendrogram /
2//! org chart).
3//!
4//! Input: parent → child links (`from` = parent, `to` = child). Nodes with
5//! no parent are roots (several roots lay out side by side). Leaves are
6//! spaced one unit apart in depth-first order and every parent is centred
7//! over its children, so subtrees never overlap. A node's level is its depth
8//! (root 0) unless `height` names a column holding each link's *parent*
9//! height (dendrograms: leaves at 0, merges at their merge height).
10//!
11//! Output: one row per node (`part = "node"`, `x`/`y`, `node`, `depth`,
12//! `leaf`) then one per link (`part = "link"`, `x`/`y` = parent, `xend`/
13//! `yend` = child, `from`, `to`), already oriented (top-down, left-right or
14//! radial), plus `xmin`/`xmax`/`ymin`/`ymax` framing the layout.
15
16use std::collections::HashMap;
17use std::f64::consts::TAU;
18
19use crate::aes::Aesthetic;
20use crate::data::{DataFrame, Value};
21use crate::scale::ScaleSet;
22
23use super::Stat;
24
25/// Tree orientation.
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
27pub enum TreeOrientation {
28    /// Root on top, children below.
29    #[default]
30    TopDown,
31    /// Root on the left, children to the right.
32    LeftRight,
33    /// Root in the centre, levels as rings.
34    Radial,
35}
36
37/// Tree layout stat behind [`GeomTree`](crate::geom::tree::GeomTree).
38#[derive(Clone, Debug, Default)]
39pub struct StatTree {
40    pub orientation: TreeOrientation,
41    /// Put every leaf on the deepest level (dendrogram style).
42    pub align_leaves: bool,
43    /// Column holding the parent's height on each link (dendrograms).
44    pub height: Option<String>,
45}
46
47impl Stat for StatTree {
48    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
49        let (Some(from), Some(to)) = (data.column("from"), data.column("to")) else {
50            return DataFrame::new();
51        };
52        let hcol = self.height.as_deref().and_then(|h| data.column(h));
53        let mut index: HashMap<String, usize> = HashMap::new();
54        let mut names: Vec<String> = Vec::new();
55        let mut id = |v: &Value| -> usize {
56            let key = v.key_str().into_owned();
57            *index.entry(key.clone()).or_insert_with(|| {
58                names.push(key);
59                names.len() - 1
60            })
61        };
62        let mut links: Vec<(usize, usize)> = Vec::new();
63        let mut height: HashMap<usize, f64> = HashMap::new();
64        for i in 0..data.nrows() {
65            if from[i].is_na() || to[i].is_na() {
66                continue;
67            }
68            let (p, c) = (id(&from[i]), id(&to[i]));
69            if p != c {
70                links.push((p, c));
71                if let Some(h) = hcol.and_then(|col| col[i].as_f64()) {
72                    height.insert(p, h);
73                }
74            }
75        }
76        let n = names.len();
77        if n == 0 {
78            return DataFrame::new();
79        }
80        let mut children: Vec<Vec<usize>> = vec![Vec::new(); n];
81        let mut has_parent = vec![false; n];
82        for &(p, c) in &links {
83            if !has_parent[c] {
84                children[p].push(c);
85                has_parent[c] = true;
86            }
87        }
88        let roots: Vec<usize> = (0..n).filter(|&i| !has_parent[i]).collect();
89        // A cycle has no root: start from the first node.
90        let roots = if roots.is_empty() { vec![0] } else { roots };
91
92        let mut x = vec![f64::NAN; n];
93        let mut depth = vec![0usize; n];
94        let mut next_leaf = 0.0;
95        let mut seen = vec![false; n];
96        fn place(
97            v: usize,
98            d: usize,
99            children: &[Vec<usize>],
100            x: &mut [f64],
101            depth: &mut [usize],
102            next_leaf: &mut f64,
103            seen: &mut [bool],
104        ) {
105            seen[v] = true;
106            depth[v] = d;
107            let kids: Vec<usize> = children[v].iter().copied().filter(|&c| !seen[c]).collect();
108            if kids.is_empty() {
109                x[v] = *next_leaf;
110                *next_leaf += 1.0;
111                return;
112            }
113            for &c in &kids {
114                place(c, d + 1, children, x, depth, next_leaf, seen);
115            }
116            x[v] = (x[kids[0]] + x[kids[kids.len() - 1]]) / 2.0;
117        }
118        for &r in &roots {
119            if !seen[r] {
120                place(
121                    r,
122                    0,
123                    &children,
124                    &mut x,
125                    &mut depth,
126                    &mut next_leaf,
127                    &mut seen,
128                );
129                next_leaf += 0.5;
130            }
131        }
132        let leaf: Vec<bool> = (0..n).map(|i| children[i].is_empty()).collect();
133        let max_depth = depth.iter().copied().max().unwrap_or(0) as f64;
134        let level = |i: usize| -> f64 {
135            if hcol.is_some() {
136                return height.get(&i).copied().unwrap_or(0.0);
137            }
138            if self.align_leaves && leaf[i] {
139                max_depth
140            } else {
141                depth[i] as f64
142            }
143        };
144        // Down-axis value: larger = further from the root.
145        let max_h = (0..n).map(level).fold(0.0, f64::max);
146        let down = |i: usize| {
147            if hcol.is_some() {
148                max_h - level(i)
149            } else {
150                level(i)
151            }
152        };
153        let span = (next_leaf - 1.5).max(1.0);
154        let reach = (0..n).map(down).fold(0.0, f64::max).max(1.0);
155        let place_xy = |i: usize| -> (f64, f64) {
156            let (along, d) = (x[i], down(i));
157            match self.orientation {
158                TreeOrientation::TopDown => (along, -d),
159                TreeOrientation::LeftRight => (d, -along),
160                TreeOrientation::Radial => {
161                    // Leaves around the full circle; depth as radius.
162                    let a = TAU * along / (span + 1.0);
163                    let r = d / reach;
164                    (r * a.sin(), r * a.cos())
165                }
166            }
167        };
168        let pos: Vec<(f64, f64)> = (0..n).map(place_xy).collect();
169        let (mut x0, mut x1, mut y0, mut y1) = (f64::MAX, f64::MIN, f64::MAX, f64::MIN);
170        for &(px, py) in &pos {
171            x0 = x0.min(px);
172            x1 = x1.max(px);
173            y0 = y0.min(py);
174            y1 = y1.max(py);
175        }
176        let pad = |lo: f64, hi: f64| 0.08 * (hi - lo).max(1e-9) + 0.3;
177        let (px, py) = (pad(x0, x1), pad(y0, y1));
178
179        let s = |v: &str| Value::Str(v.to_string());
180        let tree_links: Vec<(usize, usize)> = (0..n)
181            .flat_map(|p| children[p].iter().map(move |&c| (p, c)))
182            .collect();
183        let rows = n + tree_links.len();
184        let (mut part, mut xs, mut ys, mut xe, mut ye) = (vec![], vec![], vec![], vec![], vec![]);
185        let (mut node, mut f, mut t, mut dep, mut lf) = (vec![], vec![], vec![], vec![], vec![]);
186        for i in 0..n {
187            part.push(s("node"));
188            xs.push(pos[i].0);
189            ys.push(pos[i].1);
190            xe.push(pos[i].0);
191            ye.push(pos[i].1);
192            node.push(s(&names[i]));
193            f.push(s(&names[i]));
194            t.push(s(&names[i]));
195            dep.push(Value::Float(depth[i] as f64));
196            lf.push(Value::Bool(leaf[i]));
197        }
198        for &(p, c) in &tree_links {
199            part.push(s("link"));
200            xs.push(pos[p].0);
201            ys.push(pos[p].1);
202            xe.push(pos[c].0);
203            ye.push(pos[c].1);
204            node.push(s(&names[c]));
205            f.push(s(&names[p]));
206            t.push(s(&names[c]));
207            dep.push(Value::Float(depth[c] as f64));
208            lf.push(Value::Bool(leaf[c]));
209        }
210        let fl = |v: Vec<f64>| v.into_iter().map(Value::Float).collect::<Vec<_>>();
211        let mut out = DataFrame::new();
212        out.add_column("part".into(), part);
213        out.add_column("x".into(), fl(xs));
214        out.add_column("y".into(), fl(ys));
215        out.add_column("xend".into(), fl(xe));
216        out.add_column("yend".into(), fl(ye));
217        out.add_column("node".into(), node);
218        out.add_column("from".into(), f);
219        out.add_column("to".into(), t);
220        out.add_column("depth".into(), dep);
221        out.add_column("leaf".into(), lf);
222        for (c, v) in [
223            ("xmin", x0 - px),
224            ("xmax", x1 + px),
225            ("ymin", y0 - py),
226            ("ymax", y1 + py),
227        ] {
228            out.add_column(c.into(), vec![Value::Float(v); rows]);
229        }
230        out
231    }
232
233    fn required_aes(&self) -> Vec<Aesthetic> {
234        vec![Aesthetic::From, Aesthetic::To]
235    }
236
237    /// The layout needs the whole hierarchy at once.
238    fn panelwise(&self) -> bool {
239        true
240    }
241
242    fn name(&self) -> &str {
243        "tree"
244    }
245}
246
247/// Parent → child links for a dendrogram from an R `hclust`-style merge
248/// matrix: row `k` merges `a` and `b` (negative = leaf `-a`, 1-based;
249/// positive = the cluster formed at row `a`) at `heights[k]`. Returns
250/// columns `from`, `to` and `height` (the parent's merge height) ready for
251/// `geom_dendrogram`; leaves are named by `labels` (1-based leaf `i` →
252/// `labels[i - 1]`), clusters `"#k"`.
253pub fn dendrogram_links(
254    merge: &[(i64, i64)],
255    heights: &[f64],
256    labels: &[&str],
257) -> Vec<(String, Vec<Value>)> {
258    let name = |v: i64| -> String {
259        if v < 0 {
260            labels
261                .get((-v - 1) as usize)
262                .map(|s| s.to_string())
263                .unwrap_or_else(|| format!("{}", -v))
264        } else {
265            format!("#{v}")
266        }
267    };
268    let (mut from, mut to, mut h) = (vec![], vec![], vec![]);
269    for (k, &(a, b)) in merge.iter().enumerate() {
270        let parent = format!("#{}", k + 1);
271        let hk = heights.get(k).copied().unwrap_or(0.0);
272        for child in [a, b] {
273            from.push(Value::Str(parent.clone()));
274            to.push(Value::Str(name(child)));
275            h.push(Value::Float(hk));
276        }
277    }
278    vec![
279        ("from".into(), from),
280        ("to".into(), to),
281        ("height".into(), h),
282    ]
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    fn edges(e: &[(&str, &str)]) -> DataFrame {
290        let mut d = DataFrame::new();
291        d.add_column(
292            "from".into(),
293            e.iter().map(|r| Value::Str(r.0.into())).collect(),
294        );
295        d.add_column(
296            "to".into(),
297            e.iter().map(|r| Value::Str(r.1.into())).collect(),
298        );
299        d
300    }
301
302    fn col(d: &DataFrame, c: &str) -> Vec<f64> {
303        d.column(c)
304            .unwrap()
305            .iter()
306            .map(|v| v.as_f64().unwrap())
307            .collect()
308    }
309
310    #[test]
311    fn parents_centre_over_children_and_leaves_never_collide() {
312        let d = edges(&[
313            ("CEO", "CTO"),
314            ("CEO", "CFO"),
315            ("CTO", "Dev"),
316            ("CTO", "Ops"),
317            ("CFO", "Acc"),
318        ]);
319        let out = StatTree::default().compute_group(&d, &ScaleSet::new());
320        assert_eq!(out.nrows(), 6 + 5);
321        let (x, y) = (col(&out, "x"), col(&out, "y"));
322        // Node order: CEO, CTO, CFO, Dev, Ops, Acc.
323        assert_eq!(&x[3..6], &[0.0, 1.0, 2.0], "leaves one apart");
324        assert_eq!(x[1], 0.5, "CTO centred over Dev/Ops");
325        assert_eq!(x[0], (0.5 + 2.0) / 2.0);
326        assert_eq!(&y[..3], &[0.0, -1.0, -1.0]);
327    }
328
329    #[test]
330    fn dendrogram_heights_and_radial() {
331        let links = dendrogram_links(&[(-1, -2), (-3, 1)], &[1.0, 2.5], &["a", "b", "c"]);
332        let mut d = DataFrame::new();
333        for (k, v) in links {
334            d.add_column(k, v);
335        }
336        let out = StatTree {
337            height: Some("height".into()),
338            ..Default::default()
339        }
340        .compute_group(&d, &ScaleSet::new());
341        let names: Vec<String> = out.column("node").unwrap()[..5]
342            .iter()
343            .map(|v| v.key_str().into_owned())
344            .collect();
345        let y = col(&out, "y");
346        let at = |n: &str| y[names.iter().position(|m| m == n).unwrap()];
347        // Root (#2, h 2.5) on top at 0; #1 (h 1) at -1.5; leaves at -2.5.
348        assert_eq!(
349            (at("#2"), at("#1"), at("a"), at("c")),
350            (0.0, -1.5, -2.5, -2.5)
351        );
352        let radial = StatTree {
353            orientation: TreeOrientation::Radial,
354            ..Default::default()
355        }
356        .compute_group(&d, &ScaleSet::new());
357        let (rx, ry) = (col(&radial, "x"), col(&radial, "y"));
358        let root = names.iter().position(|m| m == "#2").unwrap();
359        assert!(rx[root].hypot(ry[root]) < 1e-12, "root at the centre");
360    }
361}
ggplot_rs/stat/tree.rs

ggplot_rs/stat/
tree.rs
