plotkit_core/charts/
contour.rs

1//! Contour and filled contour chart builder methods.
2//!
3//! Provides a fluent builder API for configuring [`ContourArtist`] instances.
4//! Each method returns `&mut Self`, allowing calls to be chained together
5//! for concise, readable chart construction.
6
7use crate::artist::ContourArtist;
8use crate::colormap::Colormap;
9use crate::primitives::Color;
10
11impl ContourArtist {
12    /// Sets the colormap used to map contour levels to colors.
13    pub fn colormap(&mut self, cmap: Colormap) -> &mut Self {
14        self.cmap = cmap;
15        self
16    }
17
18    /// Sets explicit contour levels.
19    ///
20    /// When `None`, levels are automatically computed by evenly spacing
21    /// between the minimum and maximum finite z values.
22    pub fn levels(&mut self, levels: Vec<f64>) -> &mut Self {
23        self.levels = Some(levels);
24        self
25    }
26
27    /// Sets the number of automatically computed contour levels.
28    ///
29    /// Only used when explicit levels are not set. Defaults to 10.
30    pub fn num_levels(&mut self, n: usize) -> &mut Self {
31        self.num_levels = n;
32        self
33    }
34
35    /// Sets the line width for contour lines (unfilled mode).
36    pub fn linewidths(&mut self, w: f64) -> &mut Self {
37        self.linewidths = w;
38        self
39    }
40
41    /// Sets explicit colors for contour levels, overriding the colormap.
42    pub fn colors(&mut self, colors: Vec<Color>) -> &mut Self {
43        self.colors = Some(colors);
44        self
45    }
46
47    /// Sets the legend label.
48    pub fn label(&mut self, label: &str) -> &mut Self {
49        self.label = Some(label.to_string());
50        self
51    }
52
53    /// Returns the effective contour levels.
54    ///
55    /// If explicit levels are set, returns those. Otherwise, computes
56    /// `num_levels` evenly spaced levels between the minimum and maximum
57    /// finite z values.
58    pub fn effective_levels(&self) -> Vec<f64> {
59        if let Some(ref lvls) = self.levels {
60            return lvls.clone();
61        }
62        let (zmin, zmax) = self.z_bounds();
63        if (zmax - zmin).abs() < f64::EPSILON {
64            return vec![zmin];
65        }
66        let n = self.num_levels.max(2);
67        (0..n)
68            .map(|i| zmin + (zmax - zmin) * (i as f64) / ((n - 1) as f64))
69            .collect()
70    }
71
72    /// Returns `(zmin, zmax)` from the finite z values.
73    ///
74    /// Falls back to `(0.0, 1.0)` when the data contains no finite values.
75    pub fn z_bounds(&self) -> (f64, f64) {
76        let mut lo = f64::INFINITY;
77        let mut hi = f64::NEG_INFINITY;
78        for row in &self.z {
79            for &val in row {
80                if val.is_finite() {
81                    if val < lo {
82                        lo = val;
83                    }
84                    if val > hi {
85                        hi = val;
86                    }
87                }
88            }
89        }
90        if lo.is_finite() && hi.is_finite() {
91            (lo, hi)
92        } else {
93            (0.0, 1.0)
94        }
95    }
96
97    /// Computes contour line segments for a single level using marching squares.
98    ///
99    /// Returns a list of line segments, each represented as `(x0, y0, x1, y1)`.
100    pub fn marching_squares(&self, level: f64) -> Vec<(f64, f64, f64, f64)> {
101        let ny = self.y.len();
102        let nx = self.x.len();
103        if ny < 2 || nx < 2 || self.z.len() < 2 {
104            return Vec::new();
105        }
106
107        let mut segments = Vec::new();
108
109        for j in 0..(ny - 1) {
110            if j + 1 >= self.z.len() {
111                break;
112            }
113            let row_bot = &self.z[j];
114            let row_top = &self.z[j + 1];
115            for i in 0..(nx - 1) {
116                if i + 1 >= row_bot.len() || i + 1 >= row_top.len() {
117                    break;
118                }
119
120                // Four corners of the cell:
121                // bl = bottom-left, br = bottom-right, tr = top-right, tl = top-left
122                let bl = row_bot[i];
123                let br = row_bot[i + 1];
124                let tr = row_top[i + 1];
125                let tl = row_top[i];
126
127                // Skip cells with NaN values.
128                if !bl.is_finite() || !br.is_finite() || !tr.is_finite() || !tl.is_finite() {
129                    continue;
130                }
131
132                // Classify corners: 1 if above level, 0 if below.
133                let case = ((bl >= level) as u8)
134                    | (((br >= level) as u8) << 1)
135                    | (((tr >= level) as u8) << 2)
136                    | (((tl >= level) as u8) << 3);
137
138                // Cell coordinates.
139                let x0 = self.x[i];
140                let x1 = self.x[i + 1];
141                let y0 = self.y[j];
142                let y1 = self.y[j + 1];
143
144                // Interpolation helper: fraction along edge where level crosses.
145                let interp = |a: f64, b: f64| -> f64 {
146                    if (b - a).abs() < f64::EPSILON {
147                        0.5
148                    } else {
149                        (level - a) / (b - a)
150                    }
151                };
152
153                // Edge midpoints (interpolated):
154                // bottom: between bl and br
155                // right:  between br and tr
156                // top:    between tl and tr
157                // left:   between bl and tl
158                let bottom = || {
159                    let t = interp(bl, br);
160                    (x0 + t * (x1 - x0), y0)
161                };
162                let right = || {
163                    let t = interp(br, tr);
164                    (x1, y0 + t * (y1 - y0))
165                };
166                let top = || {
167                    let t = interp(tl, tr);
168                    (x0 + t * (x1 - x0), y1)
169                };
170                let left = || {
171                    let t = interp(bl, tl);
172                    (x0, y0 + t * (y1 - y0))
173                };
174
175                // 16 marching squares cases.
176                match case {
177                    0 | 15 => { /* entirely below or above -- no contour */ }
178                    1 | 14 => {
179                        let (ax, ay) = bottom();
180                        let (bx, by) = left();
181                        segments.push((ax, ay, bx, by));
182                    }
183                    2 | 13 => {
184                        let (ax, ay) = bottom();
185                        let (bx, by) = right();
186                        segments.push((ax, ay, bx, by));
187                    }
188                    3 | 12 => {
189                        let (ax, ay) = left();
190                        let (bx, by) = right();
191                        segments.push((ax, ay, bx, by));
192                    }
193                    4 | 11 => {
194                        let (ax, ay) = right();
195                        let (bx, by) = top();
196                        segments.push((ax, ay, bx, by));
197                    }
198                    5 | 10 => {
199                        // Ambiguous saddle point -- use simple decomposition
200                        // (no saddle disambiguation for v0.4).
201                        let (ax, ay) = bottom();
202                        let (bx, by) = left();
203                        segments.push((ax, ay, bx, by));
204                        let (cx, cy) = right();
205                        let (dx, dy) = top();
206                        segments.push((cx, cy, dx, dy));
207                    }
208                    6 | 9 => {
209                        let (ax, ay) = bottom();
210                        let (bx, by) = top();
211                        segments.push((ax, ay, bx, by));
212                    }
213                    7 | 8 => {
214                        let (ax, ay) = left();
215                        let (bx, by) = top();
216                        segments.push((ax, ay, bx, by));
217                    }
218                    _ => unreachable!(),
219                }
220            }
221        }
222
223        segments
224    }
225
226    /// Computes the average z value for each grid cell.
227    ///
228    /// Returns a 2D vector of shape `[ny-1][nx-1]` where each element is the
229    /// mean of the four corner z values. Used for the simplified filled
230    /// contour approach.
231    pub fn cell_averages(&self) -> Vec<Vec<f64>> {
232        let ny = self.y.len();
233        let nx = self.x.len();
234        if ny < 2 || nx < 2 || self.z.len() < 2 {
235            return Vec::new();
236        }
237
238        let mut avgs = Vec::with_capacity(ny - 1);
239        for j in 0..(ny - 1) {
240            if j + 1 >= self.z.len() {
241                break;
242            }
243            let mut row = Vec::with_capacity(nx - 1);
244            let row_bot = &self.z[j];
245            let row_top = &self.z[j + 1];
246            for i in 0..(nx - 1) {
247                if i + 1 >= row_bot.len() || i + 1 >= row_top.len() {
248                    break;
249                }
250                let bl = row_bot[i];
251                let br = row_bot[i + 1];
252                let tl = row_top[i];
253                let tr = row_top[i + 1];
254                let vals = [bl, br, tl, tr];
255                let finite: Vec<f64> = vals.iter().copied().filter(|v| v.is_finite()).collect();
256                let avg = if finite.is_empty() {
257                    f64::NAN
258                } else {
259                    finite.iter().sum::<f64>() / finite.len() as f64
260                };
261                row.push(avg);
262            }
263            avgs.push(row);
264        }
265        avgs
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use crate::artist::ContourArtist;
272    use crate::colormap::Colormap;
273    use crate::primitives::Color;
274
275    fn sample_contour() -> ContourArtist {
276        ContourArtist {
277            x: vec![0.0, 1.0, 2.0],
278            y: vec![0.0, 1.0, 2.0],
279            z: vec![
280                vec![0.0, 1.0, 2.0],
281                vec![1.0, 2.0, 3.0],
282                vec![2.0, 3.0, 4.0],
283            ],
284            levels: None,
285            filled: false,
286            cmap: Colormap::Viridis,
287            colors: None,
288            linewidths: 1.0,
289            label: None,
290            color: Color::TAB_BLUE,
291            num_levels: 10,
292        }
293    }
294
295    #[test]
296    fn builder_colormap() {
297        let mut c = sample_contour();
298        c.colormap(Colormap::Plasma);
299        assert_eq!(c.cmap, Colormap::Plasma);
300    }
301
302    #[test]
303    fn builder_levels() {
304        let mut c = sample_contour();
305        c.levels(vec![1.0, 2.0, 3.0]);
306        assert_eq!(c.levels, Some(vec![1.0, 2.0, 3.0]));
307    }
308
309    #[test]
310    fn builder_num_levels() {
311        let mut c = sample_contour();
312        c.num_levels(5);
313        assert_eq!(c.num_levels, 5);
314    }
315
316    #[test]
317    fn builder_linewidths() {
318        let mut c = sample_contour();
319        c.linewidths(2.5);
320        assert!((c.linewidths - 2.5).abs() < f64::EPSILON);
321    }
322
323    #[test]
324    fn builder_colors() {
325        let mut c = sample_contour();
326        c.colors(vec![Color::TAB_RED, Color::TAB_BLUE]);
327        assert_eq!(c.colors.as_ref().unwrap().len(), 2);
328    }
329
330    #[test]
331    fn builder_label() {
332        let mut c = sample_contour();
333        c.label("my contour");
334        assert_eq!(c.label.as_deref(), Some("my contour"));
335    }
336
337    #[test]
338    fn builder_chaining() {
339        let mut c = sample_contour();
340        c.colormap(Colormap::Plasma)
341            .linewidths(2.0)
342            .num_levels(5)
343            .label("chained");
344        assert_eq!(c.cmap, Colormap::Plasma);
345        assert!((c.linewidths - 2.0).abs() < f64::EPSILON);
346        assert_eq!(c.num_levels, 5);
347        assert_eq!(c.label.as_deref(), Some("chained"));
348    }
349
350    #[test]
351    fn z_bounds_basic() {
352        let c = sample_contour();
353        let (lo, hi) = c.z_bounds();
354        assert!((lo - 0.0).abs() < f64::EPSILON);
355        assert!((hi - 4.0).abs() < f64::EPSILON);
356    }
357
358    #[test]
359    fn z_bounds_empty() {
360        let c = ContourArtist {
361            x: vec![], y: vec![], z: vec![],
362            levels: None, filled: false, cmap: Colormap::Viridis,
363            colors: None, linewidths: 1.0, label: None,
364            color: Color::TAB_BLUE, num_levels: 10,
365        };
366        let (lo, hi) = c.z_bounds();
367        assert!((lo - 0.0).abs() < f64::EPSILON);
368        assert!((hi - 1.0).abs() < f64::EPSILON);
369    }
370
371    #[test]
372    fn z_bounds_with_nan() {
373        let c = ContourArtist {
374            x: vec![0.0, 1.0],
375            y: vec![0.0, 1.0],
376            z: vec![
377                vec![f64::NAN, 3.0],
378                vec![1.0, f64::NAN],
379            ],
380            levels: None, filled: false, cmap: Colormap::Viridis,
381            colors: None, linewidths: 1.0, label: None,
382            color: Color::TAB_BLUE, num_levels: 10,
383        };
384        let (lo, hi) = c.z_bounds();
385        assert!((lo - 1.0).abs() < f64::EPSILON);
386        assert!((hi - 3.0).abs() < f64::EPSILON);
387    }
388
389    #[test]
390    fn effective_levels_auto() {
391        let c = sample_contour();
392        let levels = c.effective_levels();
393        assert_eq!(levels.len(), 10);
394        assert!((levels[0] - 0.0).abs() < f64::EPSILON);
395        assert!((levels[9] - 4.0).abs() < f64::EPSILON);
396    }
397
398    #[test]
399    fn effective_levels_explicit() {
400        let mut c = sample_contour();
401        c.levels(vec![1.0, 2.0, 3.0]);
402        let levels = c.effective_levels();
403        assert_eq!(levels, vec![1.0, 2.0, 3.0]);
404    }
405
406    #[test]
407    fn effective_levels_constant_z() {
408        let c = ContourArtist {
409            x: vec![0.0, 1.0],
410            y: vec![0.0, 1.0],
411            z: vec![vec![5.0, 5.0], vec![5.0, 5.0]],
412            levels: None, filled: false, cmap: Colormap::Viridis,
413            colors: None, linewidths: 1.0, label: None,
414            color: Color::TAB_BLUE, num_levels: 10,
415        };
416        let levels = c.effective_levels();
417        assert_eq!(levels.len(), 1);
418        assert!((levels[0] - 5.0).abs() < f64::EPSILON);
419    }
420
421    #[test]
422    fn data_bounds_basic() {
423        let c = sample_contour();
424        let (xmin, xmax, ymin, ymax) = c.data_bounds();
425        assert!((xmin - 0.0).abs() < f64::EPSILON);
426        assert!((xmax - 2.0).abs() < f64::EPSILON);
427        assert!((ymin - 0.0).abs() < f64::EPSILON);
428        assert!((ymax - 2.0).abs() < f64::EPSILON);
429    }
430
431    #[test]
432    fn data_bounds_empty() {
433        let c = ContourArtist {
434            x: vec![], y: vec![], z: vec![],
435            levels: None, filled: false, cmap: Colormap::Viridis,
436            colors: None, linewidths: 1.0, label: None,
437            color: Color::TAB_BLUE, num_levels: 10,
438        };
439        assert_eq!(c.data_bounds(), (0.0, 1.0, 0.0, 1.0));
440    }
441
442    #[test]
443    fn data_bounds_single_point() {
444        let c = ContourArtist {
445            x: vec![5.0], y: vec![3.0], z: vec![vec![1.0]],
446            levels: None, filled: false, cmap: Colormap::Viridis,
447            colors: None, linewidths: 1.0, label: None,
448            color: Color::TAB_BLUE, num_levels: 10,
449        };
450        let (xmin, xmax, ymin, ymax) = c.data_bounds();
451        assert!((xmin - 5.0).abs() < f64::EPSILON);
452        assert!((xmax - 5.0).abs() < f64::EPSILON);
453        assert!((ymin - 3.0).abs() < f64::EPSILON);
454        assert!((ymax - 3.0).abs() < f64::EPSILON);
455    }
456
457    #[test]
458    fn marching_squares_no_contour() {
459        let c = ContourArtist {
460            x: vec![0.0, 1.0],
461            y: vec![0.0, 1.0],
462            z: vec![vec![0.0, 0.0], vec![0.0, 0.0]],
463            levels: None, filled: false, cmap: Colormap::Viridis,
464            colors: None, linewidths: 1.0, label: None,
465            color: Color::TAB_BLUE, num_levels: 10,
466        };
467        let segs = c.marching_squares(1.0);
468        assert!(segs.is_empty());
469    }
470
471    #[test]
472    fn marching_squares_all_above() {
473        let c = ContourArtist {
474            x: vec![0.0, 1.0],
475            y: vec![0.0, 1.0],
476            z: vec![vec![5.0, 5.0], vec![5.0, 5.0]],
477            levels: None, filled: false, cmap: Colormap::Viridis,
478            colors: None, linewidths: 1.0, label: None,
479            color: Color::TAB_BLUE, num_levels: 10,
480        };
481        let segs = c.marching_squares(1.0);
482        assert!(segs.is_empty());
483    }
484
485    #[test]
486    fn marching_squares_produces_segments() {
487        let c = ContourArtist {
488            x: vec![0.0, 1.0, 2.0],
489            y: vec![0.0, 1.0, 2.0],
490            z: vec![
491                vec![0.0, 1.0, 2.0],
492                vec![1.0, 2.0, 3.0],
493                vec![2.0, 3.0, 4.0],
494            ],
495            levels: None, filled: false, cmap: Colormap::Viridis,
496            colors: None, linewidths: 1.0, label: None,
497            color: Color::TAB_BLUE, num_levels: 10,
498        };
499        let segs = c.marching_squares(1.5);
500        assert!(!segs.is_empty(), "should produce at least one segment for level 1.5");
501    }
502
503    #[test]
504    fn marching_squares_with_nan() {
505        let c = ContourArtist {
506            x: vec![0.0, 1.0],
507            y: vec![0.0, 1.0],
508            z: vec![vec![0.0, f64::NAN], vec![1.0, 2.0]],
509            levels: None, filled: false, cmap: Colormap::Viridis,
510            colors: None, linewidths: 1.0, label: None,
511            color: Color::TAB_BLUE, num_levels: 10,
512        };
513        let segs = c.marching_squares(0.5);
514        assert!(segs.is_empty());
515    }
516
517    #[test]
518    fn marching_squares_small_grid() {
519        let c = ContourArtist {
520            x: vec![0.0],
521            y: vec![0.0],
522            z: vec![vec![1.0]],
523            levels: None, filled: false, cmap: Colormap::Viridis,
524            colors: None, linewidths: 1.0, label: None,
525            color: Color::TAB_BLUE, num_levels: 10,
526        };
527        let segs = c.marching_squares(0.5);
528        assert!(segs.is_empty());
529    }
530
531    #[test]
532    fn cell_averages_basic() {
533        let c = sample_contour();
534        let avgs = c.cell_averages();
535        assert_eq!(avgs.len(), 2);
536        assert_eq!(avgs[0].len(), 2);
537        // avg of (0, 1, 1, 2) = 1.0
538        assert!((avgs[0][0] - 1.0).abs() < f64::EPSILON);
539    }
540
541    #[test]
542    fn cell_averages_empty() {
543        let c = ContourArtist {
544            x: vec![], y: vec![], z: vec![],
545            levels: None, filled: false, cmap: Colormap::Viridis,
546            colors: None, linewidths: 1.0, label: None,
547            color: Color::TAB_BLUE, num_levels: 10,
548        };
549        let avgs = c.cell_averages();
550        assert!(avgs.is_empty());
551    }
552
553    #[test]
554    fn cell_averages_with_nan() {
555        let c = ContourArtist {
556            x: vec![0.0, 1.0],
557            y: vec![0.0, 1.0],
558            z: vec![
559                vec![f64::NAN, f64::NAN],
560                vec![f64::NAN, f64::NAN],
561            ],
562            levels: None, filled: false, cmap: Colormap::Viridis,
563            colors: None, linewidths: 1.0, label: None,
564            color: Color::TAB_BLUE, num_levels: 10,
565        };
566        let avgs = c.cell_averages();
567        assert_eq!(avgs.len(), 1);
568        assert_eq!(avgs[0].len(), 1);
569        assert!(avgs[0][0].is_nan());
570    }
571
572    #[test]
573    fn marching_squares_diagonal_contour() {
574        let c = ContourArtist {
575            x: vec![0.0, 1.0],
576            y: vec![0.0, 1.0],
577            z: vec![
578                vec![0.0, 1.0],
579                vec![1.0, 2.0],
580            ],
581            levels: None, filled: false, cmap: Colormap::Viridis,
582            colors: None, linewidths: 1.0, label: None,
583            color: Color::TAB_BLUE, num_levels: 10,
584        };
585        let segs = c.marching_squares(0.5);
586        assert_eq!(segs.len(), 1);
587        // bl=0 < 0.5, br=1 >= 0.5, tr=2 >= 0.5, tl=1 >= 0.5
588        // case = 0 | 2 | 4 | 8 = 14 -> same as case 1:
589        // bottom edge and left edge
590        let (x0, y0, x1, y1) = segs[0];
591        // bottom edge: interp(0, 1) at level 0.5 -> t=0.5 -> x=0.5, y=0
592        // left edge: interp(0, 1) at level 0.5 -> t=0.5 -> x=0, y=0.5
593        assert!((x0 - 0.5).abs() < f64::EPSILON);
594        assert!((y0 - 0.0).abs() < f64::EPSILON);
595        assert!((x1 - 0.0).abs() < f64::EPSILON);
596        assert!((y1 - 0.5).abs() < f64::EPSILON);
597    }
598}
plotkit_core/charts/contour.rs

plotkit_core/charts/
contour.rs
