plotters-statistical 0.1.0

//! Missingness heatmap figure: a present/absent map across observations (rows)
//! and variables (columns), for EDA data-quality checks. Column labels carry the
//! per-variable missing percentage.

use plotters::coord::Shift;
use plotters::prelude::*;
use plotters::style::text_anchor::{HPos, VPos};

use super::text_style;
use crate::stats::StatsError;

const PRESENT_COLOR: RGBColor = RGBColor(38, 70, 108); // dark blue
const MISSING_COLOR: RGBColor = RGBColor(232, 232, 232); // light gray

/// A present/absent map: one column per variable, observations down the rows.
#[derive(Debug, Clone)]
pub struct MissingnessHeatmap {
    // present[row][col] — true if the value is present.
    present: Vec<Vec<bool>>,
    labels: Vec<String>,
    nrows: usize,
    present_color: RGBColor,
    missing_color: RGBColor,
    show_percent: bool,
    title: Option<String>,
}

impl MissingnessHeatmap {
    /// Build from `columns` of `Option<f64>` (one column per variable; `None` =
    /// missing). Columns may differ in length; shorter columns are treated as
    /// missing past their end. `labels` names the variables.
    ///
    /// # Errors
    /// * [`StatsError::EmptyInput`] if there are no columns.
    /// * [`StatsError::LengthMismatch`] if `labels` and `columns` differ in count.
    pub fn from_columns(
        columns: &[Vec<Option<f64>>],
        labels: Vec<String>,
    ) -> Result<Self, StatsError> {
        if columns.is_empty() {
            return Err(StatsError::EmptyInput);
        }
        if labels.len() != columns.len() {
            return Err(StatsError::LengthMismatch {
                scores: columns.len(),
                labels: labels.len(),
            });
        }
        let nrows = columns.iter().map(|c| c.len()).max().unwrap_or(0);
        let ncols = columns.len();
        let mut present = vec![vec![false; ncols]; nrows];
        for (j, col) in columns.iter().enumerate() {
            for (i, v) in col.iter().enumerate() {
                present[i][j] = v.is_some();
            }
        }
        Ok(Self {
            present,
            labels,
            nrows,
            present_color: PRESENT_COLOR,
            missing_color: MISSING_COLOR,
            show_percent: true,
            title: None,
        })
    }

    /// Set the color used for present values.
    pub fn present_color(mut self, color: RGBColor) -> Self {
        self.present_color = color;
        self
    }

    /// Set the color used for missing values.
    pub fn missing_color(mut self, color: RGBColor) -> Self {
        self.missing_color = color;
        self
    }

    /// Toggle the per-column missing-percentage labels.
    pub fn show_percent(mut self, on: bool) -> Self {
        self.show_percent = on;
        self
    }

    /// Set a title drawn above the grid.
    pub fn title(mut self, title: impl Into<String>) -> Self {
        self.title = Some(title.into());
        self
    }

    fn missing_fraction(&self, col: usize) -> f64 {
        if self.nrows == 0 {
            return 0.0;
        }
        let missing = self.present.iter().filter(|row| !row[col]).count();
        missing as f64 / self.nrows as f64
    }

    /// Render the heatmap onto `area`.
    pub fn draw<DB: DrawingBackend>(
        &self,
        area: &DrawingArea<DB, Shift>,
    ) -> Result<(), Box<dyn std::error::Error>>
    where
        DB::ErrorType: 'static,
    {
        let ncols = self.labels.len();
        if ncols == 0 || self.nrows == 0 {
            return Ok(());
        }
        let (w, h) = area.dim_in_pixel();
        let (w, h) = (w as i32, h as i32);
        let top = if self.title.is_some() { 34 } else { 12 };
        let left = 12;
        let bottom = 70;
        let gx0 = left;
        let gy0 = top;
        let gx1 = (w - 12).max(gx0 + 1);
        let gy1 = (h - bottom).max(gy0 + 1);
        let gw = gx1 - gx0;
        let gh = gy1 - gy0;
        let nc = ncols as i32;

        if let Some(t) = &self.title {
            area.draw(&Text::new(
                t.clone(),
                ((gx0 + gx1) / 2, 8),
                text_style(HPos::Center, VPos::Top, 18, &RGBColor(0, 0, 0)),
            ))?;
        }

        // Render by pixel row: map each screen row to a data observation. This
        // bounds work at ~gh*ncols regardless of how many observations there are.
        for py in 0..gh {
            let row = (py as usize * self.nrows) / gh as usize;
            let y0 = gy0 + py;
            for j in 0..ncols {
                let x0 = gx0 + gw * j as i32 / nc;
                let x1 = gx0 + gw * (j as i32 + 1) / nc;
                let color = if self.present[row][j] {
                    self.present_color
                } else {
                    self.missing_color
                };
                area.draw(&Rectangle::new([(x0, y0), (x1, y0 + 1)], color.filled()))?;
            }
        }

        // Column labels + missing percentage.
        for (j, label) in self.labels.iter().enumerate() {
            let xc = gx0 + gw * (2 * j as i32 + 1) / (2 * nc);
            area.draw(&Text::new(
                label.clone(),
                (xc, gy1 + 8),
                text_style(HPos::Center, VPos::Top, 12, &RGBColor(0, 0, 0)),
            ))?;
            if self.show_percent {
                area.draw(&Text::new(
                    format!("{:.0}%", 100.0 * self.missing_fraction(j)),
                    (xc, gy1 + 26),
                    text_style(HPos::Center, VPos::Top, 11, &RGBColor(120, 120, 120)),
                ))?;
            }
        }
        Ok(())
    }
}