Skip to main content

ggplot_rs/stat/
network.rs

1//! `StatNetwork`: node-link layouts for network graphs.
2//!
3//! Input: an edge list (`from`, `to`, optional `weight`, 1 when unmapped;
4//! repeated pairs summed, self-loops dropped). Nodes are laid out by a
5//! seeded Fruchterman–Reingold force simulation (deterministic for a seed),
6//! on a circle, or on a grid.
7//!
8//! Output: one row per node (`part = "node"`, `x`/`y`, `node`, `degree` =
9//! summed incident weight, `component` = connected-component number) then
10//! one per edge (`part = "edge"`, `x`/`y` → `xend`/`yend`, `from`, `to`,
11//! `value`), plus constant `xmin`/`xmax`/`ymin`/`ymax` framing a square.
12
13use std::collections::HashMap;
14use std::f64::consts::TAU;
15
16use crate::aes::Aesthetic;
17use crate::data::{DataFrame, Value};
18use crate::rng::SplitMix64;
19use crate::scale::ScaleSet;
20
21use super::Stat;
22
23/// Node placement.
24#[derive(Clone, Copy, Debug, PartialEq)]
25pub enum NetworkLayout {
26    /// Fruchterman–Reingold forces from a seeded start.
27    ForceDirected { iterations: usize, seed: u64 },
28    /// Nodes evenly on a circle, in first-appearance order.
29    Circle,
30    /// Nodes on a square grid, row by row.
31    Grid,
32}
33
34impl Default for NetworkLayout {
35    fn default() -> Self {
36        NetworkLayout::ForceDirected {
37            iterations: 300,
38            seed: 42,
39        }
40    }
41}
42
43/// Network layout stat behind [`GeomNetwork`](crate::geom::network::GeomNetwork).
44#[derive(Clone, Debug, Default)]
45pub struct StatNetwork {
46    pub layout: NetworkLayout,
47}
48
49/// Fruchterman–Reingold in the unit square, with a weak pull to the centre
50/// so disconnected components stay together.
51pub(crate) fn force_directed(
52    n: usize,
53    edges: &[(usize, usize, f64)],
54    iterations: usize,
55    seed: u64,
56) -> Vec<(f64, f64)> {
57    let mut rng = SplitMix64::new(seed);
58    let mut pos: Vec<(f64, f64)> = (0..n)
59        .map(|_| (rng.range_f64(-0.5, 0.5), rng.range_f64(-0.5, 0.5)))
60        .collect();
61    if n < 2 {
62        return vec![(0.0, 0.0); n];
63    }
64    let k = (1.0 / n as f64).sqrt();
65    let iterations = iterations.max(1);
66    for it in 0..iterations {
67        let temp = 0.1 * (1.0 - it as f64 / iterations as f64) + 1e-4;
68        let mut disp = vec![(0.0, 0.0); n];
69        for i in 0..n {
70            for j in i + 1..n {
71                let (dx, dy) = (pos[i].0 - pos[j].0, pos[i].1 - pos[j].1);
72                let d = dx.hypot(dy).max(1e-3);
73                let f = k * k / d;
74                let (fx, fy) = (dx / d * f, dy / d * f);
75                disp[i].0 += fx;
76                disp[i].1 += fy;
77                disp[j].0 -= fx;
78                disp[j].1 -= fy;
79            }
80        }
81        for &(a, b, _) in edges {
82            let (dx, dy) = (pos[a].0 - pos[b].0, pos[a].1 - pos[b].1);
83            let d = dx.hypot(dy).max(1e-3);
84            let f = d * d / k;
85            let (fx, fy) = (dx / d * f, dy / d * f);
86            disp[a].0 -= fx;
87            disp[a].1 -= fy;
88            disp[b].0 += fx;
89            disp[b].1 += fy;
90        }
91        for i in 0..n {
92            disp[i].0 -= 0.5 * k * pos[i].0;
93            disp[i].1 -= 0.5 * k * pos[i].1;
94            let len = disp[i].0.hypot(disp[i].1);
95            if len > 0.0 {
96                let step = len.min(temp);
97                pos[i].0 += disp[i].0 / len * step;
98                pos[i].1 += disp[i].1 / len * step;
99            }
100        }
101    }
102    pos
103}
104
105/// Force-directed layout per connected component (so components don't
106/// drift apart), each scaled by its node count and packed in shelf rows.
107fn layout_components(
108    n: usize,
109    edges: &[(usize, usize, f64)],
110    comp: &[usize],
111    iterations: usize,
112    seed: u64,
113) -> Vec<(f64, f64)> {
114    let n_comp = comp.iter().copied().max().unwrap_or(0);
115    let mut pos = vec![(0.0, 0.0); n];
116    // (component, members, size) — larger components first.
117    let mut parts: Vec<(Vec<usize>, f64)> = (1..=n_comp)
118        .map(|c| {
119            let members: Vec<usize> = (0..n).filter(|&i| comp[i] == c).collect();
120            let size = (members.len() as f64).sqrt();
121            (members, size)
122        })
123        .collect();
124    parts.sort_by(|a, b| b.1.total_cmp(&a.1));
125    let row_width = parts.iter().map(|p| p.1 * p.1).sum::<f64>().sqrt() * 1.3;
126    let (mut x, mut y, mut row_h) = (0.0, 0.0, 0.0f64);
127    for (k, (members, size)) in parts.iter().enumerate() {
128        let size = *size;
129        let local: std::collections::HashMap<usize, usize> =
130            members.iter().enumerate().map(|(j, &i)| (i, j)).collect();
131        let sub: Vec<(usize, usize, f64)> = edges
132            .iter()
133            .filter_map(|&(a, b, w)| Some((*local.get(&a)?, *local.get(&b)?, w)))
134            .collect();
135        let p = force_directed(members.len(), &sub, iterations, seed.wrapping_add(k as u64));
136        // Normalise to a box of side `size` (a single node is a point).
137        let (mut x0, mut x1, mut y0, mut y1) = (f64::MAX, f64::MIN, f64::MAX, f64::MIN);
138        for &(px, py) in &p {
139            x0 = x0.min(px);
140            x1 = x1.max(px);
141            y0 = y0.min(py);
142            y1 = y1.max(py);
143        }
144        let span = (x1 - x0).max(y1 - y0);
145        let scale = if span > 0.0 { size / span } else { 0.0 };
146        let side = if span > 0.0 { size } else { 0.0 };
147        if x > 0.0 && x + side > row_width {
148            x = 0.0;
149            y -= row_h + 0.6;
150            row_h = 0.0;
151        }
152        for (&i, &(px, py)) in members.iter().zip(&p) {
153            pos[i] = (x + (px - x0) * scale, y - (py - y0) * scale);
154        }
155        x += side + 0.6;
156        row_h = row_h.max(side);
157    }
158    pos
159}
160
161/// Connected-component number (1-based, in first-appearance order) per node.
162pub(crate) fn components(n: usize, edges: &[(usize, usize, f64)]) -> Vec<usize> {
163    let mut parent: Vec<usize> = (0..n).collect();
164    fn find(p: &mut [usize], mut x: usize) -> usize {
165        while p[x] != x {
166            p[x] = p[p[x]];
167            x = p[x];
168        }
169        x
170    }
171    for &(a, b, _) in edges {
172        let (ra, rb) = (find(&mut parent, a), find(&mut parent, b));
173        if ra != rb {
174            parent[ra.max(rb)] = ra.min(rb);
175        }
176    }
177    let mut label: HashMap<usize, usize> = HashMap::new();
178    (0..n)
179        .map(|i| {
180            let r = find(&mut parent, i);
181            let next = label.len() + 1;
182            *label.entry(r).or_insert(next)
183        })
184        .collect()
185}
186
187impl Stat for StatNetwork {
188    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
189        let (Some(from), Some(to)) = (data.column("from"), data.column("to")) else {
190            return DataFrame::new();
191        };
192        let weight = data.column("weight");
193        let mut index: HashMap<String, usize> = HashMap::new();
194        let mut names: Vec<String> = Vec::new();
195        let mut id = |v: &Value| -> Option<usize> {
196            if v.is_na() {
197                return None;
198            }
199            let key = v.key_str().into_owned();
200            Some(*index.entry(key.clone()).or_insert_with(|| {
201                names.push(key);
202                names.len() - 1
203            }))
204        };
205        let mut pairs: indexmap::IndexMap<(usize, usize), f64> = indexmap::IndexMap::new();
206        for i in 0..data.nrows() {
207            let (Some(a), Some(b)) = (id(&from[i]), id(&to[i])) else {
208                continue;
209            };
210            let w = match weight {
211                Some(c) => match c[i].as_f64() {
212                    Some(w) if w.is_finite() && w > 0.0 => w,
213                    _ => continue,
214                },
215                None => 1.0,
216            };
217            if a != b {
218                *pairs.entry((a, b)).or_insert(0.0) += w;
219            }
220        }
221        let n = names.len();
222        if n == 0 {
223            return DataFrame::new();
224        }
225        let edges: Vec<(usize, usize, f64)> =
226            pairs.into_iter().map(|((a, b), w)| (a, b, w)).collect();
227        let comp = components(n, &edges);
228        let pos = match self.layout {
229            NetworkLayout::ForceDirected { iterations, seed } => {
230                layout_components(n, &edges, &comp, iterations, seed)
231            }
232            NetworkLayout::Circle => (0..n)
233                .map(|i| {
234                    let a = TAU * i as f64 / n as f64;
235                    (0.5 * a.sin(), 0.5 * a.cos())
236                })
237                .collect(),
238            NetworkLayout::Grid => {
239                let cols = (n as f64).sqrt().ceil() as usize;
240                (0..n)
241                    .map(|i| ((i % cols) as f64, -((i / cols) as f64)))
242                    .collect()
243            }
244        };
245        let mut degree = vec![0.0; n];
246        for &(a, b, w) in &edges {
247            degree[a] += w;
248            degree[b] += w;
249        }
250        // Square frame around the layout with a margin for nodes/labels.
251        let (mut x0, mut x1, mut y0, mut y1) = (f64::MAX, f64::MIN, f64::MAX, f64::MIN);
252        for &(x, y) in &pos {
253            x0 = x0.min(x);
254            x1 = x1.max(x);
255            y0 = y0.min(y);
256            y1 = y1.max(y);
257        }
258        let half = ((x1 - x0).max(y1 - y0) / 2.0).max(0.5) * 1.12;
259        let (cx, cy) = ((x0 + x1) / 2.0, (y0 + y1) / 2.0);
260
261        let s = |x: &str| Value::Str(x.to_string());
262        let rows = n + edges.len();
263        let (mut part, mut x, mut y, mut xe, mut ye) = (vec![], vec![], vec![], vec![], vec![]);
264        let (mut node, mut f, mut t, mut value, mut deg, mut cmp) =
265            (vec![], vec![], vec![], vec![], vec![], vec![]);
266        for i in 0..n {
267            part.push(s("node"));
268            x.push(pos[i].0);
269            y.push(pos[i].1);
270            xe.push(pos[i].0);
271            ye.push(pos[i].1);
272            node.push(s(&names[i]));
273            f.push(s(&names[i]));
274            t.push(s(&names[i]));
275            value.push(Value::Float(degree[i]));
276            deg.push(Value::Float(degree[i]));
277            cmp.push(s(&comp[i].to_string()));
278        }
279        for &(a, b, w) in &edges {
280            part.push(s("edge"));
281            x.push(pos[a].0);
282            y.push(pos[a].1);
283            xe.push(pos[b].0);
284            ye.push(pos[b].1);
285            node.push(s(&names[a]));
286            f.push(s(&names[a]));
287            t.push(s(&names[b]));
288            value.push(Value::Float(w));
289            deg.push(Value::Float(degree[a]));
290            cmp.push(s(&comp[a].to_string()));
291        }
292        let fl = |v: Vec<f64>| v.into_iter().map(Value::Float).collect::<Vec<_>>();
293        let mut out = DataFrame::new();
294        out.add_column("part".into(), part);
295        out.add_column("x".into(), fl(x));
296        out.add_column("y".into(), fl(y));
297        out.add_column("xend".into(), fl(xe));
298        out.add_column("yend".into(), fl(ye));
299        out.add_column("node".into(), node);
300        out.add_column("from".into(), f);
301        out.add_column("to".into(), t);
302        out.add_column("value".into(), value);
303        out.add_column("degree".into(), deg);
304        out.add_column("component".into(), cmp);
305        for (c, v) in [
306            ("xmin", cx - half),
307            ("xmax", cx + half),
308            ("ymin", cy - half),
309            ("ymax", cy + half),
310        ] {
311            out.add_column(c.into(), vec![Value::Float(v); rows]);
312        }
313        out
314    }
315
316    fn required_aes(&self) -> Vec<Aesthetic> {
317        vec![Aesthetic::From, Aesthetic::To]
318    }
319
320    /// The layout needs every edge of the panel at once.
321    fn panelwise(&self) -> bool {
322        true
323    }
324
325    fn name(&self) -> &str {
326        "network"
327    }
328}
329
330#[cfg(test)]
331mod tests {
332    use super::*;
333
334    #[test]
335    fn force_layout_is_deterministic_and_separates_nodes() {
336        let edges = vec![(0, 1, 1.0), (1, 2, 1.0), (2, 0, 1.0), (3, 4, 1.0)];
337        let a = force_directed(5, &edges, 200, 7);
338        let b = force_directed(5, &edges, 200, 7);
339        assert_eq!(a, b, "same seed, same layout");
340        assert_ne!(a, force_directed(5, &edges, 200, 8));
341        for i in 0..5 {
342            for j in i + 1..5 {
343                let d = (a[i].0 - a[j].0).hypot(a[i].1 - a[j].1);
344                assert!(d > 0.05, "nodes {i},{j} collapsed: {d}");
345            }
346        }
347        // Linked nodes end up closer than the average pair.
348        let dist = |i: usize, j: usize| (a[i].0 - a[j].0).hypot(a[i].1 - a[j].1);
349        assert!(dist(3, 4) < dist(0, 3).max(dist(1, 4)));
350    }
351
352    #[test]
353    fn components_and_degrees() {
354        assert_eq!(
355            components(5, &[(0, 1, 1.0), (3, 4, 1.0)]),
356            vec![1, 1, 2, 3, 3]
357        );
358        let mut d = DataFrame::new();
359        let s = |v: &[&str]| v.iter().map(|x| Value::Str(x.to_string())).collect();
360        d.add_column("from".into(), s(&["a", "a", "b", "c"]));
361        d.add_column("to".into(), s(&["b", "c", "c", "c"]));
362        let out = StatNetwork {
363            layout: NetworkLayout::Circle,
364        }
365        .compute_group(&d, &ScaleSet::new());
366        assert_eq!(out.nrows(), 3 + 3, "self-loop dropped");
367        let deg: Vec<f64> = out.column("degree").unwrap()[..3]
368            .iter()
369            .map(|v| v.as_f64().unwrap())
370            .collect();
371        assert_eq!(deg, vec![2.0, 2.0, 2.0]);
372    }
373}