1use std::collections::HashMap;
15
16use crate::aes::Aesthetic;
17use crate::data::{DataFrame, Value};
18use crate::scale::ScaleSet;
19
20use super::Stat;
21
22#[derive(Clone, Debug, Default)]
24pub struct StatArcdiagram {
25 pub sort_by_value: bool,
27 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 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 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 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 assert_eq!(names, ["B", "A", "C"]);
198 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}