src/plot.rs

changeset 52
f0e8704d3f0e
parent 35
b087e3eab191
equal deleted inserted replaced
31:6105b5cd8d89 52:f0e8704d3f0e
1 //! Plotting helper utilities 1 //! Plotting helper utilities
2 2
3 use numeric_literals::replace_float_literals; 3 use numeric_literals::replace_float_literals;
4 use std::io::Write; 4 use serde::Serialize;
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::*; 5 use alg_tools::types::*;
14 use alg_tools::lingrid::LinGrid; 6 use alg_tools::lingrid::LinGrid;
15 use alg_tools::mapping::RealMapping; 7 use alg_tools::mapping::RealMapping;
16 use alg_tools::loc::Loc; 8 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; 9 use alg_tools::tabledump::write_csv;
20 use crate::measures::*; 10 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 11
27 /// Helper trait for implementing dimension-dependent plotting routines. 12 /// Helper trait for implementing dimension-dependent plotting routines.
28 pub trait Plotting<const N : usize> { 13 pub trait Plotting<const N : usize> {
29 /// Plot several mappings and a discrete measure into a file. 14 /// Plot several mappings and a discrete measure into a file.
30 fn plot_into_file_spikes< 15 fn plot_into_file_spikes<
31 F : Float, 16 F : Float,
32 T1 : RealMapping<F, N>, 17 T1 : RealMapping<F, N>,
33 T2 : RealMapping<F, N> 18 T2 : RealMapping<F, N>
34 > ( 19 > (
35 g_explanation : String, 20 g : Option<&T1>,
36 g : &T1,
37 ω_explanation : String,
38 ω : Option<&T2>, 21 ω : Option<&T2>,
39 grid : LinGrid<F, N>, 22 grid : LinGrid<F, N>,
40 bnd : Option<Bounds<F>>, 23 μ : &RNDM<F, N>,
41 μ : &DiscreteMeasure<Loc<F, N>, F>,
42 filename : String, 24 filename : String,
43 ); 25 );
44 26
45 /// Plot a mapping into a file, sampling values on a given grid. 27 /// Plot a mapping into a file, sampling values on a given grid.
46 fn plot_into_file< 28 fn plot_into_file<
48 T1 : RealMapping<F, N>, 30 T1 : RealMapping<F, N>,
49 > ( 31 > (
50 g : &T1, 32 g : &T1,
51 grid : LinGrid<F, N>, 33 grid : LinGrid<F, N>,
52 filename : String, 34 filename : String,
53 explanation : String
54 ); 35 );
55 } 36 }
56 37
57 /// Helper type for looking up a [`Plotting`] based on dimension. 38 /// Helper type for looking up a [`Plotting`] based on dimension.
58 pub struct PlotLookup; 39 pub struct PlotLookup;
59 40
41 #[derive(Serialize)]
42 struct CSVHelper1<F : Float> {
43 x : F,
44 f : F,
45 }
46
47 #[derive(Serialize)]
48 struct CSVHelper1_2<F : Float>{
49 x : F,
50 g : Option<F>,
51 omega : Option<F>
52 }
53
54 #[derive(Serialize)]
55 struct CSVSpike1<F : Float> {
56 x : F,
57 alpha : F,
58 }
59
60 impl Plotting<1> for PlotLookup { 60 impl Plotting<1> for PlotLookup {
61 fn plot_into_file_spikes< 61 fn plot_into_file_spikes<
62 F : Float, 62 F : Float,
63 T1 : RealMapping<F, 1>, 63 T1 : RealMapping<F, 1>,
64 T2 : RealMapping<F, 1> 64 T2 : RealMapping<F, 1>
65 > ( 65 > (
66 g_explanation : String, 66 g0 : Option<&T1>,
67 g : &T1,
68 ω_explanation : String,
69 ω0 : Option<&T2>, 67 ω0 : Option<&T2>,
70 grid : LinGrid<F, 1>, 68 grid : LinGrid<F, 1>,
71 bnd0 : Option<Bounds<F>>,
72 μ : &DiscreteMeasure<Loc<F, 1>, F>, 69 μ : &DiscreteMeasure<Loc<F, 1>, F>,
73 filename : String, 70 filename : String,
74 ) { 71 ) {
75 let start = grid.start[0].as_(); 72 let data = grid.into_iter().map(|p@Loc([x]) : Loc<F, 1>| CSVHelper1_2 {
76 let end = grid.end[0].as_(); 73 x,
77 let m = μ.iter_masses().fold(F::ZERO, |m, α| m.max(α)); 74 g : g0.map(|g| g.apply(&p)),
78 let s = μ.iter_masses().fold(F::ZERO, |m, α| m.add(α)); 75 omega : ω0.map(|ω| ω.apply(&p))
79 let mut spike_scale = F::ONE; 76 });
80 77 let csv_f = format!("{}_functions.csv", filename);
81 let mut plotter = poloto::plot( 78 write_csv(data, csv_f).expect("CSV save error");
82 "f", "x", 79
83 format!("f(x); spike max={:.4}, n={}, ∑={:.4}", m, μ.len(), s) 80 let spikes = μ.iter_spikes().map(|δ| {
84 ).move_into(); 81 let Loc([x]) = δ.x;
85 82 CSVSpike1 { x, alpha : δ.α }
86 if let Some(ω) = ω0 { 83 });
87 let graph_ω = grid.into_iter().map(|x@Loc([x0]) : Loc<F, 1>| { 84 let csv_f = format!("{}_spikes.csv", filename);
88 [x0.as_(), ω.apply(&x).as_()] 85 write_csv(spikes, csv_f).expect("CSV save error");
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 } 86 }
126 87
127 fn plot_into_file< 88 fn plot_into_file<
128 F : Float, 89 F : Float,
129 T1 : RealMapping<F, 1>, 90 T1 : RealMapping<F, 1>,
130 > ( 91 > (
131 g : &T1, 92 g : &T1,
132 grid : LinGrid<F, 1>, 93 grid : LinGrid<F, 1>,
133 filename : String, 94 filename : String,
134 explanation : String 95 ) {
135 ) { 96 let data = grid.into_iter().map(|p@Loc([x]) : Loc<F, 1>| CSVHelper1 {
136 let graph_g = grid.into_iter().map(|x@Loc([x0]) : Loc<F, 1>| { 97 x,
137 [x0.as_(), g.apply(&x).as_()] 98 f : g.apply(&p),
138 }); 99 });
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); 100 let csv_f = format!("{}.txt", filename);
152 write_csv(graph_g, csv_f).expect("CSV save error"); 101 write_csv(data, csv_f).expect("CSV save error");
153 } 102 }
154 103
155 } 104 }
156 105
157 /// Convert $[0, 1] ∈ F$ to $\\\{0, …, M\\\} ∈ F$ where $M=$`F::RANGE_MAX`. 106 #[derive(Serialize)]
158 #[inline] 107 struct CSVHelper2<F : Float> {
159 fn scale_uint<F, U>(v : F) -> U 108 x : F,
160 where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, 109 y : F,
161 U : Unsigned { 110 f : F,
162 (v*F::cast_from(U::RANGE_MAX)).as_() 111 }
163 } 112
164 113 #[derive(Serialize)]
165 /// Convert $[a, b] ∈ F$ to $\\\{0, …, M\\\} ∈ F$ where $M=$`F::RANGE_MAX`. 114 struct CSVHelper2_2<F : Float>{
166 #[replace_float_literals(F::cast_from(literal))] 115 x : F,
167 #[inline] 116 y : F,
168 fn scale_range_uint<F, U>(v : F, &Bounds(a, b) : &Bounds<F>) -> U 117 g : Option<F>,
169 where F : Float + CastFrom<U> + num_traits::cast::AsPrimitive<U>, 118 omega : Option<F>
170 U : Unsigned { 119 }
171 debug_assert!(a < b); 120
172 scale_uint(((v - a)/(b - a)).max(0.0).min(1.0)) 121 #[derive(Serialize)]
173 } 122 struct CSVSpike2<F : Float> {
174 123 x : F,
175 124 y : F,
176 /// Sample a mapping on a grid. 125 alpha : F,
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 } 126 }
279 127
280 128
281 impl Plotting<2> for PlotLookup { 129 impl Plotting<2> for PlotLookup {
282 #[replace_float_literals(F::cast_from(literal))] 130 #[replace_float_literals(F::cast_from(literal))]
283 fn plot_into_file_spikes< 131 fn plot_into_file_spikes<
284 F : Float, 132 F : Float,
285 T1 : RealMapping<F, 2>, 133 T1 : RealMapping<F, 2>,
286 T2 : RealMapping<F, 2> 134 T2 : RealMapping<F, 2>
287 > ( 135 > (
288 _g_explanation : String, 136 g0 : Option<&T1>,
289 g : &T1,
290 _ω_explanation : String,
291 ω0 : Option<&T2>, 137 ω0 : Option<&T2>,
292 grid : LinGrid<F, 2>, 138 grid : LinGrid<F, 2>,
293 _bnd0 : Option<Bounds<F>>,
294 μ : &DiscreteMeasure<Loc<F, 2>, F>, 139 μ : &DiscreteMeasure<Loc<F, 2>, F>,
295 filename : String, 140 filename : String,
296 ) { 141 ) {
297 let [w, h] = grid.count; 142 let data = grid.into_iter().map(|p@Loc([x, y]) : Loc<F, 2>| CSVHelper2_2 {
298 let (rawdata_g, range_g) = rawdata_and_range(&grid, g); 143 x,
299 let (rawdata_ω, range) = match ω0 { 144 y,
300 Some(ω) => { 145 g : g0.map(|g| g.apply(&p)),
301 let (rawdata_ω, range_ω) = rawdata_and_range(&grid, ω); 146 omega : ω0.map(|ω| ω.apply(&p))
302 (rawdata_ω, range_g.common(&range_ω)) 147 });
303 }, 148 let csv_f = format!("{}_functions.csv", filename);
304 None => { 149 write_csv(data, csv_f).expect("CSV save error");
305 let mut zeros = Vec::new(); 150
306 zeros.resize(rawdata_g.len(), 0.0); 151 let spikes = μ.iter_spikes().map(|δ| {
307 (zeros, range_g) 152 let Loc([x, y]) = δ.x;
308 } 153 CSVSpike2 { x, y, alpha : δ.α }
309 }; 154 });
310 let ramp = get_ramp(RAMP); 155 let csv_f = format!("{}_spikes.csv", filename);
311 let base_im_iter = to_ramp::<F, u16, _>(&range_g, &ramp, rawdata_g.iter().cloned()); 156 write_csv(spikes, csv_f).expect("CSV save error");
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 } 157 }
338 158
339 fn plot_into_file< 159 fn plot_into_file<
340 F : Float, 160 F : Float,
341 T1 : RealMapping<F, 2>, 161 T1 : RealMapping<F, 2>,
342 > ( 162 > (
343 g : &T1, 163 g : &T1,
344 grid : LinGrid<F, 2>, 164 grid : LinGrid<F, 2>,
345 filename : String, 165 filename : String,
346 _explanation : String 166 ) {
347 ) { 167 let data = grid.into_iter().map(|p@Loc([x, y]) : Loc<F, 2>| CSVHelper2 {
348 let [w, h] = grid.count; 168 x,
349 let (rawdata, range) = rawdata_and_range(&grid, g); 169 y,
350 let ramp = get_ramp(RAMP); 170 f : g.apply(&p),
351 let im_iter = to_ramp::<F, u16, _>(&range, &ramp, rawdata.iter().cloned()); 171 });
352 let mut img = ImageBuffer::new(w as u32, h as u32); 172 let csv_f = format!("{}.txt", filename);
353 img.pixels_mut() 173 write_csv(data, csv_f).expect("CSV save error");
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 } 174 }
363 175
364 } 176 }
365 177
366 /// A helper structure for plotting a sequence of images. 178 /// A helper structure for plotting a sequence of images.
384 } 196 }
385 197
386 /// This calls [`PlotLookup::plot_into_file_spikes`] with a sequentially numbered file name. 198 /// This calls [`PlotLookup::plot_into_file_spikes`] with a sequentially numbered file name.
387 pub fn plot_spikes<T1, T2>( 199 pub fn plot_spikes<T1, T2>(
388 &mut self, 200 &mut self,
389 g_explanation : String, 201 iter : usize,
390 g : &T1, 202 g : Option<&T1>,
391 ω_explanation : String,
392 ω : Option<&T2>, 203 ω : Option<&T2>,
393 tol : Option<Bounds<F>>, 204 μ : &RNDM<F, N>,
394 μ : &DiscreteMeasure<Loc<F, N>, F>,
395 ) where T1 : RealMapping<F, N>, 205 ) where T1 : RealMapping<F, N>,
396 T2 : RealMapping<F, N> 206 T2 : RealMapping<F, N>
397 { 207 {
398 if self.plot_count == 0 && self.max_plots > 0 { 208 if self.plot_count == 0 && self.max_plots > 0 {
399 std::fs::create_dir_all(&self.prefix).expect("Unable to create plot directory"); 209 std::fs::create_dir_all(&self.prefix).expect("Unable to create plot directory");
400 } 210 }
401 if self.plot_count < self.max_plots { 211 if self.plot_count < self.max_plots {
402 PlotLookup::plot_into_file_spikes( 212 PlotLookup::plot_into_file_spikes(
403 g_explanation, g, 213 g,
404 ω_explanation, ω, 214 ω,
405 self.grid, 215 self.grid,
406 tol,
407 μ, 216 μ,
408 format!("{}out{:03}", self.prefix, self.plot_count) 217 format!("{}out{:03}", self.prefix, iter)
409 ); 218 );
410 self.plot_count += 1; 219 self.plot_count += 1;
411 } 220 }
412 } 221 }
413 } 222 }

mercurial