Skip to main content

plotters_statistical/
colormap.rs

1//! Continuous color maps and value normalization for heatmaps and any other
2//! value-to-color chart.
3//!
4//! A [`GradientColorMap`] linearly interpolates between a small set of RGB color
5//! stops; several standard maps are provided as named constructors, with sources
6//! cited. A [`Normalization`] maps a data value onto the map's `[0, 1]` input.
7//!
8//! Sequential maps (`viridis`, `magma`, `blues`, `reds`) are sampled from
9//! matplotlib / ColorBrewer. The diverging maps (`rd_bu`, `coolwarm`) are meant
10//! for values centered on a midpoint (e.g. correlations around 0) and pair with
11//! [`Normalization::Symmetric`].
12
13use plotters::style::RGBColor;
14
15/// A color map built from sorted RGB stops, interpolated linearly in RGB space.
16#[derive(Debug, Clone)]
17pub struct GradientColorMap {
18    /// `(position in [0,1], color)` stops, sorted ascending by position.
19    stops: Vec<(f64, RGBColor)>,
20}
21
22impl GradientColorMap {
23    /// Build from explicit `(position, color)` stops. Positions are clamped to
24    /// `[0, 1]` and sorted; at least one stop is required (an empty list falls
25    /// back to mid-gray).
26    pub fn new(mut stops: Vec<(f64, RGBColor)>) -> Self {
27        if stops.is_empty() {
28            stops.push((0.0, RGBColor(128, 128, 128)));
29        }
30        for s in &mut stops {
31            s.0 = s.0.clamp(0.0, 1.0);
32        }
33        stops.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
34        Self { stops }
35    }
36
37    /// The color at `t` (clamped to `[0, 1]`).
38    pub fn color(&self, t: f64) -> RGBColor {
39        let t = t.clamp(0.0, 1.0);
40        if t <= self.stops[0].0 {
41            return self.stops[0].1;
42        }
43        let last = self.stops.len() - 1;
44        if t >= self.stops[last].0 {
45            return self.stops[last].1;
46        }
47        // Find the bracketing stops.
48        let mut hi = 1;
49        while hi < self.stops.len() && self.stops[hi].0 < t {
50            hi += 1;
51        }
52        let (p0, c0) = self.stops[hi - 1];
53        let (p1, c1) = self.stops[hi];
54        let f = if p1 > p0 { (t - p0) / (p1 - p0) } else { 0.0 };
55        RGBColor(
56            lerp_u8(c0.0, c1.0, f),
57            lerp_u8(c0.1, c1.1, f),
58            lerp_u8(c0.2, c1.2, f),
59        )
60    }
61
62    /// The `viridis` perceptually-uniform sequential map (matplotlib).
63    pub fn viridis() -> Self {
64        Self::new(vec![
65            (0.0, RGBColor(68, 1, 84)),
66            (0.125, RGBColor(72, 40, 120)),
67            (0.25, RGBColor(62, 74, 137)),
68            (0.375, RGBColor(49, 104, 142)),
69            (0.5, RGBColor(38, 130, 142)),
70            (0.625, RGBColor(31, 158, 137)),
71            (0.75, RGBColor(53, 183, 121)),
72            (0.875, RGBColor(110, 206, 88)),
73            (1.0, RGBColor(253, 231, 37)),
74        ])
75    }
76
77    /// The `magma` sequential map (matplotlib).
78    pub fn magma() -> Self {
79        Self::new(vec![
80            (0.0, RGBColor(0, 0, 4)),
81            (0.25, RGBColor(81, 18, 124)),
82            (0.5, RGBColor(183, 55, 121)),
83            (0.75, RGBColor(252, 137, 97)),
84            (1.0, RGBColor(252, 253, 191)),
85        ])
86    }
87
88    /// ColorBrewer sequential `Blues`.
89    pub fn blues() -> Self {
90        Self::new(vec![
91            (0.0, RGBColor(247, 251, 255)),
92            (0.25, RGBColor(198, 219, 239)),
93            (0.5, RGBColor(107, 174, 214)),
94            (0.75, RGBColor(33, 113, 181)),
95            (1.0, RGBColor(8, 48, 107)),
96        ])
97    }
98
99    /// ColorBrewer sequential `Reds`.
100    pub fn reds() -> Self {
101        Self::new(vec![
102            (0.0, RGBColor(255, 245, 240)),
103            (0.25, RGBColor(252, 187, 161)),
104            (0.5, RGBColor(251, 106, 74)),
105            (0.75, RGBColor(203, 24, 29)),
106            (1.0, RGBColor(103, 0, 13)),
107        ])
108    }
109
110    /// ColorBrewer diverging `RdBu` (red → white → blue). Pairs with
111    /// [`Normalization::Symmetric`] for correlation matrices.
112    pub fn rd_bu() -> Self {
113        Self::new(vec![
114            (0.0, RGBColor(178, 24, 43)),
115            (0.25, RGBColor(239, 138, 98)),
116            (0.5, RGBColor(247, 247, 247)),
117            (0.75, RGBColor(103, 169, 207)),
118            (1.0, RGBColor(33, 102, 172)),
119        ])
120    }
121
122    /// A blue → light → red diverging map (matplotlib `coolwarm` style).
123    pub fn coolwarm() -> Self {
124        Self::new(vec![
125            (0.0, RGBColor(59, 76, 192)),
126            (0.5, RGBColor(221, 221, 221)),
127            (1.0, RGBColor(180, 4, 38)),
128        ])
129    }
130
131    /// Simple white → black grayscale.
132    pub fn grayscale() -> Self {
133        Self::new(vec![
134            (0.0, RGBColor(255, 255, 255)),
135            (1.0, RGBColor(0, 0, 0)),
136        ])
137    }
138}
139
140impl Default for GradientColorMap {
141    fn default() -> Self {
142        Self::viridis()
143    }
144}
145
146fn lerp_u8(a: u8, b: u8, f: f64) -> u8 {
147    (a as f64 + (b as f64 - a as f64) * f)
148        .round()
149        .clamp(0.0, 255.0) as u8
150}
151
152/// Maps a data value onto the `[0, 1]` input of a [`GradientColorMap`].
153#[derive(Debug, Clone, Copy, PartialEq)]
154pub enum Normalization {
155    /// Linear map: `min -> 0`, `max -> 1`.
156    Linear {
157        /// Value mapped to 0.
158        min: f64,
159        /// Value mapped to 1.
160        max: f64,
161    },
162    /// Symmetric (diverging) map centered on `center`: `center -> 0.5`, and
163    /// `center ± half` -> `0` / `1`. Ideal for values around zero.
164    Symmetric {
165        /// The value placed at the midpoint (0.5).
166        center: f64,
167        /// Half-range: `center + half` maps to 1, `center - half` to 0.
168        half: f64,
169    },
170}
171
172impl Normalization {
173    /// Normalize `v` to `[0, 1]` (clamped). Non-finite `v` returns `NaN`, which
174    /// callers can treat as "no data".
175    pub fn t(&self, v: f64) -> f64 {
176        if !v.is_finite() {
177            return f64::NAN;
178        }
179        match *self {
180            Normalization::Linear { min, max } => {
181                if max > min {
182                    ((v - min) / (max - min)).clamp(0.0, 1.0)
183                } else {
184                    0.5
185                }
186            }
187            Normalization::Symmetric { center, half } => {
188                if half > 0.0 {
189                    (0.5 + (v - center) / (2.0 * half)).clamp(0.0, 1.0)
190                } else {
191                    0.5
192                }
193            }
194        }
195    }
196
197    /// Inverse of [`Normalization::t`]: the data value at position `t` in
198    /// `[0, 1]`. Used to label a colorbar.
199    pub fn value(&self, t: f64) -> f64 {
200        match *self {
201            Normalization::Linear { min, max } => min + t * (max - min),
202            Normalization::Symmetric { center, half } => center + (2.0 * t - 1.0) * half,
203        }
204    }
205}