src/plot.rs

//! Plotting helper utilities

use crate::measures::*;
use alg_tools::lingrid::LinGrid;
use alg_tools::loc::Loc;
use alg_tools::mapping::RealMapping;
use alg_tools::tabledump::write_csv;
use alg_tools::types::*;
use numeric_literals::replace_float_literals;
use serde::Serialize;

/// Helper trait for implementing dimension-dependent plotting routines.
pub trait Plotting<const N: usize> {
    /// Plot several mappings and a discrete measure into a file.
    fn plot_into_file_spikes<F: Float, T1: RealMapping<N, F>, T2: RealMapping<N, F>>(
        g: Option<&T1>,
        ω: Option<&T2>,
        grid: LinGrid<N, F>,
        μ: &RNDM<N, F>,
        filename: String,
    );

    /// Plot a mapping into a file, sampling values on a given grid.
    fn plot_into_file<F: Float, T1: RealMapping<N, F>>(
        g: &T1,
        grid: LinGrid<N, F>,
        filename: String,
    );
}

/// Helper type for looking up a [`Plotting`] based on dimension.
pub struct PlotLookup;

#[derive(Serialize)]
struct CSVHelper1<F: Float> {
    x: F,
    f: F,
}

#[derive(Serialize)]
struct CSVHelper1_2<F: Float> {
    x: F,
    g: Option<F>,
    omega: Option<F>,
}

#[derive(Serialize)]
struct CSVSpike1<F: Float> {
    x: F,
    alpha: F,
}

impl Plotting<1> for PlotLookup {
    fn plot_into_file_spikes<F: Float, T1: RealMapping<1, F>, T2: RealMapping<1, F>>(
        g0: Option<&T1>,
        ω0: Option<&T2>,
        grid: LinGrid<1, F>,
        μ: &DiscreteMeasure<Loc<1, F>, F>,
        filename: String,
    ) {
        let data = grid
            .into_iter()
            .map(|p @ Loc([x]): Loc<1, F>| CSVHelper1_2 {
                x,
                g: g0.map(|g| g.apply(&p)),
                omega: ω0.map(|ω| ω.apply(&p)),
            });
        let csv_f = format!("{}_functions.csv", filename);
        write_csv(data, csv_f).expect("CSV save error");

        let spikes = μ.iter_spikes().map(|δ| {
            let Loc([x]) = δ.x;
            CSVSpike1 { x, alpha: δ.α }
        });
        let csv_f = format!("{}_spikes.csv", filename);
        write_csv(spikes, csv_f).expect("CSV save error");
    }

    fn plot_into_file<F: Float, T1: RealMapping<1, F>>(
        g: &T1,
        grid: LinGrid<1, F>,
        filename: String,
    ) {
        let data = grid
            .into_iter()
            .map(|p @ Loc([x]): Loc<1, F>| CSVHelper1 { x, f: g.apply(&p) });
        let csv_f = format!("{}.txt", filename);
        write_csv(data, csv_f).expect("CSV save error");
    }
}

#[derive(Serialize)]
struct CSVHelper2<F: Float> {
    x: F,
    y: F,
    f: F,
}

#[derive(Serialize)]
struct CSVHelper2_2<F: Float> {
    x: F,
    y: F,
    g: Option<F>,
    omega: Option<F>,
}

#[derive(Serialize)]
struct CSVSpike2<F: Float> {
    x: F,
    y: F,
    alpha: F,
}

impl Plotting<2> for PlotLookup {
    #[replace_float_literals(F::cast_from(literal))]
    fn plot_into_file_spikes<F: Float, T1: RealMapping<2, F>, T2: RealMapping<2, F>>(
        g0: Option<&T1>,
        ω0: Option<&T2>,
        grid: LinGrid<2, F>,
        μ: &DiscreteMeasure<Loc<2, F>, F>,
        filename: String,
    ) {
        let data = grid
            .into_iter()
            .map(|p @ Loc([x, y]): Loc<2, F>| CSVHelper2_2 {
                x,
                y,
                g: g0.map(|g| g.apply(&p)),
                omega: ω0.map(|ω| ω.apply(&p)),
            });
        let csv_f = format!("{}_functions.csv", filename);
        write_csv(data, csv_f).expect("CSV save error");

        let spikes = μ.iter_spikes().map(|δ| {
            let Loc([x, y]) = δ.x;
            CSVSpike2 { x, y, alpha: δ.α }
        });
        let csv_f = format!("{}_spikes.csv", filename);
        write_csv(spikes, csv_f).expect("CSV save error");
    }

    fn plot_into_file<F: Float, T1: RealMapping<2, F>>(
        g: &T1,
        grid: LinGrid<2, F>,
        filename: String,
    ) {
        let data = grid
            .into_iter()
            .map(|p @ Loc([x, y]): Loc<2, F>| CSVHelper2 {
                x,
                y,
                f: g.apply(&p),
            });
        let csv_f = format!("{}.txt", filename);
        write_csv(data, csv_f).expect("CSV save error");
    }
}

/// Trait for plotters
pub trait Plotter<T1, T2, M> {
    /// Plot the functions `g` and `ω` as well as the spikes of `μ`.
    fn plot_spikes(&mut self, iter: usize, g: Option<&T1>, ω: Option<&T2>, μ: &M);
}

/// A plotter that does nothing.
pub struct NoPlotter;

impl<T1, T2, M> Plotter<T1, T2, M> for NoPlotter {
    fn plot_spikes(&mut self, _iter: usize, _g: Option<&T1>, _ω: Option<&T2>, _μ: &M) {}
}

/// A basic plotter.
///
/// This calls [`PlotLookup::plot_into_file_spikes`] with a sequentially numbered file name.
#[derive(Clone, Debug)]
pub struct SeqPlotter<const N: usize, F: Float = f64> {
    /// File name prefix
    prefix: String,
    /// Maximum number of plots to perform
    max_plots: usize,
    /// Sampling grid
    grid: LinGrid<N, F>,
    /// Current plot count
    plot_count: usize,
}

impl<F: Float, const N: usize> SeqPlotter<N, F>
where
    PlotLookup: Plotting<N>,
{
    /// Creates a new sequence plotter instance
    pub fn new(prefix: String, max_plots: usize, grid: LinGrid<N, F>) -> Self {
        SeqPlotter {
            prefix,
            max_plots,
            grid,
            plot_count: 0,
        }
    }
}

impl<F, T1, T2, const N: usize> Plotter<T1, T2, RNDM<N, F>> for SeqPlotter<N, F>
where
    F: Float,
    T1: RealMapping<N, F>,
    T2: RealMapping<N, F>,
    PlotLookup: Plotting<N>,
{
    fn plot_spikes(&mut self, iter: usize, g: Option<&T1>, ω: Option<&T2>, μ: &RNDM<N, F>) {
        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,
                ω,
                self.grid,
                μ,
                format!("{}out{:03}", self.prefix, iter),
            );
            self.plot_count += 1;
        }
    }
}