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/// 
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}