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![],
362            y: vec![],
363            z: vec![],
364            levels: None,
365            filled: false,
366            cmap: Colormap::Viridis,
367            colors: None,
368            linewidths: 1.0,
369            label: None,
370            color: Color::TAB_BLUE,
371            num_levels: 10,
372        };
373        let (lo, hi) = c.z_bounds();
374        assert!((lo - 0.0).abs() < f64::EPSILON);
375        assert!((hi - 1.0).abs() < f64::EPSILON);
376    }
377
378    #[test]
379    fn z_bounds_with_nan() {
380        let c = ContourArtist {
381            x: vec![0.0, 1.0],
382            y: vec![0.0, 1.0],
383            z: vec![vec![f64::NAN, 3.0], vec![1.0, f64::NAN]],
384            levels: None,
385            filled: false,
386            cmap: Colormap::Viridis,
387            colors: None,
388            linewidths: 1.0,
389            label: None,
390            color: Color::TAB_BLUE,
391            num_levels: 10,
392        };
393        let (lo, hi) = c.z_bounds();
394        assert!((lo - 1.0).abs() < f64::EPSILON);
395        assert!((hi - 3.0).abs() < f64::EPSILON);
396    }
397
398    #[test]
399    fn effective_levels_auto() {
400        let c = sample_contour();
401        let levels = c.effective_levels();
402        assert_eq!(levels.len(), 10);
403        assert!((levels[0] - 0.0).abs() < f64::EPSILON);
404        assert!((levels[9] - 4.0).abs() < f64::EPSILON);
405    }
406
407    #[test]
408    fn effective_levels_explicit() {
409        let mut c = sample_contour();
410        c.levels(vec![1.0, 2.0, 3.0]);
411        let levels = c.effective_levels();
412        assert_eq!(levels, vec![1.0, 2.0, 3.0]);
413    }
414
415    #[test]
416    fn effective_levels_constant_z() {
417        let c = ContourArtist {
418            x: vec![0.0, 1.0],
419            y: vec![0.0, 1.0],
420            z: vec![vec![5.0, 5.0], vec![5.0, 5.0]],
421            levels: None,
422            filled: false,
423            cmap: Colormap::Viridis,
424            colors: None,
425            linewidths: 1.0,
426            label: None,
427            color: Color::TAB_BLUE,
428            num_levels: 10,
429        };
430        let levels = c.effective_levels();
431        assert_eq!(levels.len(), 1);
432        assert!((levels[0] - 5.0).abs() < f64::EPSILON);
433    }
434
435    #[test]
436    fn data_bounds_basic() {
437        let c = sample_contour();
438        let (xmin, xmax, ymin, ymax) = c.data_bounds();
439        assert!((xmin - 0.0).abs() < f64::EPSILON);
440        assert!((xmax - 2.0).abs() < f64::EPSILON);
441        assert!((ymin - 0.0).abs() < f64::EPSILON);
442        assert!((ymax - 2.0).abs() < f64::EPSILON);
443    }
444
445    #[test]
446    fn data_bounds_empty() {
447        let c = ContourArtist {
448            x: vec![],
449            y: vec![],
450            z: vec![],
451            levels: None,
452            filled: false,
453            cmap: Colormap::Viridis,
454            colors: None,
455            linewidths: 1.0,
456            label: None,
457            color: Color::TAB_BLUE,
458            num_levels: 10,
459        };
460        assert_eq!(c.data_bounds(), (0.0, 1.0, 0.0, 1.0));
461    }
462
463    #[test]
464    fn data_bounds_single_point() {
465        let c = ContourArtist {
466            x: vec![5.0],
467            y: vec![3.0],
468            z: vec![vec![1.0]],
469            levels: None,
470            filled: false,
471            cmap: Colormap::Viridis,
472            colors: None,
473            linewidths: 1.0,
474            label: None,
475            color: Color::TAB_BLUE,
476            num_levels: 10,
477        };
478        let (xmin, xmax, ymin, ymax) = c.data_bounds();
479        assert!((xmin - 5.0).abs() < f64::EPSILON);
480        assert!((xmax - 5.0).abs() < f64::EPSILON);
481        assert!((ymin - 3.0).abs() < f64::EPSILON);
482        assert!((ymax - 3.0).abs() < f64::EPSILON);
483    }
484
485    #[test]
486    fn marching_squares_no_contour() {
487        let c = ContourArtist {
488            x: vec![0.0, 1.0],
489            y: vec![0.0, 1.0],
490            z: vec![vec![0.0, 0.0], vec![0.0, 0.0]],
491            levels: None,
492            filled: false,
493            cmap: Colormap::Viridis,
494            colors: None,
495            linewidths: 1.0,
496            label: None,
497            color: Color::TAB_BLUE,
498            num_levels: 10,
499        };
500        let segs = c.marching_squares(1.0);
501        assert!(segs.is_empty());
502    }
503
504    #[test]
505    fn marching_squares_all_above() {
506        let c = ContourArtist {
507            x: vec![0.0, 1.0],
508            y: vec![0.0, 1.0],
509            z: vec![vec![5.0, 5.0], vec![5.0, 5.0]],
510            levels: None,
511            filled: false,
512            cmap: Colormap::Viridis,
513            colors: None,
514            linewidths: 1.0,
515            label: None,
516            color: Color::TAB_BLUE,
517            num_levels: 10,
518        };
519        let segs = c.marching_squares(1.0);
520        assert!(segs.is_empty());
521    }
522
523    #[test]
524    fn marching_squares_produces_segments() {
525        let c = ContourArtist {
526            x: vec![0.0, 1.0, 2.0],
527            y: vec![0.0, 1.0, 2.0],
528            z: vec![
529                vec![0.0, 1.0, 2.0],
530                vec![1.0, 2.0, 3.0],
531                vec![2.0, 3.0, 4.0],
532            ],
533            levels: None,
534            filled: false,
535            cmap: Colormap::Viridis,
536            colors: None,
537            linewidths: 1.0,
538            label: None,
539            color: Color::TAB_BLUE,
540            num_levels: 10,
541        };
542        let segs = c.marching_squares(1.5);
543        assert!(
544            !segs.is_empty(),
545            "should produce at least one segment for level 1.5"
546        );
547    }
548
549    #[test]
550    fn marching_squares_with_nan() {
551        let c = ContourArtist {
552            x: vec![0.0, 1.0],
553            y: vec![0.0, 1.0],
554            z: vec![vec![0.0, f64::NAN], vec![1.0, 2.0]],
555            levels: None,
556            filled: false,
557            cmap: Colormap::Viridis,
558            colors: None,
559            linewidths: 1.0,
560            label: None,
561            color: Color::TAB_BLUE,
562            num_levels: 10,
563        };
564        let segs = c.marching_squares(0.5);
565        assert!(segs.is_empty());
566    }
567
568    #[test]
569    fn marching_squares_small_grid() {
570        let c = ContourArtist {
571            x: vec![0.0],
572            y: vec![0.0],
573            z: vec![vec![1.0]],
574            levels: None,
575            filled: false,
576            cmap: Colormap::Viridis,
577            colors: None,
578            linewidths: 1.0,
579            label: None,
580            color: Color::TAB_BLUE,
581            num_levels: 10,
582        };
583        let segs = c.marching_squares(0.5);
584        assert!(segs.is_empty());
585    }
586
587    #[test]
588    fn cell_averages_basic() {
589        let c = sample_contour();
590        let avgs = c.cell_averages();
591        assert_eq!(avgs.len(), 2);
592        assert_eq!(avgs[0].len(), 2);
593        // avg of (0, 1, 1, 2) = 1.0
594        assert!((avgs[0][0] - 1.0).abs() < f64::EPSILON);
595    }
596
597    #[test]
598    fn cell_averages_empty() {
599        let c = ContourArtist {
600            x: vec![],
601            y: vec![],
602            z: vec![],
603            levels: None,
604            filled: false,
605            cmap: Colormap::Viridis,
606            colors: None,
607            linewidths: 1.0,
608            label: None,
609            color: Color::TAB_BLUE,
610            num_levels: 10,
611        };
612        let avgs = c.cell_averages();
613        assert!(avgs.is_empty());
614    }
615
616    #[test]
617    fn cell_averages_with_nan() {
618        let c = ContourArtist {
619            x: vec![0.0, 1.0],
620            y: vec![0.0, 1.0],
621            z: vec![vec![f64::NAN, f64::NAN], vec![f64::NAN, f64::NAN]],
622            levels: None,
623            filled: false,
624            cmap: Colormap::Viridis,
625            colors: None,
626            linewidths: 1.0,
627            label: None,
628            color: Color::TAB_BLUE,
629            num_levels: 10,
630        };
631        let avgs = c.cell_averages();
632        assert_eq!(avgs.len(), 1);
633        assert_eq!(avgs[0].len(), 1);
634        assert!(avgs[0][0].is_nan());
635    }
636
637    #[test]
638    fn marching_squares_diagonal_contour() {
639        let c = ContourArtist {
640            x: vec![0.0, 1.0],
641            y: vec![0.0, 1.0],
642            z: vec![vec![0.0, 1.0], vec![1.0, 2.0]],
643            levels: None,
644            filled: false,
645            cmap: Colormap::Viridis,
646            colors: None,
647            linewidths: 1.0,
648            label: None,
649            color: Color::TAB_BLUE,
650            num_levels: 10,
651        };
652        let segs = c.marching_squares(0.5);
653        assert_eq!(segs.len(), 1);
654        // bl=0 < 0.5, br=1 >= 0.5, tr=2 >= 0.5, tl=1 >= 0.5
655        // case = 0 | 2 | 4 | 8 = 14 -> same as case 1:
656        // bottom edge and left edge
657        let (x0, y0, x1, y1) = segs[0];
658        // bottom edge: interp(0, 1) at level 0.5 -> t=0.5 -> x=0.5, y=0
659        // left edge: interp(0, 1) at level 0.5 -> t=0.5 -> x=0, y=0.5
660        assert!((x0 - 0.5).abs() < f64::EPSILON);
661        assert!((y0 - 0.0).abs() < f64::EPSILON);
662        assert!((x1 - 0.0).abs() < f64::EPSILON);
663        assert!((y1 - 0.5).abs() < f64::EPSILON);
664    }
665}
plotkit_core/charts/contour.rs

plotkit_core/charts/
contour.rs
