plotters-statistical 0.1.0

//! Pair plot (scatterplot matrix): every variable against every other, with a
//! distribution view on the diagonal. Renders onto a drawing area by splitting
//! it into an n×n grid and reusing this crate's series plus plain `plotters`
//! scatter marks.

use plotters::coord::Shift;
use plotters::prelude::*;

use crate::series::Ecdf;
use crate::stats::{histogram, BinRule};
use crate::style::{palette_color, translucent_fill};

/// What to draw on the diagonal panels (variable vs itself).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Diagonal {
    /// A histogram of the variable.
    Histogram,
    /// The variable's empirical CDF.
    Ecdf,
}

/// A scatterplot matrix.
#[derive(Debug, Clone)]
pub struct PairPlot {
    columns: Vec<Vec<f64>>,
    labels: Vec<String>,
    diagonal: Diagonal,
    marker_radius: u32,
    hue: Option<Vec<usize>>,
    bins: BinRule,
}

impl PairPlot {
    /// Build from `columns` (one slice per variable, all the same length) and
    /// their `labels`. Panics-free: mismatched lengths are simply reflected in
    /// the empty result of [`PairPlot::draw`] (which returns `Ok(())` when there
    /// is nothing coherent to plot).
    pub fn new(columns: Vec<Vec<f64>>, labels: Vec<String>) -> Self {
        Self {
            columns,
            labels,
            diagonal: Diagonal::Histogram,
            marker_radius: 2,
            hue: None,
            bins: BinRule::Sturges,
        }
    }

    /// Choose what the diagonal panels show.
    pub fn diagonal(mut self, diagonal: Diagonal) -> Self {
        self.diagonal = diagonal;
        self
    }

    /// Set the scatter marker radius in pixels.
    pub fn marker_radius(mut self, radius: u32) -> Self {
        self.marker_radius = radius;
        self
    }

    /// Color scatter points by a per-observation group index (its length should
    /// match the column length). Colors cycle the shared palette.
    pub fn hue(mut self, groups: Vec<usize>) -> Self {
        self.hue = Some(groups);
        self
    }

    /// Set the histogram bin rule used on the diagonal.
    pub fn bins(mut self, bins: BinRule) -> Self {
        self.bins = bins;
        self
    }

    /// Render the matrix onto `area`.
    pub fn draw<DB: DrawingBackend>(
        &self,
        area: &DrawingArea<DB, Shift>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        DB::ErrorType: 'static,
    {
        let n = self.columns.len();
        if n == 0 || self.labels.len() != n {
            return Ok(());
        }
        let len = self.columns[0].len();
        if self.columns.iter().any(|c| c.len() != len) || len == 0 {
            return Ok(());
        }

        // Per-column display range with a little padding.
        let ranges: Vec<(f64, f64)> = self.columns.iter().map(|c| padded_range(c)).collect();

        let panels = area.split_evenly((n, n));
        for i in 0..n {
            for j in 0..n {
                let panel = &panels[i * n + j];
                if i == j {
                    self.draw_diagonal(panel, i, ranges[i], i == n - 1, j == 0)?;
                } else {
                    self.draw_scatter(panel, j, i, ranges[j], ranges[i], i == n - 1, j == 0)?;
                }
            }
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn draw_scatter<DB: DrawingBackend>(
        &self,
        panel: &DrawingArea<DB, Shift>,
        xcol: usize,
        ycol: usize,
        xr: (f64, f64),
        yr: (f64, f64),
        bottom_row: bool,
        left_col: bool,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        DB::ErrorType: 'static,
    {
        let mut chart = self.panel_chart(panel, xr, yr, xcol, ycol, bottom_row, left_col)?;
        let r = self.marker_radius;
        let default_fill = translucent_fill(&palette_color(0), 0.55);
        chart.draw_series(
            self.columns[xcol]
                .iter()
                .zip(&self.columns[ycol])
                .enumerate()
                .map(|(k, (&x, &y))| {
                    let style = match &self.hue {
                        Some(g) if k < g.len() => translucent_fill(&palette_color(g[k]), 0.6),
                        _ => default_fill,
                    };
                    Circle::new((x, y), r, style)
                }),
        )?;
        Ok(())
    }

    fn draw_diagonal<DB: DrawingBackend>(
        &self,
        panel: &DrawingArea<DB, Shift>,
        col: usize,
        xr: (f64, f64),
        bottom_row: bool,
        left_col: bool,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        DB::ErrorType: 'static,
    {
        let data = &self.columns[col];
        match self.diagonal {
            Diagonal::Ecdf => {
                let mut chart =
                    self.panel_chart(panel, xr, (0.0, 1.0), col, col, bottom_row, left_col)?;
                if let Ok(e) = Ecdf::from_data(data) {
                    chart.draw_series(std::iter::once(e))?;
                }
            }
            Diagonal::Histogram => {
                if let Ok(h) = histogram(data, self.bins) {
                    let ymax = h.counts.iter().copied().max().unwrap_or(1).max(1) as f64;
                    let mut chart = self.panel_chart(
                        panel,
                        xr,
                        (0.0, ymax * 1.05),
                        col,
                        col,
                        bottom_row,
                        left_col,
                    )?;
                    let fill = translucent_fill(&palette_color(0), 0.6);
                    chart.draw_series(
                        h.edges
                            .windows(2)
                            .zip(&h.counts)
                            .map(|(e, &c)| Rectangle::new([(e[0], 0.0), (e[1], c as f64)], fill)),
                    )?;
                }
            }
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn panel_chart<'b, DB: DrawingBackend>(
        &self,
        panel: &'b DrawingArea<DB, Shift>,
        xr: (f64, f64),
        yr: (f64, f64),
        xcol: usize,
        ycol: usize,
        bottom_row: bool,
        left_col: bool,
    ) -> Result<
        ChartContext<
            'b,
            DB,
            Cartesian2d<
                plotters::coord::types::RangedCoordf64,
                plotters::coord::types::RangedCoordf64,
            >,
        >,
        Box<dyn std::error::Error>,
    >
    where
        DB::ErrorType: 'static,
    {
        let mut builder = ChartBuilder::on(panel);
        builder.margin(3);
        if left_col {
            builder.set_label_area_size(LabelAreaPosition::Left, 38);
        }
        if bottom_row {
            builder.set_label_area_size(LabelAreaPosition::Bottom, 26);
        }
        let mut chart = builder.build_cartesian_2d(xr.0..xr.1, yr.0..yr.1)?;
        let x_desc = if bottom_row {
            self.labels[xcol].clone()
        } else {
            String::new()
        };
        let y_desc = if left_col {
            self.labels[ycol].clone()
        } else {
            String::new()
        };
        chart
            .configure_mesh()
            .disable_mesh()
            .x_labels(if bottom_row { 4 } else { 0 })
            .y_labels(if left_col { 4 } else { 0 })
            .x_desc(x_desc)
            .y_desc(y_desc)
            .label_style(("sans-serif", 11))
            .draw()?;
        Ok(chart)
    }
}

fn padded_range(data: &[f64]) -> (f64, f64) {
    let mut lo = f64::INFINITY;
    let mut hi = f64::NEG_INFINITY;
    for &v in data {
        if v.is_finite() {
            lo = lo.min(v);
            hi = hi.max(v);
        }
    }
    if !lo.is_finite() {
        return (-1.0, 1.0);
    }
    if hi <= lo {
        return (lo - 0.5, lo + 0.5);
    }
    let pad = (hi - lo) * 0.05;
    (lo - pad, hi + pad)
}