ggplot_rs/stat/
contour.rs

1use crate::aes::Aesthetic;
2use crate::data::DataFrame;
3use crate::scale::ScaleSet;
4
5use super::marching_squares::extract_contours;
6use super::Stat;
7
8/// Compute contour lines from gridded (x, y, z) data using Marching Squares.
9pub struct StatContour {
10    pub bins: usize,
11    pub n_levels: usize,
12}
13
14impl Default for StatContour {
15    fn default() -> Self {
16        StatContour {
17            bins: 50,
18            n_levels: 10,
19        }
20    }
21}
22
23impl Stat for StatContour {
24    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
25        let x_col = match data.column("x") {
26            Some(c) => c,
27            None => return DataFrame::new(),
28        };
29        let y_col = match data.column("y") {
30            Some(c) => c,
31            None => return DataFrame::new(),
32        };
33        let z_col = match data.column("z") {
34            Some(c) => c,
35            None => return DataFrame::new(),
36        };
37
38        let points: Vec<(f64, f64, f64)> = x_col
39            .iter()
40            .zip(y_col.iter())
41            .zip(z_col.iter())
42            .filter_map(|((x, y), z)| match (x.as_f64(), y.as_f64(), z.as_f64()) {
43                (Some(xv), Some(yv), Some(zv))
44                    if xv.is_finite() && yv.is_finite() && zv.is_finite() =>
45                {
46                    Some((xv, yv, zv))
47                }
48                _ => None,
49            })
50            .collect();
51
52        if points.is_empty() {
53            return DataFrame::new();
54        }
55
56        let x_min = points.iter().map(|p| p.0).fold(f64::INFINITY, f64::min);
57        let x_max = points.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max);
58        let y_min = points.iter().map(|p| p.1).fold(f64::INFINITY, f64::min);
59        let y_max = points.iter().map(|p| p.1).fold(f64::NEG_INFINITY, f64::max);
60        let z_min = points.iter().map(|p| p.2).fold(f64::INFINITY, f64::min);
61        let z_max = points.iter().map(|p| p.2).fold(f64::NEG_INFINITY, f64::max);
62
63        if (x_max - x_min).abs() < f64::EPSILON || (y_max - y_min).abs() < f64::EPSILON {
64            return DataFrame::new();
65        }
66
67        let nx = self.bins;
68        let ny = self.bins;
69        let dx = (x_max - x_min) / nx as f64;
70        let dy = (y_max - y_min) / ny as f64;
71
72        let grid = build_grid_from_xyz(&points, nx, ny, x_min, y_min, dx, dy, z_min, z_max);
73
74        let z_range = z_max - z_min;
75        if z_range.abs() < f64::EPSILON {
76            return DataFrame::new();
77        }
78
79        let n_levels = self.n_levels;
80        let level_step = z_range / (n_levels + 1) as f64;
81        let levels: Vec<f64> = (1..=n_levels)
82            .map(|li| z_min + li as f64 * level_step)
83            .collect();
84
85        let (x_vals, y_vals, level_vals, group_vals) =
86            extract_contours(&grid, nx, ny, x_min, y_min, dx, dy, &levels);
87
88        let mut result = DataFrame::new();
89        result.add_column("x".to_string(), x_vals);
90        result.add_column("y".to_string(), y_vals);
91        result.add_column("level".to_string(), level_vals);
92        result.add_column("group".to_string(), group_vals);
93        result
94    }
95
96    fn required_aes(&self) -> Vec<Aesthetic> {
97        vec![Aesthetic::X, Aesthetic::Y]
98    }
99
100    fn name(&self) -> &str {
101        "contour"
102    }
103}
104
105/// Build a regular grid from scattered (x, y, z) points via nearest-neighbor binning.
106#[allow(clippy::too_many_arguments)]
107pub(crate) fn build_grid_from_xyz(
108    points: &[(f64, f64, f64)],
109    nx: usize,
110    ny: usize,
111    x_min: f64,
112    y_min: f64,
113    dx: f64,
114    dy: f64,
115    z_min: f64,
116    z_max: f64,
117) -> Vec<f64> {
118    let mut grid = vec![0.0f64; (nx + 1) * (ny + 1)];
119    let mut counts = vec![0usize; (nx + 1) * (ny + 1)];
120
121    for &(x, y, z) in points {
122        let ix = ((x - x_min) / dx).round() as usize;
123        let iy = ((y - y_min) / dy).round() as usize;
124        let ix = ix.min(nx);
125        let iy = iy.min(ny);
126        let idx = iy * (nx + 1) + ix;
127        grid[idx] += z;
128        counts[idx] += 1;
129    }
130
131    for i in 0..grid.len() {
132        if counts[i] > 0 {
133            grid[i] /= counts[i] as f64;
134        }
135    }
136
137    // Fill empty cells with nearest filled neighbor
138    for iy in 0..=ny {
139        for ix in 0..=nx {
140            let idx = iy * (nx + 1) + ix;
141            if counts[idx] == 0 {
142                let mut best_dist = f64::INFINITY;
143                let mut best_val = (z_min + z_max) / 2.0;
144                for dy2 in -3i32..=3 {
145                    for dx2 in -3i32..=3 {
146                        let nx2 = ix as i32 + dx2;
147                        let ny2 = iy as i32 + dy2;
148                        if nx2 >= 0 && nx2 <= nx as i32 && ny2 >= 0 && ny2 <= ny as i32 {
149                            let idx2 = ny2 as usize * (nx + 1) + nx2 as usize;
150                            if counts[idx2] > 0 {
151                                let d = ((dx2 * dx2 + dy2 * dy2) as f64).sqrt();
152                                if d < best_dist {
153                                    best_dist = d;
154                                    best_val = grid[idx2];
155                                }
156                            }
157                        }
158                    }
159                }
160                grid[idx] = best_val;
161            }
162        }
163    }
164
165    grid
166}
ggplot_rs/stat/contour.rs

ggplot_rs/stat/
contour.rs
