Skip to main content

ggplot_rs/scale/
gradient_n.rs

1use crate::aes::Aesthetic;
2use crate::data::Value;
3
4use super::color::RGBAColor;
5use super::util::nice_step;
6use super::Scale;
7
8/// N-stop continuous color gradient scale.
9/// Interpolates linearly between user-defined color stops.
10#[derive(Clone, Debug)]
11pub struct ScaleColorGradientN {
12    aesthetic: Aesthetic,
13    name: String,
14    /// Color stops as (position_0_to_1, color) pairs, sorted by position.
15    stops: Vec<(f64, RGBAColor)>,
16    min: f64,
17    max: f64,
18    /// Locale for the default break labels (`None` = default marks).
19    locale: Option<crate::locale::Locale>,
20    /// Fixed `(lo, hi)` domain ([`with_limits`](Self::with_limits)):
21    /// `train` no longer widens it, and `map_to_color` maps against it
22    /// instead of the trained min/max.
23    limits: Option<(f64, f64)>,
24    /// Out-of-bounds policy for `limits` ([`with_oob`](Self::with_oob)).
25    /// Default [`Oob::Censor`](super::Oob::Censor).
26    oob: super::Oob,
27    /// Colour for a genuinely missing or out-of-bounds-and-censored value
28    /// ([`with_na_value`](Self::with_na_value)). Default ggplot2's
29    /// `"grey50"`, `(127, 127, 127)`.
30    na_value: (u8, u8, u8),
31}
32
33impl ScaleColorGradientN {
34    /// Create a new N-stop gradient for the given aesthetic.
35    /// Stops are `(position, color)` where position is in [0, 1].
36    pub fn new(aesthetic: Aesthetic, stops: Vec<(f64, RGBAColor)>) -> Self {
37        let mut stops = stops;
38        stops.sort_by(|a, b| a.0.total_cmp(&b.0));
39        ScaleColorGradientN {
40            aesthetic,
41            name: String::new(),
42            stops,
43            min: f64::INFINITY,
44            max: f64::NEG_INFINITY,
45            locale: None,
46            limits: None,
47            oob: super::Oob::Censor,
48            na_value: super::DEFAULT_NA_VALUE,
49        }
50    }
51
52    /// Fix the colour domain to `[lo, hi]` instead of whatever the data
53    /// happens to span. Out-of-range values follow
54    /// [`with_oob`](Self::with_oob) (default: `na_value`).
55    pub fn with_limits(mut self, lo: f64, hi: f64) -> Self {
56        self.limits = Some((lo, hi));
57        self
58    }
59
60    /// Out-of-bounds policy for [`with_limits`](Self::with_limits):
61    /// [`Oob::Censor`](super::Oob::Censor) (default) or
62    /// [`Oob::Squish`](super::Oob::Squish).
63    pub fn with_oob(mut self, oob: super::Oob) -> Self {
64        self.oob = oob;
65        self
66    }
67
68    /// Colour for a genuinely missing (`NA`) value, or one censored by
69    /// `with_limits`'s default [`Oob::Censor`]. Default ggplot2's
70    /// `"grey50"`, `(127, 127, 127)`.
71    pub fn with_na_value(mut self, color: (u8, u8, u8)) -> Self {
72        self.na_value = color;
73        self
74    }
75
76    /// Create a continuous viridis palette.
77    pub fn viridis(aesthetic: Aesthetic) -> Self {
78        Self::new(aesthetic, viridis_stops())
79    }
80
81    /// Create a continuous magma palette.
82    pub fn magma(aesthetic: Aesthetic) -> Self {
83        Self::new(aesthetic, magma_stops())
84    }
85
86    /// Create a continuous plasma palette.
87    pub fn plasma(aesthetic: Aesthetic) -> Self {
88        Self::new(aesthetic, plasma_stops())
89    }
90
91    /// Create a continuous inferno palette.
92    pub fn inferno(aesthetic: Aesthetic) -> Self {
93        Self::new(aesthetic, inferno_stops())
94    }
95
96    /// Interpolate the color at a normalized position t in [0, 1].
97    fn color_at(&self, t: f64) -> RGBAColor {
98        let t = t.clamp(0.0, 1.0);
99        if self.stops.is_empty() {
100            return RGBAColor::new(127, 127, 127);
101        }
102        if self.stops.len() == 1 {
103            return self.stops[0].1;
104        }
105        // Find the two surrounding stops
106        if t <= self.stops[0].0 {
107            return self.stops[0].1;
108        }
109        if t >= self.stops[self.stops.len() - 1].0 {
110            return self.stops[self.stops.len() - 1].1;
111        }
112        for i in 0..self.stops.len() - 1 {
113            let (p0, c0) = &self.stops[i];
114            let (p1, c1) = &self.stops[i + 1];
115            if t >= *p0 && t <= *p1 {
116                let range = p1 - p0;
117                let local_t = if range.abs() < f64::EPSILON {
118                    0.0
119                } else {
120                    (t - p0) / range
121                };
122                return c0.lerp(c1, local_t);
123            }
124        }
125        self.stops[self.stops.len() - 1].1
126    }
127}
128
129impl Scale for ScaleColorGradientN {
130    fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
131        if self.locale.is_none() {
132            self.locale = Some(locale.clone());
133        }
134    }
135
136    fn locale(&self) -> Option<&crate::locale::Locale> {
137        self.locale.as_ref()
138    }
139
140    fn format_value(&self, value: &Value) -> Option<String> {
141        let f = value.as_f64()?;
142        if !f.is_finite() {
143            return None;
144        }
145        Some(super::util::format_number_opt(f, self.locale.as_ref()))
146    }
147
148    fn aesthetic(&self) -> Aesthetic {
149        self.aesthetic.clone()
150    }
151
152    fn train(&mut self, values: &[Value]) {
153        if self.limits.is_some() {
154            return;
155        }
156        for v in values {
157            if let Some(f) = v.as_f64() {
158                if f.is_finite() {
159                    if f < self.min {
160                        self.min = f;
161                    }
162                    if f > self.max {
163                        self.max = f;
164                    }
165                }
166            }
167        }
168    }
169
170    fn map(&self, value: &Value) -> f64 {
171        let f = match value.as_f64() {
172            Some(f) => f,
173            None => return 0.0,
174        };
175        let (lo, hi) = self.limits.unwrap_or((self.min, self.max));
176        let range = hi - lo;
177        if range.abs() < f64::EPSILON {
178            0.5
179        } else {
180            (f - lo) / range
181        }
182    }
183
184    fn breaks(&self) -> Vec<(f64, String)> {
185        let (lo, hi) = self.limits.unwrap_or((self.min, self.max));
186        if lo > hi || !lo.is_finite() || !hi.is_finite() {
187            return vec![];
188        }
189        let range = hi - lo;
190        if super::util::is_degenerate_range(lo, hi) {
191            return vec![(
192                0.5,
193                super::util::format_number_opt(lo, self.locale.as_ref()),
194            )];
195        }
196        let n_breaks = 5;
197        let raw_step = range / n_breaks as f64;
198        let step = nice_step(raw_step);
199        let start = (lo / step).ceil() * step;
200        let values = super::util::stepped_breaks(start, hi, step);
201        let labels = super::util::format_numbers_opt(&values, self.locale.as_ref());
202        values
203            .iter()
204            .zip(labels)
205            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
206            .collect()
207    }
208
209    fn name(&self) -> &str {
210        &self.name
211    }
212
213    fn set_name(&mut self, name: &str) {
214        self.name = name.to_string();
215    }
216
217    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
218        // A genuinely missing or non-finite value, and a value censored
219        // by `with_limits`, render as `na_value` (ggplot2's own
220        // behaviour) instead of a plausible-looking colour (review PR#141
221        // finding 3). A non-numeric-but-present value (e.g. the wrong
222        // column type) keeps the pre-existing contract of falling through
223        // to `map`'s own 0.0 fallback instead — this type's `map_to_color`
224        // always returned `Some` for it before `with_limits`/`Oob` existed.
225        if value.is_na() {
226            return Some(self.na_value);
227        }
228        if let Some(f) = value.as_f64() {
229            if !f.is_finite() {
230                return Some(self.na_value);
231            }
232            if super::resolve_oob(f, self.limits, self.oob).is_none() {
233                return Some(self.na_value);
234            }
235        }
236        let t = self.map(value);
237        let c = self.color_at(t);
238        Some((c.r, c.g, c.b))
239    }
240
241    fn domain(&self) -> Option<(f64, f64)> {
242        if let Some((lo, hi)) = self.limits {
243            return Some((lo, hi));
244        }
245        if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
246            Some((self.min, self.max))
247        } else {
248            None
249        }
250    }
251
252    fn clone_box(&self) -> Box<dyn Scale> {
253        Box::new(self.clone())
254    }
255
256    fn reset_training(&mut self) {
257        self.min = f64::INFINITY;
258        self.max = f64::NEG_INFINITY;
259    }
260
261    fn validate(&self, warnings: &mut Vec<String>) {
262        if let Some((lo, hi)) = self.limits {
263            if lo > hi {
264                warnings.push(format!(
265                    "scale_color/fill_gradientn: with_limits({lo}, {hi}) has lo > hi; every value maps to na_value"
266                ));
267            }
268        }
269    }
270}
271
272// ─── Continuous palette color stops ──────────────────────────────
273
274fn c(r: u8, g: u8, b: u8) -> RGBAColor {
275    RGBAColor::new(r, g, b)
276}
277
278fn viridis_stops() -> Vec<(f64, RGBAColor)> {
279    let colors = [
280        c(68, 1, 84),
281        c(72, 26, 108),
282        c(71, 47, 126),
283        c(65, 68, 135),
284        c(57, 86, 140),
285        c(47, 104, 142),
286        c(38, 121, 142),
287        c(31, 138, 141),
288        c(30, 155, 138),
289        c(42, 172, 130),
290        c(70, 188, 115),
291        c(109, 202, 93),
292        c(155, 213, 67),
293        c(200, 222, 39),
294        c(240, 229, 30),
295        c(253, 231, 37),
296    ];
297    evenly_spaced_stops(&colors)
298}
299
300fn magma_stops() -> Vec<(f64, RGBAColor)> {
301    let colors = [
302        c(0, 0, 4),
303        c(16, 12, 50),
304        c(41, 17, 90),
305        c(72, 12, 110),
306        c(101, 19, 110),
307        c(131, 29, 103),
308        c(160, 42, 93),
309        c(187, 55, 84),
310        c(213, 72, 72),
311        c(232, 99, 62),
312        c(247, 131, 57),
313        c(254, 167, 69),
314        c(254, 203, 99),
315        c(252, 235, 141),
316        c(252, 254, 188),
317        c(252, 253, 191),
318    ];
319    evenly_spaced_stops(&colors)
320}
321
322fn plasma_stops() -> Vec<(f64, RGBAColor)> {
323    let colors = [
324        c(13, 8, 135),
325        c(53, 5, 157),
326        c(82, 1, 163),
327        c(109, 1, 159),
328        c(133, 7, 147),
329        c(156, 23, 127),
330        c(175, 42, 106),
331        c(192, 61, 85),
332        c(206, 82, 66),
333        c(218, 105, 46),
334        c(228, 130, 24),
335        c(236, 157, 6),
336        c(240, 185, 11),
337        c(239, 213, 38),
338        c(232, 240, 73),
339        c(240, 249, 33),
340    ];
341    evenly_spaced_stops(&colors)
342}
343
344fn inferno_stops() -> Vec<(f64, RGBAColor)> {
345    let colors = [
346        c(0, 0, 4),
347        c(14, 11, 49),
348        c(39, 15, 90),
349        c(67, 10, 107),
350        c(95, 13, 106),
351        c(122, 21, 97),
352        c(149, 33, 81),
353        c(174, 49, 60),
354        c(196, 69, 38),
355        c(215, 95, 15),
356        c(231, 124, 3),
357        c(243, 155, 7),
358        c(250, 189, 28),
359        c(252, 222, 67),
360        c(247, 252, 118),
361        c(252, 255, 164),
362    ];
363    evenly_spaced_stops(&colors)
364}
365
366fn evenly_spaced_stops(colors: &[RGBAColor]) -> Vec<(f64, RGBAColor)> {
367    let n = colors.len();
368    if n == 0 {
369        return vec![];
370    }
371    if n == 1 {
372        return vec![(0.0, colors[0])];
373    }
374    colors
375        .iter()
376        .enumerate()
377        .map(|(i, c)| (i as f64 / (n - 1) as f64, *c))
378        .collect()
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384
385    #[test]
386    fn format_value_formats_a_finite_number_and_rejects_non_finite() {
387        let g = ScaleColorGradientN::new(Aesthetic::Color, vec![(0.0, RGBAColor::new(0, 0, 0))]);
388        assert_eq!(g.format_value(&Value::Float(42.0)), Some("42".to_string()));
389        assert_eq!(g.format_value(&Value::Float(f64::INFINITY)), None);
390    }
391
392    #[test]
393    fn with_limits_fixes_the_domain_and_training_does_not_widen_it() {
394        let mut g = ScaleColorGradientN::new(
395            Aesthetic::Color,
396            vec![
397                (0.0, RGBAColor::new(0, 0, 0)),
398                (1.0, RGBAColor::new(255, 255, 255)),
399            ],
400        )
401        .with_limits(0.0, 10.0);
402        g.train(&[Value::Float(-1000.0), Value::Float(1000.0)]);
403        assert_eq!(g.domain(), Some((0.0, 10.0)));
404    }
405
406    #[test]
407    fn test_gradient_n_interpolation() {
408        let g = ScaleColorGradientN::new(
409            Aesthetic::Color,
410            vec![
411                (0.0, RGBAColor::new(0, 0, 0)),
412                (0.5, RGBAColor::new(255, 0, 0)),
413                (1.0, RGBAColor::new(255, 255, 255)),
414            ],
415        );
416        // At t=0 should be black
417        let c0 = g.color_at(0.0);
418        assert_eq!((c0.r, c0.g, c0.b), (0, 0, 0));
419        // At t=0.5 should be red
420        let c5 = g.color_at(0.5);
421        assert_eq!((c5.r, c5.g, c5.b), (255, 0, 0));
422        // At t=1.0 should be white
423        let c1 = g.color_at(1.0);
424        assert_eq!((c1.r, c1.g, c1.b), (255, 255, 255));
425        // At t=0.25 should be ~midpoint between black and red
426        let c25 = g.color_at(0.25);
427        assert_eq!(c25.r, 127); // roughly half of 255
428    }
429
430    #[test]
431    fn test_viridis_continuous_endpoints() {
432        let g = ScaleColorGradientN::viridis(Aesthetic::Fill);
433        let c0 = g.color_at(0.0);
434        assert_eq!((c0.r, c0.g, c0.b), (68, 1, 84));
435        let c1 = g.color_at(1.0);
436        assert_eq!((c1.r, c1.g, c1.b), (253, 231, 37));
437    }
438}