Thu, 08 Dec 2022 14:10:07 +0200
Save more CSV files when iteration-wise plotting is enabled.
This helps to generate TikZ illustrations for presentations.
0 | 1 | //! Plotting helper utilities |
2 | ||
3 | use numeric_literals::replace_float_literals; | |
4 | use std::io::Write; | |
5 | use image::{ | |
6 | ImageFormat, | |
7 | ImageBuffer, | |
8 | Rgb | |
9 | }; | |
10 | use itertools::izip; | |
11 | use colorbrewer::Palette as CbPalette; | |
12 | ||
13 | use alg_tools::types::*; | |
14 | use alg_tools::lingrid::LinGrid; | |
15 | use alg_tools::mapping::RealMapping; | |
16 | use alg_tools::loc::Loc; | |
17 | use alg_tools::bisection_tree::Bounds; | |
18 | use alg_tools::maputil::map4; | |
21
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
19 | use alg_tools::tabledump::{write_csv, TableDump}; |
0 | 20 | use crate::measures::*; |
21 | ||
22 | /// Default RGB ramp from [`colorbrewer`]. | |
23 | /// | |
24 | /// This is a tuple of parameters to [`colorbrewer::get_color_ramp`]. | |
25 | const RAMP : (CbPalette, u32) = (CbPalette::RdBu, 11); | |
26 | ||
27 | /// Helper trait for implementing dimension-dependent plotting routines. | |
28 | pub trait Plotting<const N : usize> { | |
29 | /// Plot several mappings and a discrete measure into a file. | |
30 | fn plot_into_file_spikes< | |
31 | F : Float, | |
32 | T1 : RealMapping<F, N>, | |
33 | T2 : RealMapping<F, N> | |
34 | > ( | |
35 | g_explanation : String, | |
36 | g : &T1, | |
37 | ω_explanation : String, | |
38 | ω : Option<&T2>, | |
39 | grid : LinGrid<F, N>, | |
40 | bnd : Option<Bounds<F>>, | |
41 | μ : &DiscreteMeasure<Loc<F, N>, F>, | |
42 | filename : String, | |
43 | ); | |
44 | ||
45 | /// Plot a mapping into a file, sampling values on a given grid. | |
46 | fn plot_into_file< | |
47 | F : Float, | |
48 | T1 : RealMapping<F, N>, | |
49 | > ( | |
50 | g : &T1, | |
51 | grid : LinGrid<F, N>, | |
52 | filename : String, | |
53 | explanation : String | |
54 | ); | |
55 | } | |
56 | ||
57 | /// Helper type for looking up a [`Plotting`] based on dimension. | |
58 | pub struct PlotLookup; | |
59 | ||
60 | impl Plotting<1> for PlotLookup { | |
61 | fn plot_into_file_spikes< | |
62 | F : Float, | |
63 | T1 : RealMapping<F, 1>, | |
64 | T2 : RealMapping<F, 1> | |
65 | > ( | |
66 | g_explanation : String, | |
67 | g : &T1, | |
68 | ω_explanation : String, | |
69 | ω0 : Option<&T2>, | |
70 | grid : LinGrid<F, 1>, | |
71 | bnd0 : Option<Bounds<F>>, | |
72 | μ : &DiscreteMeasure<Loc<F, 1>, F>, | |
73 | filename : String, | |
74 | ) { | |
75 | let start = grid.start[0].as_(); | |
76 | let end = grid.end[0].as_(); | |
77 | let m = μ.iter_masses().fold(F::ZERO, |m, α| m.max(α)); | |
78 | let s = μ.iter_masses().fold(F::ZERO, |m, α| m.add(α)); | |
79 | let mut spike_scale = F::ONE; | |
80 | ||
21
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
81 | μ.write_csv(format!("{filename}_mu.txt").into()).expect("CSV save error"); |
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
82 | |
0 | 83 | let mut plotter = poloto::plot( |
84 | "f", "x", | |
85 | format!("f(x); spike max={:.4}, n={}, ∑={:.4}", m, μ.len(), s) | |
86 | ).move_into(); | |
87 | ||
88 | if let Some(ω) = ω0 { | |
89 | let graph_ω = grid.into_iter().map(|x@Loc([x0]) : Loc<F, 1>| { | |
90 | [x0.as_(), ω.apply(&x).as_()] | |
91 | }); | |
92 | plotter.line(ω_explanation.as_str(), graph_ω.clone()); | |
21
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
93 | write_csv(graph_ω, format!("{filename}_omega.txt")).expect("CSV save error"); |
0 | 94 | } |
95 | ||
96 | let graph_g = grid.into_iter().map(|x@Loc([x0]) : Loc<F, 1>| { | |
97 | [x0.as_(), g.apply(&x).as_()] | |
98 | }); | |
99 | plotter.line(g_explanation.as_str(), graph_g.clone()); | |
21
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
100 | write_csv(graph_g, format!("{filename}_g.txt")).expect("CSV save error"); |
0 | 101 | |
102 | bnd0.map(|bnd| { | |
103 | let upperb = bnd.upper().as_(); | |
104 | let lowerb = bnd.lower().as_(); | |
105 | let upper : [[f64; 2]; 2] = [[start, upperb], [end, upperb]]; | |
106 | let lower = [[start, lowerb], [end, lowerb]]; | |
107 | spike_scale *= bnd.upper(); | |
108 | ||
21
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
109 | write_csv([[lowerb, upperb]].into_iter(), |
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
110 | format!("{filename}_bounds.txt")).expect("CSV save error"); |
0771706f472f
Save more CSV files when iteration-wise plotting is enabled.
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
111 | |
0 | 112 | plotter.line("upper bound", upper) |
113 | .line("lower bound", lower) | |
114 | .ymarker(lowerb) | |
115 | .ymarker(upperb); | |
116 | }); | |
117 | ||
118 | for &DeltaMeasure{ α, x : Loc([x]) } in μ.iter_spikes() { | |
119 | let spike = [[x.as_(), 0.0], [x.as_(), (α/m * spike_scale).as_()]]; | |
120 | plotter.line("", spike); | |
121 | } | |
122 | ||
123 | let svg = format!("{}", poloto::disp(|a| poloto::simple_theme(a, plotter))); | |
124 | ||
125 | std::fs::File::create(filename + ".svg").and_then(|mut file| | |
126 | file.write_all(svg.as_bytes()) | |
127 | ).expect("SVG save error"); | |
128 | } | |
129 | ||
130 | fn plot_into_file< | |
131 | F : Float, | |
132 | T1 : RealMapping<F, 1>, | |
133 | > ( | |
134 | g : &T1, | |
135 | grid : LinGrid<F, 1>, | |
136 | filename : String, | |
137 | explanation : String | |
138 | ) { | |
139 | let graph_g = grid.into_iter().map(|x@Loc([x0]) : Loc<F, 1>| { | |
140 | [x0.as_(), g.apply(&x).as_()] | |
141 | }); | |
142 | ||
143 | let plotter: poloto::Plotter<'_, float, float> = poloto::plot("f", "x", "f(x)") | |
144 | .line(explanation.as_str(), graph_g.clone()) | |
145 | .move_into(); | |
146 | ||
147 | let svg = format!("{}", poloto::disp(|a| poloto::simple_theme(a, plotter))); | |
148 | ||
149 | let svg_f = format!("{}.svg", filename); | |
150 | std::fs::File::create(svg_f).and_then(|mut file| | |
151 | file.write_all(svg.as_bytes()) | |
152 | ).expect("SVG save error"); | |
153 | ||
154 | let csv_f = format!("{}.txt", filename); | |
155 | write_csv(graph_g, csv_f).expect("CSV save error"); | |
156 | } | |
157 | ||
158 | } | |
159 | ||
160 | /// Convert $[0, 1] ∈ F$ to $\\\{0, …, M\\\} ∈ F$ where $M=$`F::RANGE_MAX`. | |
161 | #[inline] | |
162 | fn scale_uint<F, U>(v : F) -> U | |
163 | where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, | |
164 | U : Unsigned { | |
165 | (v*F::cast_from(U::RANGE_MAX)).as_() | |
166 | } | |
167 | ||
168 | /// Convert $[a, b] ∈ F$ to $\\\{0, …, M\\\} ∈ F$ where $M=$`F::RANGE_MAX`. | |
169 | #[replace_float_literals(F::cast_from(literal))] | |
170 | #[inline] | |
171 | fn scale_range_uint<F, U>(v : F, &Bounds(a, b) : &Bounds<F>) -> U | |
172 | where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, | |
173 | U : Unsigned { | |
174 | debug_assert!(a < b); | |
175 | scale_uint(((v - a)/(b - a)).max(0.0).min(1.0)) | |
176 | } | |
177 | ||
178 | ||
179 | /// Sample a mapping on a grid. | |
180 | /// | |
181 | /// Returns a vector of values as well as upper and lower bounds of the values. | |
182 | fn rawdata_and_range<F, T>(grid : &LinGrid<F, 2>, g :&T) -> (Vec<F>, Bounds<F>) | |
183 | where F : Float, | |
184 | T : RealMapping<F, 2> { | |
185 | let rawdata : Vec<F> = grid.into_iter().map(|x| g.apply(&x)).collect(); | |
186 | let range = rawdata.iter() | |
187 | .map(|&v| Bounds(v, v)) | |
188 | .reduce(|b1, b2| b1.common(&b2)) | |
189 | .unwrap(); | |
190 | (rawdata, range) | |
191 | } | |
192 | ||
193 | /*fn to_range<'a, F, U>(rawdata : &'a Vec<F>, range : &'a Bounds<F>) | |
194 | -> std::iter::Map<std::slice::Iter<'a, F>, impl FnMut(&'a F) -> U> | |
195 | where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, | |
196 | U : Unsigned { | |
197 | rawdata.iter().map(move |&v| scale_range_uint(v, range)) | |
198 | }*/ | |
199 | ||
200 | /// Convert a scalar value to an RGB triplet. | |
201 | /// | |
202 | /// Converts the value `v` supposed to be within the range `[a, b]` to an rgb value according | |
203 | /// to the given `ramp` of equally-spaced rgb interpolation points. | |
204 | #[replace_float_literals(F::cast_from(literal))] | |
205 | fn one_to_ramp<F, U>( | |
206 | &Bounds(a, b) : &Bounds<F>, | |
207 | ramp : &Vec<Loc<F, 3>>, | |
208 | v : F, | |
209 | ) -> Rgb<U> | |
210 | where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, | |
211 | U : Unsigned { | |
212 | ||
213 | let n = ramp.len() - 1; | |
214 | let m = F::cast_from(U::RANGE_MAX); | |
215 | let ramprange = move |v : F| {let m : usize = v.as_(); m.min(n).max(0) }; | |
216 | ||
217 | let w = F::cast_from(n) * (v - a) / (b - a); // convert [0, 1] to [0, n] | |
218 | let (l, u) = (w.floor(), w.ceil()); // Find closest integers | |
219 | let (rl, ru) = (ramprange(l), ramprange(u)); | |
220 | let (cl, cu) = (ramp[rl], ramp[ru]); // Get corresponding colours | |
221 | let λ = match rl==ru { // Interpolation factor | |
222 | true => 0.0, | |
223 | false => (u - w) / (u - l), | |
224 | }; | |
225 | let Loc(rgb) = cl * λ + cu * (1.0 - λ); // Interpolate | |
226 | ||
227 | Rgb(rgb.map(|v| (v * m).round().min(m).max(0.0).as_())) | |
228 | } | |
229 | ||
230 | /// Convert a an iterator over scalar values to an iterator over RGB triplets. | |
231 | /// | |
232 | /// The conversion is that performed by [`one_to_ramp`]. | |
233 | #[replace_float_literals(F::cast_from(literal))] | |
234 | fn to_ramp<'a, F, U, I>( | |
235 | bounds : &'a Bounds<F>, | |
236 | ramp : &'a Vec<Loc<F, 3>>, | |
237 | iter : I, | |
238 | ) -> std::iter::Map<I, impl FnMut(F) -> Rgb<U> + 'a> | |
239 | where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, | |
240 | U : Unsigned, | |
241 | I : Iterator<Item = F> + 'a { | |
242 | iter.map(move |v| one_to_ramp(bounds, ramp, v)) | |
243 | } | |
244 | ||
245 | /// Convert a [`colorbrewer`] sepcification to a ramp of rgb triplets. | |
246 | fn get_ramp<F : Float>((palette, nb) : (CbPalette, u32)) -> Vec<Loc<F, 3>> { | |
247 | let m = F::cast_from(u8::MAX); | |
248 | colorbrewer::get_color_ramp(palette, nb) | |
249 | .expect("Invalid colorbrewer ramp") | |
250 | .into_iter() | |
251 | .map(|rgb::RGB{r, g, b}| { | |
252 | [r, g, b].map(|c| F::cast_from(c) / m).into() | |
253 | }).collect() | |
254 | } | |
255 | ||
256 | /// Perform hue shifting of an RGB value. | |
257 | /// | |
258 | // The hue `ω` is in radians. | |
259 | #[replace_float_literals(F::cast_from(literal))] | |
260 | fn hueshift<F, U>(ω : F, Rgb([r_in, g_in, b_in]) : Rgb<U>) -> Rgb<U> | |
261 | where F : Float + CastFrom<U>, | |
262 | U : Unsigned { | |
263 | let m = F::cast_from(U::RANGE_MAX); | |
264 | let r = F::cast_from(r_in) / m; | |
265 | let g = F::cast_from(g_in) / m; | |
266 | let b = F::cast_from(b_in) / m; | |
267 | let u = ω.cos(); | |
268 | let w = ω.sin(); | |
269 | ||
270 | let nr = (0.299 + 0.701*u + 0.168*w) * r | |
271 | + (0.587 - 0.587*u + 0.330*w) * g | |
272 | + (0.114 - 0.114*u - 0.497*w) * b; | |
273 | let ng = (0.299 - 0.299*u - 0.328*w) * r | |
274 | + (0.587 + 0.413*u + 0.035*w) * g | |
275 | + (0.114 - 0.114*u + 0.292*w) *b; | |
276 | let nb = (0.299 - 0.3*u + 1.25*w) * r | |
277 | + (0.587 - 0.588*u - 1.05*w) * g | |
278 | + (0.114 + 0.886*u - 0.203*w) * b; | |
279 | ||
280 | Rgb([nr, ng, nb].map(scale_uint)) | |
281 | } | |
282 | ||
283 | ||
284 | impl Plotting<2> for PlotLookup { | |
285 | #[replace_float_literals(F::cast_from(literal))] | |
286 | fn plot_into_file_spikes< | |
287 | F : Float, | |
288 | T1 : RealMapping<F, 2>, | |
289 | T2 : RealMapping<F, 2> | |
290 | > ( | |
291 | _g_explanation : String, | |
292 | g : &T1, | |
293 | _ω_explanation : String, | |
294 | ω0 : Option<&T2>, | |
295 | grid : LinGrid<F, 2>, | |
296 | _bnd0 : Option<Bounds<F>>, | |
297 | μ : &DiscreteMeasure<Loc<F, 2>, F>, | |
298 | filename : String, | |
299 | ) { | |
300 | let [w, h] = grid.count; | |
301 | let (rawdata_g, range_g) = rawdata_and_range(&grid, g); | |
302 | let (rawdata_ω, range) = match ω0 { | |
303 | Some(ω) => { | |
304 | let (rawdata_ω, range_ω) = rawdata_and_range(&grid, ω); | |
305 | (rawdata_ω, range_g.common(&range_ω)) | |
306 | }, | |
307 | None => { | |
308 | let mut zeros = Vec::new(); | |
309 | zeros.resize(rawdata_g.len(), 0.0); | |
310 | (zeros, range_g) | |
311 | } | |
312 | }; | |
313 | let ramp = get_ramp(RAMP); | |
314 | let base_im_iter = to_ramp::<F, u16, _>(&range_g, &ramp, rawdata_g.iter().cloned()); | |
315 | let im_iter = izip!(base_im_iter, rawdata_g.iter(), rawdata_ω.iter()) | |
316 | .map(|(rgb, &v, &w)| { | |
317 | hueshift(2.0 * F::PI * (v - w).abs() / range.upper(), rgb) | |
318 | }); | |
319 | let mut img = ImageBuffer::new(w as u32, h as u32); | |
320 | img.pixels_mut() | |
321 | .zip(im_iter) | |
322 | .for_each(|(p, v)| *p = v); | |
323 | ||
324 | // Add spikes | |
325 | let m = μ.iter_masses().fold(F::ZERO, |m, α| m.max(α)); | |
326 | let μ_range = Bounds(F::ZERO, m); | |
327 | for &DeltaMeasure{ ref x, α } in μ.iter_spikes() { | |
328 | let [a, b] = map4(x, &grid.start, &grid.end, &grid.count, |&ξ, &a, &b, &n| { | |
329 | ((ξ-a)/(b-a)*F::cast_from(n)).as_() | |
330 | }); | |
331 | if a < w.as_() && b < h.as_() { | |
332 | let sc : u16 = scale_range_uint(α, &μ_range); | |
333 | // TODO: use max of points that map to this pixel. | |
334 | img[(a, b)] = Rgb([u16::MAX, u16::MAX, sc/2]); | |
335 | } | |
336 | } | |
337 | ||
338 | img.save_with_format(filename + ".png", ImageFormat::Png) | |
339 | .expect("Image save error"); | |
340 | } | |
341 | ||
342 | fn plot_into_file< | |
343 | F : Float, | |
344 | T1 : RealMapping<F, 2>, | |
345 | > ( | |
346 | g : &T1, | |
347 | grid : LinGrid<F, 2>, | |
348 | filename : String, | |
349 | _explanation : String | |
350 | ) { | |
351 | let [w, h] = grid.count; | |
352 | let (rawdata, range) = rawdata_and_range(&grid, g); | |
353 | let ramp = get_ramp(RAMP); | |
354 | let im_iter = to_ramp::<F, u16, _>(&range, &ramp, rawdata.iter().cloned()); | |
355 | let mut img = ImageBuffer::new(w as u32, h as u32); | |
356 | img.pixels_mut() | |
357 | .zip(im_iter) | |
358 | .for_each(|(p, v)| *p = v); | |
359 | img.save_with_format(filename.clone() + ".png", ImageFormat::Png) | |
360 | .expect("Image save error"); | |
361 | ||
362 | let csv_iter = grid.into_iter().zip(rawdata.iter()).map(|(Loc(x), &v)| (x, v)); | |
363 | let csv_f = filename + ".txt"; | |
364 | write_csv(csv_iter, csv_f).expect("CSV save error"); | |
365 | } | |
366 | ||
367 | } | |
368 | ||
369 | /// A helper structure for plotting a sequence of images. | |
370 | #[derive(Clone,Debug)] | |
371 | pub struct SeqPlotter<F : Float, const N : usize> { | |
372 | /// File name prefix | |
373 | prefix : String, | |
374 | /// Maximum number of plots to perform | |
375 | max_plots : usize, | |
376 | /// Sampling grid | |
377 | grid : LinGrid<F, N>, | |
378 | /// Current plot count | |
379 | plot_count : usize, | |
380 | } | |
381 | ||
382 | impl<F : Float, const N : usize> SeqPlotter<F, N> | |
383 | where PlotLookup : Plotting<N> { | |
384 | /// Creates a new sequence plotter instance | |
385 | pub fn new(prefix : String, max_plots : usize, grid : LinGrid<F, N>) -> Self { | |
386 | SeqPlotter { prefix, max_plots, grid, plot_count : 0 } | |
387 | } | |
388 | ||
389 | /// This calls [`PlotLookup::plot_into_file_spikes`] with a sequentially numbered file name. | |
390 | pub fn plot_spikes<T1, T2>( | |
391 | &mut self, | |
392 | g_explanation : String, | |
393 | g : &T1, | |
394 | ω_explanation : String, | |
395 | ω : Option<&T2>, | |
396 | tol : Option<Bounds<F>>, | |
397 | μ : &DiscreteMeasure<Loc<F, N>, F>, | |
398 | ) where T1 : RealMapping<F, N>, | |
399 | T2 : RealMapping<F, N> | |
400 | { | |
401 | if self.plot_count == 0 && self.max_plots > 0 { | |
402 | std::fs::create_dir_all(&self.prefix).expect("Unable to create plot directory"); | |
403 | } | |
404 | if self.plot_count < self.max_plots { | |
405 | PlotLookup::plot_into_file_spikes( | |
406 | g_explanation, g, | |
407 | ω_explanation, ω, | |
408 | self.grid, | |
409 | tol, | |
410 | μ, | |
411 | format!("{}out{:03}", self.prefix, self.plot_count) | |
412 | ); | |
413 | self.plot_count += 1; | |
414 | } | |
415 | } | |
416 | } |