Skip to main content

ggplot_rs/scale/
color.rs

1use crate::aes::Aesthetic;
2use crate::data::Value;
3use std::collections::HashMap;
4
5use super::util::nice_step;
6use super::Scale;
7
8/// RGBA color representation.
9#[derive(Clone, Debug, Copy)]
10pub struct RGBAColor {
11    pub r: u8,
12    pub g: u8,
13    pub b: u8,
14    pub a: f64,
15}
16
17impl RGBAColor {
18    pub fn new(r: u8, g: u8, b: u8) -> Self {
19        RGBAColor { r, g, b, a: 1.0 }
20    }
21
22    pub fn with_alpha(mut self, a: f64) -> Self {
23        self.a = a;
24        self
25    }
26
27    /// Interpolate between two colors.
28    pub fn lerp(&self, other: &RGBAColor, t: f64) -> RGBAColor {
29        let t = t.clamp(0.0, 1.0);
30        RGBAColor {
31            r: (self.r as f64 * (1.0 - t) + other.r as f64 * t) as u8,
32            g: (self.g as f64 * (1.0 - t) + other.g as f64 * t) as u8,
33            b: (self.b as f64 * (1.0 - t) + other.b as f64 * t) as u8,
34            a: self.a * (1.0 - t) + other.a * t,
35        }
36    }
37}
38
39/// Default discrete color palette (8 colors, similar to ggplot2 default).
40pub const DEFAULT_PALETTE: &[RGBAColor] = &[
41    RGBAColor {
42        r: 248,
43        g: 118,
44        b: 109,
45        a: 1.0,
46    }, // red
47    RGBAColor {
48        r: 0,
49        g: 186,
50        b: 56,
51        a: 1.0,
52    }, // green
53    RGBAColor {
54        r: 97,
55        g: 156,
56        b: 255,
57        a: 1.0,
58    }, // blue
59    RGBAColor {
60        r: 163,
61        g: 103,
62        b: 203,
63        a: 1.0,
64    }, // purple
65    RGBAColor {
66        r: 231,
67        g: 138,
68        b: 0,
69        a: 1.0,
70    }, // orange
71    RGBAColor {
72        r: 0,
73        g: 191,
74        b: 196,
75        a: 1.0,
76    }, // cyan
77    RGBAColor {
78        r: 199,
79        g: 124,
80        b: 255,
81        a: 1.0,
82    }, // violet
83    RGBAColor {
84        r: 127,
85        g: 127,
86        b: 127,
87        a: 1.0,
88    }, // gray
89];
90
91/// Discrete color scale — maps categories to distinct colors.
92#[derive(Clone, Debug)]
93pub struct ScaleColorDiscrete {
94    aesthetic: Aesthetic,
95    name: String,
96    levels: Vec<String>,
97    /// level → index in `levels`, so training/lookup are O(1) (many-category
98    /// fills would otherwise be quadratic). Kept in sync with `levels`.
99    level_index: HashMap<String, usize>,
100    palette: Vec<RGBAColor>,
101    /// Keep levels in sorted (lexicographic) order instead of first-seen
102    /// order — see [`ScaleColorDiscrete::sorted`].
103    sorted: bool,
104}
105
106impl ScaleColorDiscrete {
107    pub fn new(aesthetic: Aesthetic) -> Self {
108        ScaleColorDiscrete {
109            aesthetic,
110            name: String::new(),
111            levels: Vec::new(),
112            level_index: HashMap::new(),
113            palette: DEFAULT_PALETTE.to_vec(),
114            sorted: false,
115        }
116    }
117
118    /// Order levels lexicographically (instead of first appearance) so a given
119    /// level always gets the same palette colour and legend position across
120    /// charts, whatever order the data arrives in — no need to pre-sort levels
121    /// and repeat a manual palette per chart.
122    pub fn sorted(mut self) -> Self {
123        self.sorted = true;
124        self.resort();
125        self
126    }
127
128    fn resort(&mut self) {
129        if !self.sorted {
130            return;
131        }
132        self.levels.sort();
133        self.level_index = self
134            .levels
135            .iter()
136            .enumerate()
137            .map(|(i, l)| (l.clone(), i))
138            .collect();
139    }
140
141    pub fn with_palette(mut self, colors: Vec<RGBAColor>) -> Self {
142        self.palette = colors;
143        self
144    }
145
146    pub fn with_named_palette(mut self, name: &super::palettes::PaletteName) -> Self {
147        self.palette = super::palettes::palette(name).to_vec();
148        self
149    }
150
151    /// Pre-seed the factor levels (order), so each level keeps a fixed palette
152    /// color regardless of which levels are present in the data — e.g. for a
153    /// legend whose series can be toggled without the colors reshuffling.
154    pub fn with_levels(mut self, levels: Vec<String>) -> Self {
155        self.levels.clear();
156        self.level_index.clear();
157        for l in levels {
158            self.push_level(l);
159        }
160        self.resort();
161        self
162    }
163
164    fn push_level(&mut self, key: String) {
165        if !self.level_index.contains_key(&key) {
166            self.level_index.insert(key.clone(), self.levels.len());
167            self.levels.push(key);
168        }
169    }
170
171    fn level_position(&self, key: &str) -> Option<usize> {
172        self.level_index.get(key).copied()
173    }
174
175    /// Get color for a given level index.
176    pub fn color_for_index(&self, idx: usize) -> RGBAColor {
177        if self.palette.is_empty() {
178            return DEFAULT_PALETTE[idx % DEFAULT_PALETTE.len()];
179        }
180        self.palette[idx % self.palette.len()]
181    }
182
183    /// Get color for a value.
184    pub fn color_for_value(&self, value: &Value) -> RGBAColor {
185        let key = value.key_str();
186        let idx = self.level_position(&key).unwrap_or(0);
187        self.color_for_index(idx)
188    }
189
190    pub fn levels(&self) -> &[String] {
191        &self.levels
192    }
193}
194
195impl Scale for ScaleColorDiscrete {
196    fn aesthetic(&self) -> Aesthetic {
197        self.aesthetic.clone()
198    }
199
200    fn train(&mut self, values: &[Value]) {
201        let before = self.levels.len();
202        for v in values {
203            let key = v.key_str();
204            if !self.level_index.contains_key(key.as_ref()) {
205                self.push_level(key.into_owned());
206            }
207        }
208        if self.levels.len() != before {
209            self.resort();
210        }
211    }
212
213    fn map(&self, value: &Value) -> f64 {
214        let key = value.key_str();
215        self.level_position(&key).map(|i| i as f64).unwrap_or(0.0)
216    }
217
218    fn breaks(&self) -> Vec<(f64, String)> {
219        self.levels
220            .iter()
221            .enumerate()
222            .map(|(i, label)| (i as f64, label.clone()))
223            .collect()
224    }
225
226    fn name(&self) -> &str {
227        &self.name
228    }
229
230    fn set_name(&mut self, name: &str) {
231        self.name = name.to_string();
232    }
233
234    fn is_discrete(&self) -> bool {
235        true
236    }
237
238    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
239        let c = self.color_for_value(value);
240        Some((c.r, c.g, c.b))
241    }
242
243    fn clone_box(&self) -> Box<dyn Scale> {
244        Box::new(self.clone())
245    }
246
247    fn reset_training(&mut self) {
248        self.levels.clear();
249        self.level_index.clear();
250    }
251}
252
253/// Continuous gradient color scale.
254#[derive(Clone, Debug)]
255pub struct ScaleColorContinuous {
256    aesthetic: Aesthetic,
257    name: String,
258    low: RGBAColor,
259    high: RGBAColor,
260    min: f64,
261    max: f64,
262    /// Locale for the default break labels (`None` = default marks).
263    locale: Option<crate::locale::Locale>,
264    /// Fixed `(lo, hi)` domain ([`with_limits`](Self::with_limits)):
265    /// `train` no longer widens it, and `map_to_color` maps against it
266    /// instead of the trained min/max.
267    limits: Option<(f64, f64)>,
268    /// Out-of-bounds policy for `limits` ([`with_oob`](Self::with_oob)).
269    /// Default [`Oob::Censor`].
270    oob: super::Oob,
271    /// Colour for a genuinely missing or out-of-bounds-and-censored value
272    /// ([`with_na_value`](Self::with_na_value)). Default ggplot2's
273    /// `"grey50"`, `(127, 127, 127)`.
274    na_value: (u8, u8, u8),
275}
276
277impl ScaleColorContinuous {
278    pub fn new(aesthetic: Aesthetic) -> Self {
279        ScaleColorContinuous {
280            aesthetic,
281            name: String::new(),
282            low: RGBAColor::new(0, 0, 255),
283            high: RGBAColor::new(255, 0, 0),
284            min: f64::INFINITY,
285            max: f64::NEG_INFINITY,
286            locale: None,
287            limits: None,
288            oob: super::Oob::Censor,
289            na_value: super::DEFAULT_NA_VALUE,
290        }
291    }
292
293    pub fn with_colors(mut self, low: RGBAColor, high: RGBAColor) -> Self {
294        self.low = low;
295        self.high = high;
296        self
297    }
298
299    /// Fix the colour domain to `[lo, hi]` — e.g. `(-1.0, 1.0)` for a
300    /// correlation matrix — instead of whatever the data happens to span.
301    /// Out-of-range values follow [`with_oob`](Self::with_oob) (default:
302    /// `na_value`, not the geom's own default colour).
303    pub fn with_limits(mut self, lo: f64, hi: f64) -> Self {
304        self.limits = Some((lo, hi));
305        self
306    }
307
308    /// Out-of-bounds policy for [`with_limits`](Self::with_limits):
309    /// [`Oob::Censor`] (default, ggplot2's own default) or
310    /// [`Oob::Squish`].
311    pub fn with_oob(mut self, oob: super::Oob) -> Self {
312        self.oob = oob;
313        self
314    }
315
316    /// Colour for a genuinely missing (`NA`) value, or one censored by
317    /// `with_limits`'s default [`Oob::Censor`]. Default ggplot2's
318    /// `"grey50"`, `(127, 127, 127)`.
319    pub fn with_na_value(mut self, color: (u8, u8, u8)) -> Self {
320        self.na_value = color;
321        self
322    }
323
324    pub fn color_at(&self, t: f64) -> RGBAColor {
325        self.low.lerp(&self.high, t)
326    }
327}
328
329impl Scale for ScaleColorContinuous {
330    fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
331        if self.locale.is_none() {
332            self.locale = Some(locale.clone());
333        }
334    }
335
336    fn locale(&self) -> Option<&crate::locale::Locale> {
337        self.locale.as_ref()
338    }
339
340    fn format_value(&self, value: &Value) -> Option<String> {
341        let f = value.as_f64()?;
342        if !f.is_finite() {
343            return None;
344        }
345        Some(super::util::format_number_opt(f, self.locale.as_ref()))
346    }
347
348    fn aesthetic(&self) -> Aesthetic {
349        self.aesthetic.clone()
350    }
351
352    fn train(&mut self, values: &[Value]) {
353        if self.limits.is_some() {
354            return;
355        }
356        for v in values {
357            if let Some(f) = v.as_f64() {
358                if f.is_finite() {
359                    if f < self.min {
360                        self.min = f;
361                    }
362                    if f > self.max {
363                        self.max = f;
364                    }
365                }
366            }
367        }
368    }
369
370    fn map(&self, value: &Value) -> f64 {
371        let f = match value.as_f64() {
372            Some(f) => f,
373            None => return 0.0,
374        };
375        let (lo, hi) = self.limits.unwrap_or((self.min, self.max));
376        let range = hi - lo;
377        if range.abs() < f64::EPSILON {
378            0.5
379        } else {
380            (f - lo) / range
381        }
382    }
383
384    fn breaks(&self) -> Vec<(f64, String)> {
385        let (lo, hi) = self.limits.unwrap_or((self.min, self.max));
386        if lo > hi || !lo.is_finite() || !hi.is_finite() {
387            return vec![];
388        }
389
390        let range = hi - lo;
391        if super::util::is_degenerate_range(lo, hi) {
392            return vec![(
393                0.5,
394                super::util::format_number_opt(lo, self.locale.as_ref()),
395            )];
396        }
397
398        let n_breaks = 5;
399        let raw_step = range / n_breaks as f64;
400        let step = nice_step(raw_step);
401
402        let start = (lo / step).ceil() * step;
403        let values = super::util::stepped_breaks(start, hi, step);
404        let labels = super::util::format_numbers_opt(&values, self.locale.as_ref());
405        values
406            .iter()
407            .zip(labels)
408            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
409            .collect()
410    }
411
412    fn name(&self) -> &str {
413        &self.name
414    }
415
416    fn set_name(&mut self, name: &str) {
417        self.name = name.to_string();
418    }
419
420    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
421        // A genuinely missing value, or a non-finite one (NaN/+-Inf),
422        // renders as `na_value` — ggplot2's own behaviour, and review
423        // PR#141 finding 3. A non-numeric-but-present value (e.g. the
424        // wrong column type) keeps the pre-existing contract of falling
425        // through to `map`'s own 0.0 fallback instead (this type's
426        // `map_to_color` always returned `Some` before `with_limits`/`Oob`
427        // existed, and nothing here can tell "wrong type" apart from
428        // "never had a numeric concept to begin with").
429        if value.is_na() {
430            return Some(self.na_value);
431        }
432        if let Some(f) = value.as_f64() {
433            if !f.is_finite() {
434                return Some(self.na_value);
435            }
436            if super::resolve_oob(f, self.limits, self.oob).is_none() {
437                return Some(self.na_value);
438            }
439        }
440        let t = self.map(value);
441        let c = self.color_at(t);
442        Some((c.r, c.g, c.b))
443    }
444
445    fn domain(&self) -> Option<(f64, f64)> {
446        if let Some((lo, hi)) = self.limits {
447            return Some((lo, hi));
448        }
449        if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
450            Some((self.min, self.max))
451        } else {
452            None
453        }
454    }
455
456    fn clone_box(&self) -> Box<dyn Scale> {
457        Box::new(self.clone())
458    }
459
460    fn reset_training(&mut self) {
461        self.min = f64::INFINITY;
462        self.max = f64::NEG_INFINITY;
463    }
464
465    fn validate(&self, warnings: &mut Vec<String>) {
466        if let Some((lo, hi)) = self.limits {
467            if lo > hi {
468                warnings.push(format!(
469                    "scale_color/fill_gradient: with_limits({lo}, {hi}) has lo > hi; every value maps to na_value"
470                ));
471            }
472        }
473    }
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479    use crate::aes::Aesthetic;
480
481    #[test]
482    fn continuous_format_value_rejects_non_finite() {
483        let s = ScaleColorContinuous::new(Aesthetic::Fill);
484        assert_eq!(s.format_value(&Value::Float(3.5)), Some("3.5".to_string()));
485        assert_eq!(s.format_value(&Value::Float(f64::NAN)), None);
486        assert_eq!(s.format_value(&Value::Na), None);
487    }
488}