plotters-statistical 0.1.0

//! Continuous color maps and value normalization for heatmaps and any other
//! value-to-color chart.
//!
//! A [`GradientColorMap`] linearly interpolates between a small set of RGB color
//! stops; several standard maps are provided as named constructors, with sources
//! cited. A [`Normalization`] maps a data value onto the map's `[0, 1]` input.
//!
//! Sequential maps (`viridis`, `magma`, `blues`, `reds`) are sampled from
//! matplotlib / ColorBrewer. The diverging maps (`rd_bu`, `coolwarm`) are meant
//! for values centered on a midpoint (e.g. correlations around 0) and pair with
//! [`Normalization::Symmetric`].

use plotters::style::RGBColor;

/// A color map built from sorted RGB stops, interpolated linearly in RGB space.
#[derive(Debug, Clone)]
pub struct GradientColorMap {
    /// `(position in [0,1], color)` stops, sorted ascending by position.
    stops: Vec<(f64, RGBColor)>,
}

impl GradientColorMap {
    /// Build from explicit `(position, color)` stops. Positions are clamped to
    /// `[0, 1]` and sorted; at least one stop is required (an empty list falls
    /// back to mid-gray).
    pub fn new(mut stops: Vec<(f64, RGBColor)>) -> Self {
        if stops.is_empty() {
            stops.push((0.0, RGBColor(128, 128, 128)));
        }
        for s in &mut stops {
            s.0 = s.0.clamp(0.0, 1.0);
        }
        stops.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
        Self { stops }
    }

    /// The color at `t` (clamped to `[0, 1]`).
    pub fn color(&self, t: f64) -> RGBColor {
        let t = t.clamp(0.0, 1.0);
        if t <= self.stops[0].0 {
            return self.stops[0].1;
        }
        let last = self.stops.len() - 1;
        if t >= self.stops[last].0 {
            return self.stops[last].1;
        }
        // Find the bracketing stops.
        let mut hi = 1;
        while hi < self.stops.len() && self.stops[hi].0 < t {
            hi += 1;
        }
        let (p0, c0) = self.stops[hi - 1];
        let (p1, c1) = self.stops[hi];
        let f = if p1 > p0 { (t - p0) / (p1 - p0) } else { 0.0 };
        RGBColor(
            lerp_u8(c0.0, c1.0, f),
            lerp_u8(c0.1, c1.1, f),
            lerp_u8(c0.2, c1.2, f),
        )
    }

    /// The `viridis` perceptually-uniform sequential map (matplotlib).
    pub fn viridis() -> Self {
        Self::new(vec![
            (0.0, RGBColor(68, 1, 84)),
            (0.125, RGBColor(72, 40, 120)),
            (0.25, RGBColor(62, 74, 137)),
            (0.375, RGBColor(49, 104, 142)),
            (0.5, RGBColor(38, 130, 142)),
            (0.625, RGBColor(31, 158, 137)),
            (0.75, RGBColor(53, 183, 121)),
            (0.875, RGBColor(110, 206, 88)),
            (1.0, RGBColor(253, 231, 37)),
        ])
    }

    /// The `magma` sequential map (matplotlib).
    pub fn magma() -> Self {
        Self::new(vec![
            (0.0, RGBColor(0, 0, 4)),
            (0.25, RGBColor(81, 18, 124)),
            (0.5, RGBColor(183, 55, 121)),
            (0.75, RGBColor(252, 137, 97)),
            (1.0, RGBColor(252, 253, 191)),
        ])
    }

    /// ColorBrewer sequential `Blues`.
    pub fn blues() -> Self {
        Self::new(vec![
            (0.0, RGBColor(247, 251, 255)),
            (0.25, RGBColor(198, 219, 239)),
            (0.5, RGBColor(107, 174, 214)),
            (0.75, RGBColor(33, 113, 181)),
            (1.0, RGBColor(8, 48, 107)),
        ])
    }

    /// ColorBrewer sequential `Reds`.
    pub fn reds() -> Self {
        Self::new(vec![
            (0.0, RGBColor(255, 245, 240)),
            (0.25, RGBColor(252, 187, 161)),
            (0.5, RGBColor(251, 106, 74)),
            (0.75, RGBColor(203, 24, 29)),
            (1.0, RGBColor(103, 0, 13)),
        ])
    }

    /// ColorBrewer diverging `RdBu` (red → white → blue). Pairs with
    /// [`Normalization::Symmetric`] for correlation matrices.
    pub fn rd_bu() -> Self {
        Self::new(vec![
            (0.0, RGBColor(178, 24, 43)),
            (0.25, RGBColor(239, 138, 98)),
            (0.5, RGBColor(247, 247, 247)),
            (0.75, RGBColor(103, 169, 207)),
            (1.0, RGBColor(33, 102, 172)),
        ])
    }

    /// A blue → light → red diverging map (matplotlib `coolwarm` style).
    pub fn coolwarm() -> Self {
        Self::new(vec![
            (0.0, RGBColor(59, 76, 192)),
            (0.5, RGBColor(221, 221, 221)),
            (1.0, RGBColor(180, 4, 38)),
        ])
    }

    /// Simple white → black grayscale.
    pub fn grayscale() -> Self {
        Self::new(vec![
            (0.0, RGBColor(255, 255, 255)),
            (1.0, RGBColor(0, 0, 0)),
        ])
    }
}

impl Default for GradientColorMap {
    fn default() -> Self {
        Self::viridis()
    }
}

fn lerp_u8(a: u8, b: u8, f: f64) -> u8 {
    (a as f64 + (b as f64 - a as f64) * f)
        .round()
        .clamp(0.0, 255.0) as u8
}

/// Maps a data value onto the `[0, 1]` input of a [`GradientColorMap`].
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Normalization {
    /// Linear map: `min -> 0`, `max -> 1`.
    Linear {
        /// Value mapped to 0.
        min: f64,
        /// Value mapped to 1.
        max: f64,
    },
    /// Symmetric (diverging) map centered on `center`: `center -> 0.5`, and
    /// `center ± half` -> `0` / `1`. Ideal for values around zero.
    Symmetric {
        /// The value placed at the midpoint (0.5).
        center: f64,
        /// Half-range: `center + half` maps to 1, `center - half` to 0.
        half: f64,
    },
}

impl Normalization {
    /// Normalize `v` to `[0, 1]` (clamped). Non-finite `v` returns `NaN`, which
    /// callers can treat as "no data".
    pub fn t(&self, v: f64) -> f64 {
        if !v.is_finite() {
            return f64::NAN;
        }
        match *self {
            Normalization::Linear { min, max } => {
                if max > min {
                    ((v - min) / (max - min)).clamp(0.0, 1.0)
                } else {
                    0.5
                }
            }
            Normalization::Symmetric { center, half } => {
                if half > 0.0 {
                    (0.5 + (v - center) / (2.0 * half)).clamp(0.0, 1.0)
                } else {
                    0.5
                }
            }
        }
    }

    /// Inverse of [`Normalization::t`]: the data value at position `t` in
    /// `[0, 1]`. Used to label a colorbar.
    pub fn value(&self, t: f64) -> f64 {
        match *self {
            Normalization::Linear { min, max } => min + t * (max - min),
            Normalization::Symmetric { center, half } => center + (2.0 * t - 1.0) * half,
        }
    }
}