diff -r 000000000000 -r eb3c7813b67a src/plot.rs --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/plot.rs Thu Dec 01 23:07:35 2022 +0200 @@ -0,0 +1,413 @@ +//! Plotting helper utilities + +use numeric_literals::replace_float_literals; +use std::io::Write; +use image::{ + ImageFormat, + ImageBuffer, + Rgb +}; +use itertools::izip; +use colorbrewer::Palette as CbPalette; + +use alg_tools::types::*; +use alg_tools::lingrid::LinGrid; +use alg_tools::mapping::RealMapping; +use alg_tools::loc::Loc; +use alg_tools::bisection_tree::Bounds; +use alg_tools::maputil::map4; +use alg_tools::tabledump::write_csv; +use crate::measures::*; + +/// Default RGB ramp from [`colorbrewer`]. +/// +/// This is a tuple of parameters to [`colorbrewer::get_color_ramp`]. +const RAMP : (CbPalette, u32) = (CbPalette::RdBu, 11); + +/// Helper trait for implementing dimension-dependent plotting routines. +pub trait Plotting { + /// Plot several mappings and a discrete measure into a file. + fn plot_into_file_spikes< + F : Float, + T1 : RealMapping, + T2 : RealMapping + > ( + g_explanation : String, + g : &T1, + ω_explanation : String, + ω : Option<&T2>, + grid : LinGrid, + bnd : Option>, + μ : &DiscreteMeasure, F>, + filename : String, + ); + + /// Plot a mapping into a file, sampling values on a given grid. + fn plot_into_file< + F : Float, + T1 : RealMapping, + > ( + g : &T1, + grid : LinGrid, + filename : String, + explanation : String + ); +} + +/// Helper type for looking up a [`Plotting`] based on dimension. +pub struct PlotLookup; + +impl Plotting<1> for PlotLookup { + fn plot_into_file_spikes< + F : Float, + T1 : RealMapping, + T2 : RealMapping + > ( + g_explanation : String, + g : &T1, + ω_explanation : String, + ω0 : Option<&T2>, + grid : LinGrid, + bnd0 : Option>, + μ : &DiscreteMeasure, F>, + filename : String, + ) { + let start = grid.start[0].as_(); + let end = grid.end[0].as_(); + let m = μ.iter_masses().fold(F::ZERO, |m, α| m.max(α)); + let s = μ.iter_masses().fold(F::ZERO, |m, α| m.add(α)); + let mut spike_scale = F::ONE; + + let mut plotter = poloto::plot( + "f", "x", + format!("f(x); spike max={:.4}, n={}, ∑={:.4}", m, μ.len(), s) + ).move_into(); + + if let Some(ω) = ω0 { + let graph_ω = grid.into_iter().map(|x@Loc([x0]) : Loc| { + [x0.as_(), ω.apply(&x).as_()] + }); + plotter.line(ω_explanation.as_str(), graph_ω.clone()); + // let csv_f = format!("{}.txt", filename); + // write_csv(graph_ω, csv_f).expect("CSV save error"); + } + + let graph_g = grid.into_iter().map(|x@Loc([x0]) : Loc| { + [x0.as_(), g.apply(&x).as_()] + }); + plotter.line(g_explanation.as_str(), graph_g.clone()); + // let csv_f = format!("{}.txt", filename); + // write_csv(graph_g, csv_f).expect("CSV save error"); + + bnd0.map(|bnd| { + let upperb = bnd.upper().as_(); + let lowerb = bnd.lower().as_(); + let upper : [[f64; 2]; 2] = [[start, upperb], [end, upperb]]; + let lower = [[start, lowerb], [end, lowerb]]; + spike_scale *= bnd.upper(); + + plotter.line("upper bound", upper) + .line("lower bound", lower) + .ymarker(lowerb) + .ymarker(upperb); + }); + + for &DeltaMeasure{ α, x : Loc([x]) } in μ.iter_spikes() { + let spike = [[x.as_(), 0.0], [x.as_(), (α/m * spike_scale).as_()]]; + plotter.line("", spike); + } + + let svg = format!("{}", poloto::disp(|a| poloto::simple_theme(a, plotter))); + + std::fs::File::create(filename + ".svg").and_then(|mut file| + file.write_all(svg.as_bytes()) + ).expect("SVG save error"); + } + + fn plot_into_file< + F : Float, + T1 : RealMapping, + > ( + g : &T1, + grid : LinGrid, + filename : String, + explanation : String + ) { + let graph_g = grid.into_iter().map(|x@Loc([x0]) : Loc| { + [x0.as_(), g.apply(&x).as_()] + }); + + let plotter: poloto::Plotter<'_, float, float> = poloto::plot("f", "x", "f(x)") + .line(explanation.as_str(), graph_g.clone()) + .move_into(); + + let svg = format!("{}", poloto::disp(|a| poloto::simple_theme(a, plotter))); + + let svg_f = format!("{}.svg", filename); + std::fs::File::create(svg_f).and_then(|mut file| + file.write_all(svg.as_bytes()) + ).expect("SVG save error"); + + let csv_f = format!("{}.txt", filename); + write_csv(graph_g, csv_f).expect("CSV save error"); + } + +} + +/// Convert $[0, 1] ∈ F$ to $\\\{0, …, M\\\} ∈ F$ where $M=$`F::RANGE_MAX`. +#[inline] +fn scale_uint(v : F) -> U +where F : Float + CastFrom + num_traits::cast::AsPrimitive, + U : Unsigned { + (v*F::cast_from(U::RANGE_MAX)).as_() +} + +/// Convert $[a, b] ∈ F$ to $\\\{0, …, M\\\} ∈ F$ where $M=$`F::RANGE_MAX`. +#[replace_float_literals(F::cast_from(literal))] +#[inline] +fn scale_range_uint(v : F, &Bounds(a, b) : &Bounds) -> U +where F : Float + CastFrom + num_traits::cast::AsPrimitive, + U : Unsigned { + debug_assert!(a < b); + scale_uint(((v - a)/(b - a)).max(0.0).min(1.0)) +} + + +/// Sample a mapping on a grid. +/// +/// Returns a vector of values as well as upper and lower bounds of the values. +fn rawdata_and_range(grid : &LinGrid, g :&T) -> (Vec, Bounds) +where F : Float, + T : RealMapping { + let rawdata : Vec = grid.into_iter().map(|x| g.apply(&x)).collect(); + let range = rawdata.iter() + .map(|&v| Bounds(v, v)) + .reduce(|b1, b2| b1.common(&b2)) + .unwrap(); + (rawdata, range) +} + +/*fn to_range<'a, F, U>(rawdata : &'a Vec, range : &'a Bounds) +-> std::iter::Map, impl FnMut(&'a F) -> U> +where F : Float + CastFrom + num_traits::cast::AsPrimitive, + U : Unsigned { + rawdata.iter().map(move |&v| scale_range_uint(v, range)) +}*/ + +/// Convert a scalar value to an RGB triplet. +/// +/// Converts the value `v` supposed to be within the range `[a, b]` to an rgb value according +/// to the given `ramp` of equally-spaced rgb interpolation points. +#[replace_float_literals(F::cast_from(literal))] +fn one_to_ramp( + &Bounds(a, b) : &Bounds, + ramp : &Vec>, + v : F, +) -> Rgb +where F : Float + CastFrom + num_traits::cast::AsPrimitive, + U : Unsigned { + + let n = ramp.len() - 1; + let m = F::cast_from(U::RANGE_MAX); + let ramprange = move |v : F| {let m : usize = v.as_(); m.min(n).max(0) }; + + let w = F::cast_from(n) * (v - a) / (b - a); // convert [0, 1] to [0, n] + let (l, u) = (w.floor(), w.ceil()); // Find closest integers + let (rl, ru) = (ramprange(l), ramprange(u)); + let (cl, cu) = (ramp[rl], ramp[ru]); // Get corresponding colours + let λ = match rl==ru { // Interpolation factor + true => 0.0, + false => (u - w) / (u - l), + }; + let Loc(rgb) = cl * λ + cu * (1.0 - λ); // Interpolate + + Rgb(rgb.map(|v| (v * m).round().min(m).max(0.0).as_())) +} + +/// Convert a an iterator over scalar values to an iterator over RGB triplets. +/// +/// The conversion is that performed by [`one_to_ramp`]. +#[replace_float_literals(F::cast_from(literal))] +fn to_ramp<'a, F, U, I>( + bounds : &'a Bounds, + ramp : &'a Vec>, + iter : I, +) -> std::iter::Map Rgb + 'a> +where F : Float + CastFrom + num_traits::cast::AsPrimitive, + U : Unsigned, + I : Iterator + 'a { + iter.map(move |v| one_to_ramp(bounds, ramp, v)) +} + +/// Convert a [`colorbrewer`] sepcification to a ramp of rgb triplets. +fn get_ramp((palette, nb) : (CbPalette, u32)) -> Vec> { + let m = F::cast_from(u8::MAX); + colorbrewer::get_color_ramp(palette, nb) + .expect("Invalid colorbrewer ramp") + .into_iter() + .map(|rgb::RGB{r, g, b}| { + [r, g, b].map(|c| F::cast_from(c) / m).into() + }).collect() +} + +/// Perform hue shifting of an RGB value. +/// +// The hue `ω` is in radians. +#[replace_float_literals(F::cast_from(literal))] +fn hueshift(ω : F, Rgb([r_in, g_in, b_in]) : Rgb) -> Rgb +where F : Float + CastFrom, + U : Unsigned { + let m = F::cast_from(U::RANGE_MAX); + let r = F::cast_from(r_in) / m; + let g = F::cast_from(g_in) / m; + let b = F::cast_from(b_in) / m; + let u = ω.cos(); + let w = ω.sin(); + + let nr = (0.299 + 0.701*u + 0.168*w) * r + + (0.587 - 0.587*u + 0.330*w) * g + + (0.114 - 0.114*u - 0.497*w) * b; + let ng = (0.299 - 0.299*u - 0.328*w) * r + + (0.587 + 0.413*u + 0.035*w) * g + + (0.114 - 0.114*u + 0.292*w) *b; + let nb = (0.299 - 0.3*u + 1.25*w) * r + + (0.587 - 0.588*u - 1.05*w) * g + + (0.114 + 0.886*u - 0.203*w) * b; + + Rgb([nr, ng, nb].map(scale_uint)) +} + + +impl Plotting<2> for PlotLookup { + #[replace_float_literals(F::cast_from(literal))] + fn plot_into_file_spikes< + F : Float, + T1 : RealMapping, + T2 : RealMapping + > ( + _g_explanation : String, + g : &T1, + _ω_explanation : String, + ω0 : Option<&T2>, + grid : LinGrid, + _bnd0 : Option>, + μ : &DiscreteMeasure, F>, + filename : String, + ) { + let [w, h] = grid.count; + let (rawdata_g, range_g) = rawdata_and_range(&grid, g); + let (rawdata_ω, range) = match ω0 { + Some(ω) => { + let (rawdata_ω, range_ω) = rawdata_and_range(&grid, ω); + (rawdata_ω, range_g.common(&range_ω)) + }, + None => { + let mut zeros = Vec::new(); + zeros.resize(rawdata_g.len(), 0.0); + (zeros, range_g) + } + }; + let ramp = get_ramp(RAMP); + let base_im_iter = to_ramp::(&range_g, &ramp, rawdata_g.iter().cloned()); + let im_iter = izip!(base_im_iter, rawdata_g.iter(), rawdata_ω.iter()) + .map(|(rgb, &v, &w)| { + hueshift(2.0 * F::PI * (v - w).abs() / range.upper(), rgb) + }); + let mut img = ImageBuffer::new(w as u32, h as u32); + img.pixels_mut() + .zip(im_iter) + .for_each(|(p, v)| *p = v); + + // Add spikes + let m = μ.iter_masses().fold(F::ZERO, |m, α| m.max(α)); + let μ_range = Bounds(F::ZERO, m); + for &DeltaMeasure{ ref x, α } in μ.iter_spikes() { + let [a, b] = map4(x, &grid.start, &grid.end, &grid.count, |&ξ, &a, &b, &n| { + ((ξ-a)/(b-a)*F::cast_from(n)).as_() + }); + if a < w.as_() && b < h.as_() { + let sc : u16 = scale_range_uint(α, &μ_range); + // TODO: use max of points that map to this pixel. + img[(a, b)] = Rgb([u16::MAX, u16::MAX, sc/2]); + } + } + + img.save_with_format(filename + ".png", ImageFormat::Png) + .expect("Image save error"); + } + + fn plot_into_file< + F : Float, + T1 : RealMapping, + > ( + g : &T1, + grid : LinGrid, + filename : String, + _explanation : String + ) { + let [w, h] = grid.count; + let (rawdata, range) = rawdata_and_range(&grid, g); + let ramp = get_ramp(RAMP); + let im_iter = to_ramp::(&range, &ramp, rawdata.iter().cloned()); + let mut img = ImageBuffer::new(w as u32, h as u32); + img.pixels_mut() + .zip(im_iter) + .for_each(|(p, v)| *p = v); + img.save_with_format(filename.clone() + ".png", ImageFormat::Png) + .expect("Image save error"); + + let csv_iter = grid.into_iter().zip(rawdata.iter()).map(|(Loc(x), &v)| (x, v)); + let csv_f = filename + ".txt"; + write_csv(csv_iter, csv_f).expect("CSV save error"); + } + +} + +/// A helper structure for plotting a sequence of images. +#[derive(Clone,Debug)] +pub struct SeqPlotter { + /// File name prefix + prefix : String, + /// Maximum number of plots to perform + max_plots : usize, + /// Sampling grid + grid : LinGrid, + /// Current plot count + plot_count : usize, +} + +impl SeqPlotter +where PlotLookup : Plotting { + /// Creates a new sequence plotter instance + pub fn new(prefix : String, max_plots : usize, grid : LinGrid) -> Self { + SeqPlotter { prefix, max_plots, grid, plot_count : 0 } + } + + /// This calls [`PlotLookup::plot_into_file_spikes`] with a sequentially numbered file name. + pub fn plot_spikes( + &mut self, + g_explanation : String, + g : &T1, + ω_explanation : String, + ω : Option<&T2>, + tol : Option>, + μ : &DiscreteMeasure, F>, + ) where T1 : RealMapping, + T2 : RealMapping + { + if self.plot_count == 0 && self.max_plots > 0 { + std::fs::create_dir_all(&self.prefix).expect("Unable to create plot directory"); + } + if self.plot_count < self.max_plots { + PlotLookup::plot_into_file_spikes( + g_explanation, g, + ω_explanation, ω, + self.grid, + tol, + μ, + format!("{}out{:03}", self.prefix, self.plot_count) + ); + self.plot_count += 1; + } + } +}