plotpy/
contour.rs

1use super::{generate_list_quoted, matrix_to_array, vector_to_array, AsMatrix, GraphMaker};
2use num_traits::Num;
3use std::fmt::Write;
4
5/// Generates a contour plot
6///
7/// [See Matplotlib's documentation](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.contour.html)
8///
9/// # Example
10///
11/// ```
12/// use plotpy::{Contour, Plot, generate3d};
13///
14/// fn main() -> Result<(), &'static str> {
15///     // generate (x,y,z) matrices
16///     let n = 21;
17///     let (x, y, z) = generate3d(-2.0, 2.0, -2.0, 2.0, n, n, |x, y| x * x - y * y);
18///
19///     // configure contour
20///     let mut contour = Contour::new();
21///     contour
22///         .set_colorbar_label("temperature")
23///         .set_colormap_name("terrain")
24///         .set_selected_line_color("#f1eb67")
25///         .set_selected_line_width(12.0)
26///         .set_selected_level(0.0, true);
27///
28///     // draw contour
29///     contour.draw(&x, &y, &z);
30///
31///     // add contour to plot
32///     let mut plot = Plot::new();
33///     plot.add(&contour)
34///         .set_labels("x", "y");
35///
36///     // save figure
37///     plot.save("/tmp/plotpy/doc_tests/doc_contour.svg")?;
38///     Ok(())
39/// }
40/// ```
41///
42/// ![doc_contour.svg](https://raw.githubusercontent.com/cpmech/plotpy/main/figures/doc_contour.svg)
43///
44/// See also integration tests in the [tests directory](https://github.com/cpmech/plotpy/tree/main/tests)
45///
46/// Output from some integration tests:
47///
48/// ![integ_contour.svg](https://raw.githubusercontent.com/cpmech/plotpy/main/figures/integ_contour.svg)
49pub struct Contour {
50    colors: Vec<String>,         // Colors to be used instead of colormap
51    levels: Vec<f64>,            // Pre-defined levels
52    colormap_name: String,       // Colormap name
53    no_lines: bool,              // Skip drawing a lines contour
54    no_labels: bool,             // Skip adding labels to the lines contour
55    no_inline_labels: bool,      // Do not draw labels inline
56    no_colorbar: bool,           // Skip drawing a colorbar
57    colorbar_label: String,      // Colorbar label
58    number_format_cb: String,    // Number format for the labels in lines contour
59    line_color: String,          // Line color for the lines contour
60    line_style: String,          // Line style for the lines contour
61    line_width: f64,             // Line width for the lines contour
62    fontsize_labels: f64,        // Font size for labels
63    with_selected: bool,         // Draw a line contour with a selected level
64    selected_level: f64,         // Selected level (e.g., 0.0)
65    selected_line_color: String, // Color to mark the selected level
66    selected_line_style: String, // Line style for the selected level
67    selected_line_width: f64,    // Line width for the selected level
68    extra_filled: String,        // Extra commands (comma separated) for the filled contour
69    extra_line: String,          // Extra commands (comma separated) for the line contour
70    buffer: String,              // buffer
71}
72
73impl Contour {
74    /// Creates a new Contour object
75    pub fn new() -> Self {
76        Contour {
77            colors: Vec::new(),
78            levels: Vec::new(),
79            colormap_name: "bwr".to_string(),
80            no_lines: false,
81            no_labels: false,
82            no_inline_labels: false,
83            no_colorbar: false,
84            colorbar_label: String::new(),
85            number_format_cb: String::new(),
86            line_color: "black".to_string(),
87            line_style: String::new(),
88            line_width: 0.0,
89            fontsize_labels: 0.0,
90            with_selected: false,
91            selected_level: 0.0,
92            selected_line_color: "yellow".to_string(),
93            selected_line_style: "-".to_string(),
94            selected_line_width: 2.0,
95            extra_filled: String::new(),
96            extra_line: String::new(),
97            buffer: String::new(),
98        }
99    }
100
101    /// Draws a fancy contour: filled contour with a line contour and a colorbar
102    ///
103    /// # Input
104    ///
105    /// * `x` -- matrix with x values
106    /// * `y` -- matrix with y values
107    /// * `z` -- matrix with z values
108    ///
109    /// # Flags
110    ///
111    /// The following flags control what features are not to be drawn:
112    ///
113    /// * `no_lines` -- skip drawing a lines contour on top of the filled contour
114    /// * `no_labels` -- skip adding labels to the lines contour (if enabled)
115    /// * `no_colorbar` -- skip drawing a colorbar
116    /// * `with_selected` -- draw a line contour with a selected level (e.g., 0.0) on top of everything
117    pub fn draw<'a, T, U>(&mut self, x: &'a T, y: &'a T, z: &'a T)
118    where
119        T: AsMatrix<'a, U>,
120        U: 'a + std::fmt::Display + Num,
121    {
122        matrix_to_array(&mut self.buffer, "x", x);
123        matrix_to_array(&mut self.buffer, "y", y);
124        matrix_to_array(&mut self.buffer, "z", z);
125        if self.colors.len() > 0 {
126            generate_list_quoted(&mut self.buffer, "colors", &self.colors);
127        }
128        if self.levels.len() > 0 {
129            vector_to_array(&mut self.buffer, "levels", &self.levels);
130        }
131        let opt = self.options_filled();
132        write!(&mut self.buffer, "cf=plt.contourf(x,y,z{})\n", &opt).unwrap();
133        if !self.no_lines {
134            let opt_line = self.options_line();
135            write!(&mut self.buffer, "cl=plt.contour(x,y,z{})\n", &opt_line).unwrap();
136            if !self.no_labels {
137                let opt_label = self.options_label();
138                write!(&mut self.buffer, "plt.clabel(cl{})\n", &opt_label).unwrap();
139            }
140        }
141        if !self.no_colorbar {
142            let opt_colorbar = self.options_colorbar();
143            write!(&mut self.buffer, "cb=plt.colorbar(cf{})\n", &opt_colorbar).unwrap();
144            if self.colorbar_label != "" {
145                write!(&mut self.buffer, "cb.ax.set_ylabel(r'{}')\n", self.colorbar_label).unwrap();
146            }
147        }
148        if self.with_selected {
149            let opt_selected = self.options_selected();
150            write!(&mut self.buffer, "plt.contour(x,y,z{})\n", &opt_selected).unwrap();
151        }
152    }
153
154    /// Sets the colors to be used instead of a pre-defined colormap
155    ///
156    /// Will use `colormap_index` instead if its empty.
157    pub fn set_colors(&mut self, colors: &[&str]) -> &mut Self {
158        self.colors = colors.iter().map(|color| color.to_string()).collect();
159        self
160    }
161
162    /// Sets pre-defined levels, otherwise automatically calculate levels
163    pub fn set_levels(&mut self, levels: &[f64]) -> &mut Self {
164        self.levels = levels.to_vec();
165        self
166    }
167
168    /// Sets the colormap index
169    ///
170    /// Options:
171    ///
172    /// * 0 -- bwr
173    /// * 1 -- RdBu
174    /// * 2 -- hsv
175    /// * 3 -- jet
176    /// * 4 -- terrain
177    /// * 5 -- pink
178    /// * 6 -- Greys
179    /// * `>`6 -- starts over from 0
180    pub fn set_colormap_index(&mut self, index: usize) -> &mut Self {
181        const CMAP: [&str; 7] = ["bwr", "RdBu", "hsv", "jet", "terrain", "pink", "Greys"];
182        self.colormap_name = CMAP[index % 7].to_string();
183        self.colors = Vec::new();
184        self
185    }
186
187    /// Sets the colormap name
188    ///
189    /// Colormap names:
190    ///
191    /// * see <https://matplotlib.org/stable/tutorials/colors/colormaps.html>
192    ///
193    /// Will use `colormap_index` instead if `colormap_name` is empty.
194    pub fn set_colormap_name(&mut self, name: &str) -> &mut Self {
195        self.colormap_name = String::from(name);
196        self.colors = Vec::new();
197        self
198    }
199
200    /// Sets option to skip drawing a lines contour on top of the filled contour
201    pub fn set_no_lines(&mut self, flag: bool) -> &mut Self {
202        self.no_lines = flag;
203        self
204    }
205
206    /// Sets option to skip adding labels to the lines contour (if enabled)
207    pub fn set_no_labels(&mut self, flag: bool) -> &mut Self {
208        self.no_labels = flag;
209        self
210    }
211
212    /// Sets option to skip drawing labels inline with the contour lines (if enabled)
213    pub fn set_no_inline_labels(&mut self, flag: bool) -> &mut Self {
214        self.no_inline_labels = flag;
215        self
216    }
217
218    /// Sets option to skip drawing a colorbar
219    pub fn set_no_colorbar(&mut self, flag: bool) -> &mut Self {
220        self.no_colorbar = flag;
221        self
222    }
223
224    /// Sets the colorbar label
225    pub fn set_colorbar_label(&mut self, label: &str) -> &mut Self {
226        self.colorbar_label = String::from(label);
227        self
228    }
229
230    /// Sets the number format for the labels in the colorbar (cb)
231    pub fn set_number_format_cb(&mut self, format: &str) -> &mut Self {
232        self.number_format_cb = String::from(format);
233        self
234    }
235
236    /// Sets the line color for the lines contour (default is black)
237    pub fn set_line_color(&mut self, color: &str) -> &mut Self {
238        self.line_color = String::from(color);
239        self
240    }
241
242    /// Sets the line style for the lines contour
243    ///
244    /// Options:
245    ///
246    /// * "`-`", "`:`", "`--`", "`-.`"
247    pub fn set_line_style(&mut self, style: &str) -> &mut Self {
248        self.line_style = String::from(style);
249        self
250    }
251
252    /// Sets the line width for the lines contour
253    pub fn set_line_width(&mut self, width: f64) -> &mut Self {
254        self.line_width = width;
255        self
256    }
257
258    /// Sets the font size for labels
259    pub fn set_fontsize_labels(&mut self, fontsize: f64) -> &mut Self {
260        self.fontsize_labels = fontsize;
261        self
262    }
263
264    /// Sets option to draw a line contour with a selected level (e.g., 0.0)
265    ///
266    /// Will draw the selected level (e.g., 0.0) on top of everything
267    pub fn set_selected_level(&mut self, level: f64, enabled: bool) -> &mut Self {
268        self.selected_level = level;
269        self.with_selected = enabled;
270        self
271    }
272
273    /// Sets the color to mark the selected level
274    pub fn set_selected_line_color(&mut self, color: &str) -> &mut Self {
275        self.selected_line_color = String::from(color);
276        self
277    }
278
279    /// Sets the line style for the selected level
280    ///
281    /// Options:
282    ///
283    /// * "`-`", "`:`", "`--`", "`-.`"
284    pub fn set_selected_line_style(&mut self, style: &str) -> &mut Self {
285        self.selected_line_style = String::from(style);
286        self
287    }
288
289    /// Sets the line width for the selected level
290    pub fn set_selected_line_width(&mut self, width: f64) -> &mut Self {
291        self.selected_line_width = width;
292        self
293    }
294
295    /// Sets extra matplotlib commands (comma separated) for the filled contour
296    ///
297    /// **Important:** The extra commands must be comma separated. For example:
298    ///
299    /// ```text
300    /// param1=123,param2='hello'
301    /// ```
302    ///
303    /// [See Matplotlib's documentation for extra parameters](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.contour.html)
304    pub fn set_extra_filled(&mut self, extra: &str) -> &mut Self {
305        self.extra_filled = extra.to_string();
306        self
307    }
308
309    /// Sets extra matplotlib commands (comma separated) for the line contour
310    ///
311    /// **Important:** The extra commands must be comma separated. For example:
312    ///
313    /// ```text
314    /// param1=123,param2='hello'
315    /// ```
316    ///
317    /// [See Matplotlib's documentation for extra parameters](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.contour.html)
318    pub fn set_extra_line(&mut self, extra: &str) -> &mut Self {
319        self.extra_line = extra.to_string();
320        self
321    }
322
323    /// Returns options for filled contour
324    fn options_filled(&self) -> String {
325        let mut opt = String::new();
326        if self.colors.len() > 0 {
327            write!(&mut opt, ",colors=colors",).unwrap();
328        } else {
329            if self.colormap_name != "" {
330                write!(&mut opt, ",cmap=plt.get_cmap('{}')", self.colormap_name).unwrap();
331            }
332        }
333        if self.levels.len() > 0 {
334            write!(&mut opt, ",levels=levels").unwrap();
335        }
336        if self.extra_filled != "" {
337            write!(&mut opt, ",{}", self.extra_filled).unwrap();
338        }
339        opt
340    }
341
342    /// Returns options for line contour
343    fn options_line(&self) -> String {
344        let mut opt = String::new();
345        if self.line_color != "" {
346            write!(&mut opt, ",colors=['{}']", self.line_color).unwrap();
347        }
348        if self.levels.len() > 0 {
349            write!(&mut opt, ",levels=levels").unwrap();
350        }
351        if self.line_style != "" {
352            write!(&mut opt, ",linestyles=['{}']", self.line_style).unwrap();
353        }
354        if self.line_width > 0.0 {
355            write!(&mut opt, ",linewidths=[{}]", self.line_width).unwrap();
356        }
357        if self.extra_line != "" {
358            write!(&mut opt, ",{}", self.extra_line).unwrap();
359        }
360        opt
361    }
362
363    /// Returns options for labels
364    fn options_label(&self) -> String {
365        let mut opt = String::new();
366        if self.no_inline_labels {
367            write!(&mut opt, ",inline=False").unwrap();
368        } else {
369            write!(&mut opt, ",inline=True").unwrap();
370        }
371        if self.fontsize_labels > 0.0 {
372            write!(&mut opt, ",fontsize={}", self.fontsize_labels).unwrap();
373        }
374        opt
375    }
376
377    /// Returns options for colorbar
378    fn options_colorbar(&self) -> String {
379        let mut opt = String::new();
380        if self.number_format_cb != "" {
381            write!(&mut opt, ",format='{}'", self.number_format_cb).unwrap();
382        }
383        opt
384    }
385
386    /// Returns options for selected line contour
387    fn options_selected(&self) -> String {
388        let mut opt = String::new();
389        if self.selected_line_color != "" {
390            write!(&mut opt, ",colors=['{}']", self.selected_line_color).unwrap();
391        }
392        write!(&mut opt, ",levels=[{}]", self.selected_level).unwrap();
393        if self.selected_line_style != "" {
394            write!(&mut opt, ",linestyles=['{}']", self.selected_line_style).unwrap();
395        }
396        if self.selected_line_width > 0.0 {
397            write!(&mut opt, ",linewidths=[{}]", self.selected_line_width).unwrap();
398        }
399        opt
400    }
401}
402
403impl GraphMaker for Contour {
404    fn get_buffer<'a>(&'a self) -> &'a String {
405        &self.buffer
406    }
407    fn clear_buffer(&mut self) {
408        self.buffer.clear();
409    }
410}
411
412////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
413
414#[cfg(test)]
415mod tests {
416    use super::Contour;
417    use crate::GraphMaker;
418
419    #[test]
420    fn new_works() {
421        let contour = Contour::new();
422        assert_eq!(contour.colors.len(), 0);
423        assert_eq!(contour.levels.len(), 0);
424        assert_eq!(contour.colormap_name, "bwr");
425        assert_eq!(contour.no_lines, false);
426        assert_eq!(contour.no_labels, false);
427        assert_eq!(contour.no_inline_labels, false);
428        assert_eq!(contour.no_colorbar, false);
429        assert_eq!(contour.colorbar_label.len(), 0);
430        assert_eq!(contour.number_format_cb.len(), 0);
431        assert_eq!(contour.line_color, "black".to_string());
432        assert_eq!(contour.line_style.len(), 0);
433        assert_eq!(contour.line_width, 0.0);
434        assert_eq!(contour.fontsize_labels, 0.0);
435        assert_eq!(contour.with_selected, false);
436        assert_eq!(contour.selected_level, 0.0);
437        assert_eq!(contour.selected_line_color, "yellow".to_string());
438        assert_eq!(contour.selected_line_style, "-".to_string());
439        assert_eq!(contour.selected_line_width, 2.0);
440        assert_eq!(contour.buffer.len(), 0);
441    }
442
443    #[test]
444    fn options_filled_works() {
445        let mut contour = Contour::new();
446        contour
447            .set_colors(&vec!["#f00", "#0f0", "#00f"])
448            .set_levels(&vec![0.25, 0.5, 1.0]);
449        let opt = contour.options_filled();
450        assert_eq!(
451            opt,
452            ",colors=colors\
453             ,levels=levels"
454        );
455        contour.set_colormap_index(4);
456        let opt = contour.options_filled();
457        assert_eq!(
458            opt,
459            ",cmap=plt.get_cmap('terrain')\
460             ,levels=levels"
461        );
462    }
463
464    #[test]
465    fn options_line_works() {
466        let mut contour = Contour::new();
467        contour
468            .set_levels(&vec![0.25, 0.5, 1.0])
469            .set_line_color("red")
470            .set_line_style(":")
471            .set_line_width(3.0);
472        let opt = contour.options_line();
473        assert_eq!(
474            opt,
475            ",colors=['red']\
476             ,levels=levels\
477             ,linestyles=[':']\
478             ,linewidths=[3]"
479        );
480    }
481
482    #[test]
483    fn options_label_works() {
484        let mut contour = Contour::new();
485        contour.set_no_inline_labels(false).set_fontsize_labels(5.0);
486        let opt = contour.options_label();
487        assert_eq!(
488            opt,
489            ",inline=True\
490             ,fontsize=5"
491        );
492        contour.set_no_inline_labels(true);
493        let opt = contour.options_label();
494        assert_eq!(
495            opt,
496            ",inline=False\
497             ,fontsize=5"
498        );
499    }
500
501    #[test]
502    fn options_colorbar_works() {
503        let mut contour = Contour::new();
504        contour.set_number_format_cb("%.4f");
505        let opt = contour.options_colorbar();
506        assert_eq!(opt, ",format='%.4f'");
507    }
508
509    #[test]
510    fn options_selected_works() {
511        let mut contour = Contour::new();
512        contour
513            .set_selected_level(0.75, true)
514            .set_selected_line_color("blue")
515            .set_selected_line_style("--")
516            .set_selected_line_width(2.5);
517        let opt = contour.options_selected();
518        assert_eq!(
519            opt,
520            ",colors=['blue']\
521             ,levels=[0.75]\
522             ,linestyles=['--']\
523             ,linewidths=[2.5]"
524        );
525    }
526
527    #[test]
528    fn draw_works() {
529        let mut contour = Contour::new();
530        contour
531            .set_colors(&vec!["#f00", "#0f0", "#00f"])
532            .set_levels(&vec![0.25, 0.5, 1.0])
533            .set_colorbar_label("temperature")
534            .set_selected_level(0.0, true);
535        let x = vec![vec![-0.5, 0.0, 0.5], vec![-0.5, 0.0, 0.5], vec![-0.5, 0.0, 0.5]];
536        let y = vec![vec![-0.5, -0.5, -0.5], vec![0.0, 0.0, 0.0], vec![0.5, 0.5, 0.5]];
537        let z = vec![vec![0.50, 0.25, 0.50], vec![0.25, 0.00, 0.25], vec![0.50, 0.25, 0.50]];
538        contour.draw(&x, &y, &z);
539        let b: &str = "x=np.array([[-0.5,0,0.5,],[-0.5,0,0.5,],[-0.5,0,0.5,],],dtype=float)\n\
540                       y=np.array([[-0.5,-0.5,-0.5,],[0,0,0,],[0.5,0.5,0.5,],],dtype=float)\n\
541                       z=np.array([[0.5,0.25,0.5,],[0.25,0,0.25,],[0.5,0.25,0.5,],],dtype=float)\n\
542                       colors=['#f00','#0f0','#00f',]\n\
543                       levels=np.array([0.25,0.5,1,],dtype=float)\n\
544                       cf=plt.contourf(x,y,z,colors=colors,levels=levels)\n\
545                       cl=plt.contour(x,y,z,colors=['black'],levels=levels)\n\
546                       plt.clabel(cl,inline=True)\n\
547                       cb=plt.colorbar(cf)\n\
548                       cb.ax.set_ylabel(r'temperature')\n\
549                       plt.contour(x,y,z,colors=['yellow'],levels=[0],linestyles=['-'],linewidths=[2])\n";
550        assert_eq!(contour.buffer, b);
551        contour.clear_buffer();
552        assert_eq!(contour.buffer, "");
553    }
554}
plotpy/contour.rs

plotpy/
contour.rs
