plotpy/
image.rs

1use super::{matrix_to_array, AsMatrix, GraphMaker};
2use num_traits::Num;
3use std::fmt::Write;
4
5/// Generates an image plot (imshow)
6///
7/// [See Matplotlib's documentation](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html)
8///
9/// # Examples
10///
11/// ```
12/// use plotpy::{Image, Plot, StrError};
13///
14/// fn main() -> Result<(), StrError> {
15///     // set values
16///     let data = [
17///         [0.8, 2.4, 2.5, 3.9, 0.0, 4.0, 0.0],
18///         [2.4, 0.0, 4.0, 1.0, 2.7, 0.0, 0.0],
19///         [1.1, 2.4, 0.8, 4.3, 1.9, 4.4, 0.0],
20///         [0.6, 0.0, 0.3, 0.0, 3.1, 0.0, 0.0],
21///         [0.7, 1.7, 0.6, 2.6, 2.2, 6.2, 0.0],
22///         [1.3, 1.2, 0.0, 0.0, 0.0, 3.2, 5.1],
23///         [0.1, 2.0, 0.0, 1.4, 0.0, 1.9, 6.3],
24///     ];
25///
26///     // image plot and options
27///     let mut img = Image::new();
28///     img.set_colormap_name("hsv").draw(&data);
29///
30///     // save figure
31///     let mut plot = Plot::new();
32///     plot.add(&img);
33///     plot.save("/tmp/plotpy/doc_tests/doc_image_1.svg")?;
34///     Ok(())
35/// }
36/// ```
37///
38/// ![doc_image_1.svg](https://raw.githubusercontent.com/cpmech/plotpy/main/figures/doc_image_1.svg)
39///
40/// See also integration test in the **tests** directory.
41pub struct Image {
42    colormap_name: String, // Colormap name
43    extra: String,         // Extra commands (comma separated)
44    buffer: String,        // buffer
45}
46
47impl Image {
48    /// Creates a new Image object
49    pub fn new() -> Self {
50        Image {
51            colormap_name: String::new(),
52            extra: String::new(),
53            buffer: String::new(),
54        }
55    }
56
57    /// (imshow) Displays data as an image
58    pub fn draw<'a, T, U>(&mut self, data: &'a T)
59    where
60        T: AsMatrix<'a, U>,
61        U: 'a + std::fmt::Display + Num,
62    {
63        matrix_to_array(&mut self.buffer, "data", data);
64        let opt = self.options();
65        write!(&mut self.buffer, "plt.imshow(data{})\n", &opt).unwrap();
66    }
67
68    /// Sets the colormap index
69    ///
70    /// Options:
71    ///
72    /// * 0 -- bwr
73    /// * 1 -- RdBu
74    /// * 2 -- hsv
75    /// * 3 -- jet
76    /// * 4 -- terrain
77    /// * 5 -- pink
78    /// * 6 -- Greys
79    /// * `>`6 -- starts over from 0
80    pub fn set_colormap_index(&mut self, index: usize) -> &mut Self {
81        const CMAP: [&str; 7] = ["bwr", "RdBu", "hsv", "jet", "terrain", "pink", "Greys"];
82        self.colormap_name = CMAP[index % 7].to_string();
83        self
84    }
85
86    /// Sets the colormap name
87    ///
88    /// Colormap names:
89    ///
90    /// * see <https://matplotlib.org/stable/tutorials/colors/colormaps.html>
91    ///
92    /// Will use `colormap_index` instead if `colormap_name` is empty.
93    pub fn set_colormap_name(&mut self, name: &str) -> &mut Self {
94        self.colormap_name = String::from(name);
95        self
96    }
97
98    // Sets extra python/matplotlib commands (comma separated)
99    pub fn set_extra(&mut self, extra: &str) -> &mut Self {
100        self.extra = extra.to_string();
101        self
102    }
103
104    /// Returns options for barplot
105    fn options(&self) -> String {
106        let mut opt = String::new();
107        if self.colormap_name != "" {
108            write!(&mut opt, ",cmap=plt.get_cmap('{}')", self.colormap_name).unwrap();
109        }
110        if self.extra != "" {
111            write!(&mut opt, ",{}", self.extra).unwrap();
112        }
113        opt
114    }
115}
116
117impl GraphMaker for Image {
118    fn get_buffer<'a>(&'a self) -> &'a String {
119        &self.buffer
120    }
121    fn clear_buffer(&mut self) {
122        self.buffer.clear();
123    }
124}
125
126////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
127
128#[cfg(test)]
129mod tests {
130    use super::Image;
131    use crate::GraphMaker;
132
133    #[test]
134    fn new_works() {
135        let img = Image::new();
136        assert_eq!(img.colormap_name.len(), 0);
137        assert_eq!(img.extra.len(), 0);
138        assert_eq!(img.buffer.len(), 0);
139    }
140
141    #[test]
142    fn draw_works_1() {
143        let xx = [[1, 2], [3, 2]];
144        let mut img = Image::new();
145        img.set_colormap_index(0).set_colormap_name("terrain").draw(&xx);
146        let b: &str = "data=np.array([[1,2,],[3,2,],],dtype=float)\n\
147                       plt.imshow(data,cmap=plt.get_cmap('terrain'))\n";
148        assert_eq!(img.buffer, b);
149        img.clear_buffer();
150        assert_eq!(img.buffer, "");
151    }
152}