plotters-statistical 0.1.0

//! Generic heatmap series — a grid of color-mapped cells with optional numeric
//! annotations. The building block under the correlation / missingness heatmap
//! figures, but usable directly on any chart.
//!
//! Cells are placed at integer coordinates: column `c` at `x = c`, row `r` at
//! `y = nrows - 1 - r` (so row 0 is at the top). Build the chart with ranges
//! `-0.5..ncols as f64 - 0.5` and `-0.5..nrows as f64 - 0.5`.

use plotters::element::{Drawable, PointCollection};
use plotters::style::text_anchor::{HPos, Pos, VPos};
use plotters::style::{IntoFont, RGBColor, ShapeStyle, TextStyle};
use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};

use crate::colormap::{GradientColorMap, Normalization};
use crate::style::fill_style;

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

/// Per-cell numeric annotation settings.
#[derive(Debug, Clone)]
pub struct HeatmapAnnotation {
    /// Decimal places shown.
    pub precision: usize,
    /// Font size in points.
    pub font_size: u32,
    /// Text color; `None` picks black or white automatically for contrast
    /// against each cell.
    pub text_color: Option<RGBColor>,
}

impl Default for HeatmapAnnotation {
    fn default() -> Self {
        Self {
            precision: 2,
            font_size: 12,
            text_color: None,
        }
    }
}

/// A grid of color-mapped cells.
#[derive(Debug, Clone)]
pub struct Heatmap {
    // Two corners per cell (row-major): [lower-left, upper-right] in data coords.
    corners: Vec<(f64, f64)>,
    values: Vec<f64>,
    colormap: GradientColorMap,
    norm: Normalization,
    annotate: Option<HeatmapAnnotation>,
    border: Option<ShapeStyle>,
    cell_gap: i32,
}

impl Heatmap {
    /// Build from a row-major matrix (`values[row][col]`). Rows may not be empty
    /// and must all be the same width; otherwise the shorter rows are padded
    /// with `NaN` (rendered as the "missing" color). The default normalization
    /// is linear over the finite value range.
    pub fn new(values: &[Vec<f64>]) -> Self {
        let nrows = values.len();
        let ncols = values.iter().map(|r| r.len()).max().unwrap_or(0);
        let mut flat = Vec::with_capacity(nrows * ncols);
        let mut corners = Vec::with_capacity(nrows * ncols * 2);
        for (r, row) in values.iter().enumerate() {
            let y = (nrows - 1 - r) as f64;
            for c in 0..ncols {
                let v = row.get(c).copied().unwrap_or(f64::NAN);
                flat.push(v);
                corners.push((c as f64 - 0.5, y - 0.5));
                corners.push((c as f64 + 0.5, y + 0.5));
            }
        }
        let (mut lo, mut hi) = (f64::INFINITY, f64::NEG_INFINITY);
        for &v in &flat {
            if v.is_finite() {
                lo = lo.min(v);
                hi = hi.max(v);
            }
        }
        if !lo.is_finite() {
            lo = 0.0;
            hi = 1.0;
        }
        Self {
            corners,
            values: flat,
            colormap: GradientColorMap::viridis(),
            norm: Normalization::Linear { min: lo, max: hi },
            annotate: None,
            border: None,
            cell_gap: 0,
        }
    }

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

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

    /// Annotate each cell with its value.
    pub fn annotate(mut self, annotation: HeatmapAnnotation) -> Self {
        self.annotate = Some(annotation);
        self
    }

    /// Draw a border around each cell.
    pub fn cell_border(mut self, style: ShapeStyle) -> Self {
        self.border = Some(style);
        self
    }

    /// Shrink each cell by `gap` pixels on every side (visual gutters).
    pub fn cell_gap(mut self, gap: i32) -> Self {
        self.cell_gap = gap;
        self
    }
}

fn contrast_text(bg: RGBColor) -> RGBColor {
    let l = 0.299 * bg.0 as f64 + 0.587 * bg.1 as f64 + 0.114 * bg.2 as f64;
    if l > 140.0 {
        RGBColor(0, 0, 0)
    } else {
        RGBColor(255, 255, 255)
    }
}

impl<'a> PointCollection<'a, (f64, f64)> for &'a Heatmap {
    type Point = &'a (f64, f64);
    type IntoIter = &'a [(f64, f64)];
    fn point_iter(self) -> &'a [(f64, f64)] {
        &self.corners
    }
}

impl<DB: DrawingBackend> Drawable<DB> for Heatmap {
    fn draw<I: Iterator<Item = BackendCoord>>(
        &self,
        points: I,
        backend: &mut DB,
        _parent_dim: (u32, u32),
    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
        let pix: Vec<BackendCoord> = points.collect();
        let ncells = self.values.len();
        if pix.len() < ncells * 2 {
            return Ok(());
        }
        for k in 0..ncells {
            let a = pix[2 * k];
            let b = pix[2 * k + 1];
            let (x0, x1) = (a.0.min(b.0) + self.cell_gap, a.0.max(b.0) - self.cell_gap);
            let (y0, y1) = (a.1.min(b.1) + self.cell_gap, a.1.max(b.1) - self.cell_gap);
            if x1 <= x0 || y1 <= y0 {
                continue;
            }
            let v = self.values[k];
            let color = if v.is_finite() {
                self.colormap.color(self.norm.t(v))
            } else {
                MISSING_COLOR
            };
            backend.draw_rect((x0, y0), (x1, y1), &fill_style(&color), true)?;
            if let Some(border) = &self.border {
                backend.draw_rect((x0, y0), (x1, y1), border, false)?;
            }
            if let Some(ann) = &self.annotate {
                if v.is_finite() {
                    let tc = ann.text_color.unwrap_or_else(|| contrast_text(color));
                    let style = TextStyle::from(("sans-serif", ann.font_size as i32).into_font())
                        .color(&tc)
                        .pos(Pos::new(HPos::Center, VPos::Center));
                    let cx = (x0 + x1) / 2;
                    let cy = (y0 + y1) / 2;
                    let text = format!("{:.*}", ann.precision, v);
                    backend.draw_text(&text, &style, (cx, cy))?;
                }
            }
        }
        Ok(())
    }
}