1use 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#[derive(Clone, Copy, Debug, PartialEq)]
25pub enum NetworkLayout {
26 ForceDirected { iterations: usize, seed: u64 },
28 Circle,
30 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#[derive(Clone, Debug, Default)]
45pub struct StatNetwork {
46 pub layout: NetworkLayout,
47}
48
49pub(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
105fn 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 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 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
161pub(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 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 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 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}