plotters-statistical 0.1.0

//! Correlation-matrix heatmap figure: labeled grid + colorbar, rendered onto a
//! drawing area.

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

use super::colorbar::draw_colorbar;
use super::{contrast_text, text_style};
use crate::colormap::{GradientColorMap, Normalization};
use crate::stats::{correlation_matrix, CorrelationMethod, StatsError};

const MISSING_COLOR: RGBColor = RGBColor(235, 235, 235);

/// A correlation heatmap: a square, labeled, color-mapped grid of a correlation
/// matrix, with an optional colorbar and per-cell value annotations.
///
/// Defaults suit correlations: a diverging `RdBu` colormap on a symmetric
/// `[-1, 1]` normalization (0 → white).
#[derive(Debug, Clone)]
pub struct CorrelationHeatmap {
    matrix: Vec<Vec<f64>>,
    labels: Vec<String>,
    colormap: GradientColorMap,
    norm: Normalization,
    annotate: bool,
    precision: usize,
    show_colorbar: bool,
    gridlines: bool,
    title: Option<String>,
}

impl CorrelationHeatmap {
    /// Compute the correlation matrix of `columns` (one slice per variable, all
    /// the same length) and build a heatmap. `labels` names the variables.
    ///
    /// # Errors
    /// * [`StatsError::LengthMismatch`] if `labels` and `columns` differ in count.
    /// * plus any error from [`correlation_matrix`].
    pub fn from_columns(
        columns: &[Vec<f64>],
        labels: Vec<String>,
        method: CorrelationMethod,
    ) -> Result<Self, StatsError> {
        if labels.len() != columns.len() {
            return Err(StatsError::LengthMismatch {
                scores: columns.len(),
                labels: labels.len(),
            });
        }
        let matrix = correlation_matrix(columns, method)?;
        Ok(Self::from_matrix(matrix, labels))
    }

    /// Build directly from a precomputed square matrix and its labels.
    pub fn from_matrix(matrix: Vec<Vec<f64>>, labels: Vec<String>) -> Self {
        Self {
            matrix,
            labels,
            colormap: GradientColorMap::rd_bu(),
            norm: Normalization::Symmetric {
                center: 0.0,
                half: 1.0,
            },
            annotate: true,
            precision: 2,
            show_colorbar: true,
            gridlines: true,
            title: None,
        }
    }

    /// Set the color map.
    pub fn colormap(mut self, colormap: GradientColorMap) -> Self {
        self.colormap = colormap;
        self
    }

    /// Set the value→color normalization.
    pub fn normalization(mut self, norm: Normalization) -> Self {
        self.norm = norm;
        self
    }

    /// Toggle per-cell value labels.
    pub fn annotate(mut self, on: bool) -> Self {
        self.annotate = on;
        self
    }

    /// Decimal places for cell annotations.
    pub fn precision(mut self, precision: usize) -> Self {
        self.precision = precision;
        self
    }

    /// Toggle the colorbar.
    pub fn colorbar(mut self, on: bool) -> Self {
        self.show_colorbar = on;
        self
    }

    /// Toggle white cell gridlines.
    pub fn gridlines(mut self, on: bool) -> Self {
        self.gridlines = 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
    }

    /// 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 n = self.labels.len();
        if n == 0 || self.matrix.is_empty() {
            return Ok(());
        }
        let (full_w, _) = area.dim_in_pixel();
        let (main, cbar) = if self.show_colorbar {
            let (m, c) = area.split_horizontally(full_w as i32 - 84);
            (m, Some(c))
        } else {
            (area.clone(), None)
        };

        let (w, h) = main.dim_in_pixel();
        let (w, h) = (w as i32, h as i32);
        let top = if self.title.is_some() { 34 } else { 12 };
        let left = 92;
        let bottom = 66;
        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 ni = n 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, &BLACK),
            ))?;
        }

        for i in 0..n {
            for j in 0..n {
                let x0 = gx0 + gw * j as i32 / ni;
                let x1 = gx0 + gw * (j as i32 + 1) / ni;
                let y0 = gy0 + gh * i as i32 / ni;
                let y1 = gy0 + gh * (i as i32 + 1) / ni;
                let v = self.matrix[i][j];
                let color = if v.is_finite() {
                    self.colormap.color(self.norm.t(v))
                } else {
                    MISSING_COLOR
                };
                main.draw(&Rectangle::new([(x0, y0), (x1, y1)], color.filled()))?;
                if self.gridlines {
                    main.draw(&Rectangle::new(
                        [(x0, y0), (x1, y1)],
                        RGBColor(255, 255, 255).stroke_width(1),
                    ))?;
                }
                if self.annotate && v.is_finite() {
                    main.draw(&Text::new(
                        format!("{:.*}", self.precision, v),
                        ((x0 + x1) / 2, (y0 + y1) / 2),
                        text_style(HPos::Center, VPos::Center, 12, &contrast_text(color)),
                    ))?;
                }
            }
        }

        // Row labels (right-aligned to the left of the grid).
        for (i, label) in self.labels.iter().enumerate() {
            let yc = gy0 + gh * (2 * i as i32 + 1) / (2 * ni);
            main.draw(&Text::new(
                label.clone(),
                (gx0 - 6, yc),
                text_style(HPos::Right, VPos::Center, 13, &BLACK),
            ))?;
        }
        // Column labels (below the grid).
        for (j, label) in self.labels.iter().enumerate() {
            let xc = gx0 + gw * (2 * j as i32 + 1) / (2 * ni);
            main.draw(&Text::new(
                label.clone(),
                (xc, gy1 + 8),
                text_style(HPos::Center, VPos::Top, 13, &BLACK),
            ))?;
        }

        if let Some(c) = cbar {
            draw_colorbar(&c, &self.colormap, &self.norm, 5)?;
        }
        Ok(())
    }
}