plotpy/
image.rs

1use super::{generate_nested_list_3, 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    ///
59    /// # Arguments
60    ///
61    /// * `data` - 2D matrix-like data structure
62    ///
63    /// See <https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html>
64    pub fn draw<'a, T, U>(&mut self, data: &'a T)
65    where
66        T: AsMatrix<'a, U>,
67        U: 'a + std::fmt::Display + Num,
68    {
69        matrix_to_array(&mut self.buffer, "data", data);
70        let opt = self.options();
71        write!(&mut self.buffer, "plt.imshow(data{})\n", &opt).unwrap();
72    }
73
74    /// (imshow) Displays data as an image with RGB or RGB(A) values
75    ///
76    /// # Arguments
77    ///
78    /// * `data` - 3D vector with shape (height, width, 3) for RGB or (height, width, 4) for RGBA
79    ///   The inner-most vector contains the color channels.
80    ///
81    /// See <https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html>
82    pub fn draw_rgb_or_rgba<T>(&mut self, data: &Vec<Vec<Vec<T>>>)
83    where
84        T: std::fmt::Display + Num,
85    {
86        generate_nested_list_3(&mut self.buffer, "data", data);
87        let opt = self.options();
88        write!(&mut self.buffer, "plt.imshow(data{})\n", &opt).unwrap();
89    }
90
91    /// Sets the colormap index
92    ///
93    /// Options:
94    ///
95    /// * 0 -- bwr
96    /// * 1 -- RdBu
97    /// * 2 -- hsv
98    /// * 3 -- jet
99    /// * 4 -- terrain
100    /// * 5 -- pink
101    /// * 6 -- Greys
102    /// * `>`6 -- starts over from 0
103    pub fn set_colormap_index(&mut self, index: usize) -> &mut Self {
104        const CMAP: [&str; 7] = ["bwr", "RdBu", "hsv", "jet", "terrain", "pink", "Greys"];
105        self.colormap_name = CMAP[index % 7].to_string();
106        self
107    }
108
109    /// Sets the colormap name
110    ///
111    /// Colormap names:
112    ///
113    /// * see <https://matplotlib.org/stable/tutorials/colors/colormaps.html>
114    ///
115    /// Will use `colormap_index` instead if `colormap_name` is empty.
116    pub fn set_colormap_name(&mut self, name: &str) -> &mut Self {
117        self.colormap_name = String::from(name);
118        self
119    }
120
121    // Sets extra python/matplotlib commands (comma separated)
122    pub fn set_extra(&mut self, extra: &str) -> &mut Self {
123        self.extra = extra.to_string();
124        self
125    }
126
127    /// Returns options for barplot
128    fn options(&self) -> String {
129        let mut opt = String::new();
130        if self.colormap_name != "" {
131            write!(&mut opt, ",cmap=plt.get_cmap('{}')", self.colormap_name).unwrap();
132        }
133        if self.extra != "" {
134            write!(&mut opt, ",{}", self.extra).unwrap();
135        }
136        opt
137    }
138}
139
140impl GraphMaker for Image {
141    fn get_buffer<'a>(&'a self) -> &'a String {
142        &self.buffer
143    }
144    fn clear_buffer(&mut self) {
145        self.buffer.clear();
146    }
147}
148
149////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
150
151#[cfg(test)]
152mod tests {
153    use super::Image;
154    use crate::GraphMaker;
155
156    #[test]
157    fn new_works() {
158        let img = Image::new();
159        assert_eq!(img.colormap_name.len(), 0);
160        assert_eq!(img.extra.len(), 0);
161        assert_eq!(img.buffer.len(), 0);
162    }
163
164    #[test]
165    fn draw_works_1() {
166        let xx = [[1, 2], [3, 2]];
167        let mut img = Image::new();
168        img.set_colormap_index(0).set_colormap_name("terrain").draw(&xx);
169        let b: &str = "data=np.array([[1,2,],[3,2,],])\n\
170                       plt.imshow(data,cmap=plt.get_cmap('terrain'))\n";
171        assert_eq!(img.buffer, b);
172        img.clear_buffer();
173        assert_eq!(img.buffer, "");
174    }
175}
plotpy/image.rs

plotpy/
image.rs
