Skip to main content

ggplot_rs/stat/
arcdiagram.rs

1//! `StatArcdiagram`: node positions and link spans for an arc diagram.
2//!
3//! Input: an edge list (`from`, `to`, optional `weight`, 1 when unmapped;
4//! repeated pairs summed, self-loops dropped). Nodes sit on a line at
5//! x = 0, 1, 2, … in first-appearance order (or by descending total flow);
6//! each link spans from its source's x to its target's x.
7//!
8//! Output: one row per node (`part = "node"`, `x`, `y = 0`, `value` = total
9//! in + out flow) then one per link (`part = "link"`, `x` → `xend`,
10//! `value` = weight), plus `node` (name; the source for a link), `from`,
11//! `to`, and `ymin`/`ymax` sized so the y scale holds every arc (above, or
12//! below for leftward links when `directed_below`) and room for labels.
13
14use std::collections::HashMap;
15
16use crate::aes::Aesthetic;
17use crate::data::{DataFrame, Value};
18use crate::scale::ScaleSet;
19
20use super::Stat;
21
22/// Arc-diagram layout stat behind [`GeomArcdiagram`](crate::geom::arcdiagram::GeomArcdiagram).
23#[derive(Clone, Debug, Default)]
24pub struct StatArcdiagram {
25    /// Order nodes by descending total flow instead of first appearance.
26    pub sort_by_value: bool,
27    /// Draw links that run right-to-left below the line (shows direction).
28    pub directed_below: bool,
29}
30
31impl Stat for StatArcdiagram {
32    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
33        let (Some(from), Some(to)) = (data.column("from"), data.column("to")) else {
34            return DataFrame::new();
35        };
36        let weight = data.column("weight");
37        let mut index: HashMap<String, usize> = HashMap::new();
38        let mut names: Vec<String> = Vec::new();
39        let mut id = |v: &Value| -> Option<usize> {
40            if v.is_na() {
41                return None;
42            }
43            let key = v.key_str().into_owned();
44            Some(*index.entry(key.clone()).or_insert_with(|| {
45                names.push(key);
46                names.len() - 1
47            }))
48        };
49        let mut pairs: indexmap::IndexMap<(usize, usize), f64> = indexmap::IndexMap::new();
50        for i in 0..data.nrows() {
51            let (Some(a), Some(b)) = (id(&from[i]), id(&to[i])) else {
52                continue;
53            };
54            let w = match weight {
55                Some(c) => match c[i].as_f64() {
56                    Some(w) if w.is_finite() && w > 0.0 => w,
57                    _ => continue,
58                },
59                None => 1.0,
60            };
61            if a != b {
62                *pairs.entry((a, b)).or_insert(0.0) += w;
63            }
64        }
65        let n = names.len();
66        if n == 0 {
67            return DataFrame::new();
68        }
69        let mut total = vec![0.0; n];
70        for (&(a, b), &w) in &pairs {
71            total[a] += w;
72            total[b] += w;
73        }
74        let mut order: Vec<usize> = (0..n).collect();
75        if self.sort_by_value {
76            order.sort_by(|&a, &b| total[b].total_cmp(&total[a]));
77        }
78        let mut pos = vec![0.0; n];
79        for (slot, &k) in order.iter().enumerate() {
80            pos[k] = slot as f64;
81        }
82        let span = |a: usize, b: usize| (pos[b] - pos[a]).abs() / 2.0;
83        let up = pairs
84            .keys()
85            .filter(|&&(a, b)| !self.directed_below || pos[b] > pos[a])
86            .map(|&(a, b)| span(a, b))
87            .fold(0.0, f64::max);
88        let down = pairs
89            .keys()
90            .filter(|&&(a, b)| self.directed_below && pos[b] < pos[a])
91            .map(|&(a, b)| span(a, b))
92            .fold(0.0, f64::max);
93        // Room under the line for the node labels.
94        let room = 0.12 * up.max(down).max(0.5);
95
96        let s = |x: &str| Value::Str(x.to_string());
97        let (mut part, mut x, mut xend, mut node, mut f, mut t, mut value) =
98            (vec![], vec![], vec![], vec![], vec![], vec![], vec![]);
99        for &k in &order {
100            part.push(s("node"));
101            x.push(pos[k]);
102            xend.push(pos[k]);
103            node.push(s(&names[k]));
104            f.push(s(&names[k]));
105            t.push(s(&names[k]));
106            value.push(total[k]);
107        }
108        for (&(a, b), &w) in &pairs {
109            part.push(s("link"));
110            x.push(pos[a]);
111            xend.push(pos[b]);
112            node.push(s(&names[a]));
113            f.push(s(&names[a]));
114            t.push(s(&names[b]));
115            value.push(w);
116        }
117        let rows = part.len();
118        let floats = |v: Vec<f64>| v.into_iter().map(Value::Float).collect::<Vec<_>>();
119        let mut out = DataFrame::new();
120        out.add_column("part".into(), part);
121        out.add_column("x".into(), floats(x));
122        out.add_column("xend".into(), floats(xend));
123        out.add_column("y".into(), vec![Value::Float(0.0); rows]);
124        out.add_column("ymin".into(), vec![Value::Float(-(down + room)); rows]);
125        out.add_column("ymax".into(), vec![Value::Float(up); rows]);
126        out.add_column("node".into(), node);
127        out.add_column("from".into(), f);
128        out.add_column("to".into(), t);
129        out.add_column("value".into(), floats(value));
130        out
131    }
132
133    fn required_aes(&self) -> Vec<Aesthetic> {
134        vec![Aesthetic::From, Aesthetic::To]
135    }
136
137    /// The layout needs every link of the panel at once.
138    fn panelwise(&self) -> bool {
139        true
140    }
141
142    fn name(&self) -> &str {
143        "arcdiagram"
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    fn edges(e: &[(&str, &str)]) -> DataFrame {
152        let mut d = DataFrame::new();
153        d.add_column(
154            "from".into(),
155            e.iter().map(|r| Value::Str(r.0.to_string())).collect(),
156        );
157        d.add_column(
158            "to".into(),
159            e.iter().map(|r| Value::Str(r.1.to_string())).collect(),
160        );
161        d
162    }
163
164    fn col(d: &DataFrame, c: &str) -> Vec<f64> {
165        d.column(c)
166            .unwrap()
167            .iter()
168            .map(|v| v.as_f64().unwrap())
169            .collect()
170    }
171
172    #[test]
173    fn positions_spans_and_degree() {
174        let d = edges(&[("A", "B"), ("A", "C"), ("C", "A"), ("D", "A"), ("B", "B")]);
175        let out = StatArcdiagram::default().compute_group(&d, &ScaleSet::new());
176        assert_eq!(out.nrows(), 4 + 4, "self-loop dropped");
177        assert_eq!(&col(&out, "x")[..4], &[0.0, 1.0, 2.0, 3.0]);
178        assert_eq!(col(&out, "value")[0], 4.0, "A: 2 out + 2 in");
179        // Tallest arc spans D(3) → A(0): radius 1.5.
180        assert_eq!(col(&out, "ymax")[0], 1.5);
181        assert!(col(&out, "ymin")[0] < 0.0);
182    }
183
184    #[test]
185    fn sort_and_directed_below() {
186        let d = edges(&[("B", "A"), ("C", "A"), ("C", "B")]);
187        let out = StatArcdiagram {
188            sort_by_value: true,
189            directed_below: true,
190        }
191        .compute_group(&d, &ScaleSet::new());
192        let names: Vec<String> = out.column("node").unwrap()[..3]
193            .iter()
194            .map(|v| v.key_str().into_owned())
195            .collect();
196        // Totals: A 2, B 2, C 2 → stable order of appearance (B, A, C).
197        assert_eq!(names, ["B", "A", "C"]);
198        // C(2) → A(1) and C(2) → B(0) run leftward: below the line.
199        assert_eq!(col(&out, "ymin")[0], -(1.0 + 0.12));
200        assert_eq!(col(&out, "ymax")[0], 0.5, "only B → A runs rightward");
201    }
202}