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