src/plot.rs

//! Plotting helper utilities

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

/// 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<F, N>,
        T2 : RealMapping<F, N>
    > (
        g : Option<&T1>,
        ω : Option<&T2>,
        grid : LinGrid<F, N>,
        μ : &RNDM<F, N>,
        filename : String,
    );

    /// Plot a mapping into a file, sampling values on a given grid.
    fn plot_into_file<
        F : Float,
        T1 : RealMapping<F, N>,
    > (
        g : &T1,
        grid : LinGrid<F, N>,
        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<F, 1>,
        T2 : RealMapping<F, 1>
    > (
        g0 : Option<&T1>,
        ω0 : Option<&T2>,
        grid : LinGrid<F, 1>,
        μ : &DiscreteMeasure<Loc<F, 1>, F>,
        filename : String,
    ) {
        let data = grid.into_iter().map(|p@Loc([x]) : Loc<F, 1>| 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<F, 1>,
    > (
        g : &T1,
        grid : LinGrid<F, 1>,
        filename : String,
    ) {
        let data = grid.into_iter().map(|p@Loc([x]) : Loc<F, 1>| 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<F, 2>,
        T2 : RealMapping<F, 2>
    > (
        g0 : Option<&T1>,
        ω0 : Option<&T2>,
        grid : LinGrid<F, 2>,
        μ : &DiscreteMeasure<Loc<F, 2>, F>,
        filename : String,
    ) {
        let data = grid.into_iter().map(|p@Loc([x, y]) : Loc<F, 2>| 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<F, 2>,
    > (
        g : &T1,
        grid : LinGrid<F, 2>,
        filename : String,
    ) {
        let data = grid.into_iter().map(|p@Loc([x, y]) : Loc<F, 2>| CSVHelper2 {
            x,
            y,
            f : g.apply(&p),
        });
        let csv_f = format!("{}.txt", filename);
        write_csv(data, csv_f).expect("CSV save error");
    }

}

/// A helper structure for plotting a sequence of images.
#[derive(Clone,Debug)]
pub struct SeqPlotter<F : Float, const N : usize> {
    /// File name prefix
    prefix : String,
    /// Maximum number of plots to perform
    max_plots : usize,
    /// Sampling grid
    grid : LinGrid<F, N>,
    /// Current plot count
    plot_count : usize,
}

impl<F : Float, const N : usize> SeqPlotter<F, N>
where PlotLookup : Plotting<N> {
    /// Creates a new sequence plotter instance
    pub fn new(prefix : String, max_plots : usize, grid : LinGrid<F, N>) -> 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<T1, T2>(
        &mut self,
        iter : usize,
        g : Option<&T1>,
        ω : Option<&T2>,
        μ : &RNDM<F, N>,
    ) where T1 : RealMapping<F, N>,
            T2 : RealMapping<F, N>
    {
        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;
        }
    }
}