ggplot_rs/stat/
connect.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// How consecutive points are joined (ggplot2 4.0's `stat_connect(connection)`).
8#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
9pub enum Connection {
10    /// Horizontal then vertical (`"hv"`, a step that holds `y` until the next x).
11    #[default]
12    Hv,
13    /// Vertical then horizontal (`"vh"`).
14    Vh,
15    /// Step at the midpoint between x values (`"mid"`).
16    Mid,
17    /// Straight lines — the points unchanged (`"linear"`).
18    Linear,
19}
20
21/// Inserts the corner points that join each group's consecutive points (in x
22/// order) with `connection`, for a path geom — ggplot2's `stat_connect()`.
23#[derive(Clone, Debug, Default, PartialEq)]
24pub struct StatConnect {
25    pub connection: Connection,
26}
27
28impl StatConnect {
29    pub fn new(connection: Connection) -> Self {
30        StatConnect { connection }
31    }
32}
33
34impl Stat for StatConnect {
35    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
36        let (Some(xc), Some(yc)) = (data.column("x"), data.column("y")) else {
37            return DataFrame::new();
38        };
39        let mut idx: Vec<usize> = (0..data.nrows())
40            .filter(|&i| xc[i].as_f64().is_some() && yc[i].as_f64().is_some())
41            .collect();
42        idx.sort_by(|&a, &b| xc[a].as_f64().unwrap().total_cmp(&xc[b].as_f64().unwrap()));
43        // (source row whose other columns the point carries, x, y)
44        let mut pts: Vec<(usize, f64, f64)> = Vec::new();
45        for (k, &i) in idx.iter().enumerate() {
46            let (x, y) = (xc[i].as_f64().unwrap(), yc[i].as_f64().unwrap());
47            if let Some(&prev) = k.checked_sub(1).map(|p| &idx[p]) {
48                let (px, py) = (xc[prev].as_f64().unwrap(), yc[prev].as_f64().unwrap());
49                match self.connection {
50                    Connection::Hv => pts.push((prev, x, py)),
51                    Connection::Vh => pts.push((prev, px, y)),
52                    Connection::Mid => {
53                        let m = (px + x) / 2.0;
54                        pts.push((prev, m, py));
55                        pts.push((prev, m, y));
56                    }
57                    Connection::Linear => {}
58                }
59            }
60            // `mid` (as ggplot2) keeps only the end points of the original
61            // ones: the interior ones sit on the horizontal runs.
62            let interior = k > 0 && k + 1 < idx.len();
63            if !(self.connection == Connection::Mid && interior) {
64                pts.push((i, x, y));
65            }
66        }
67        let mut out = DataFrame::new();
68        for name in data.column_names() {
69            let src = data.column(name).unwrap();
70            let col: Vec<Value> = match name {
71                "x" => pts.iter().map(|p| Value::Float(p.1)).collect(),
72                "y" => pts.iter().map(|p| Value::Float(p.2)).collect(),
73                _ => pts.iter().map(|p| src[p.0].clone()).collect(),
74            };
75            out.add_column(name.to_string(), col);
76        }
77        out
78    }
79
80    fn required_aes(&self) -> Vec<Aesthetic> {
81        vec![Aesthetic::X, Aesthetic::Y]
82    }
83
84    fn name(&self) -> &str {
85        "connect"
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    fn run(c: Connection) -> Vec<(f64, f64)> {
94        let mut df = DataFrame::new();
95        df.add_column(
96            "x".into(),
97            vec![Value::Float(3.0), Value::Float(1.0), Value::Float(2.0)],
98        );
99        df.add_column(
100            "y".into(),
101            vec![Value::Float(30.0), Value::Float(10.0), Value::Float(20.0)],
102        );
103        let out = StatConnect::new(c).compute_group(&df, &ScaleSet::new());
104        let x = out.column("x").unwrap();
105        let y = out.column("y").unwrap();
106        (0..out.nrows())
107            .map(|i| (x[i].as_f64().unwrap(), y[i].as_f64().unwrap()))
108            .collect()
109    }
110
111    #[test]
112    fn connections_match_ggplot2() {
113        // ggplot2 4.0.3 layer_data(stat_connect(connection = …)) on
114        // (1,10),(2,20),(3,30). "linear" there repeats the interior points
115        // (same drawn line); we don't.
116        assert_eq!(
117            run(Connection::Hv),
118            [
119                (1.0, 10.0),
120                (2.0, 10.0),
121                (2.0, 20.0),
122                (3.0, 20.0),
123                (3.0, 30.0)
124            ]
125        );
126        assert_eq!(
127            run(Connection::Vh),
128            [
129                (1.0, 10.0),
130                (1.0, 20.0),
131                (2.0, 20.0),
132                (2.0, 30.0),
133                (3.0, 30.0)
134            ]
135        );
136        assert_eq!(
137            run(Connection::Mid),
138            [
139                (1.0, 10.0),
140                (1.5, 10.0),
141                (1.5, 20.0),
142                (2.5, 20.0),
143                (2.5, 30.0),
144                (3.0, 30.0)
145            ]
146        );
147        assert_eq!(
148            run(Connection::Linear),
149            [(1.0, 10.0), (2.0, 20.0), (3.0, 30.0)]
150        );
151    }
152}
ggplot_rs/stat/connect.rs

ggplot_rs/stat/
connect.rs
