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)]
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    /// Formatter for [`format_value`](super::Scale::format_value) (#164,
276    /// same shape as `ScaleContinuous::with_value_formatter`), distinct
277    /// from the axis-break label formatting this scale has none of. Falls
278    /// back to the plain locale-formatted number when unset. Wrapped in
279    /// [`UnwindSafeFormatter`](super::format::UnwindSafeFormatter) so this
280    /// field doesn't drop `UnwindSafe`/`RefUnwindSafe` from the struct.
281    value_formatter: Option<super::format::UnwindSafeFormatter>,
282}
283
284// Manual `Debug`, since `value_formatter` (an `Arc<dyn Fn>`) isn't
285// `Debug`; every other field is printed as the derive would.
286impl std::fmt::Debug for ScaleColorContinuous {
287    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
288        f.debug_struct("ScaleColorContinuous")
289            .field("aesthetic", &self.aesthetic)
290            .field("name", &self.name)
291            .field("low", &self.low)
292            .field("high", &self.high)
293            .field("min", &self.min)
294            .field("max", &self.max)
295            .field("locale", &self.locale)
296            .field("limits", &self.limits)
297            .field("oob", &self.oob)
298            .field("na_value", &self.na_value)
299            .field("value_formatter", &self.value_formatter.is_some())
300            .finish()
301    }
302}
303
304impl ScaleColorContinuous {
305    pub fn new(aesthetic: Aesthetic) -> Self {
306        ScaleColorContinuous {
307            aesthetic,
308            name: String::new(),
309            low: RGBAColor::new(0, 0, 255),
310            high: RGBAColor::new(255, 0, 0),
311            min: f64::INFINITY,
312            max: f64::NEG_INFINITY,
313            locale: None,
314            limits: None,
315            oob: super::Oob::Censor,
316            na_value: super::DEFAULT_NA_VALUE,
317            value_formatter: None,
318        }
319    }
320
321    /// Set a formatter for [`format_value`](super::Scale::format_value) —
322    /// a per-mark `data-fvalue`/tooltip-template value for this fill/colour
323    /// scale (#164). Falls back to the plain locale-formatted number when
324    /// unset (byte-identical default output).
325    pub fn with_value_formatter<F>(mut self, f: F) -> Self
326    where
327        F: Fn(f64) -> String + Send + Sync + 'static,
328    {
329        self.value_formatter = Some(super::format::UnwindSafeFormatter(std::sync::Arc::new(f)));
330        self
331    }
332
333    pub fn with_colors(mut self, low: RGBAColor, high: RGBAColor) -> Self {
334        self.low = low;
335        self.high = high;
336        self
337    }
338
339    /// Fix the colour domain to `[lo, hi]` — e.g. `(-1.0, 1.0)` for a
340    /// correlation matrix — instead of whatever the data happens to span.
341    /// Out-of-range values follow [`with_oob`](Self::with_oob) (default:
342    /// `na_value`, not the geom's own default colour).
343    pub fn with_limits(mut self, lo: f64, hi: f64) -> Self {
344        self.limits = Some((lo, hi));
345        self
346    }
347
348    /// Out-of-bounds policy for [`with_limits`](Self::with_limits):
349    /// [`Oob::Censor`] (default, ggplot2's own default) or
350    /// [`Oob::Squish`].
351    pub fn with_oob(mut self, oob: super::Oob) -> Self {
352        self.oob = oob;
353        self
354    }
355
356    /// Colour for a genuinely missing (`NA`) value, or one censored by
357    /// `with_limits`'s default [`Oob::Censor`]. Default ggplot2's
358    /// `"grey50"`, `(127, 127, 127)`.
359    pub fn with_na_value(mut self, color: (u8, u8, u8)) -> Self {
360        self.na_value = color;
361        self
362    }
363
364    pub fn color_at(&self, t: f64) -> RGBAColor {
365        self.low.lerp(&self.high, t)
366    }
367}
368
369impl Scale for ScaleColorContinuous {
370    fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
371        if self.locale.is_none() {
372            self.locale = Some(locale.clone());
373        }
374    }
375
376    fn locale(&self) -> Option<&crate::locale::Locale> {
377        self.locale.as_ref()
378    }
379
380    fn format_value(&self, value: &Value) -> Option<String> {
381        let f = value.as_f64()?;
382        if !f.is_finite() {
383            return None;
384        }
385        Some(match &self.value_formatter {
386            Some(fmt) => fmt.call(f),
387            None => super::util::format_number_opt(f, self.locale.as_ref()),
388        })
389    }
390
391    fn aesthetic(&self) -> Aesthetic {
392        self.aesthetic.clone()
393    }
394
395    fn train(&mut self, values: &[Value]) {
396        if self.limits.is_some() {
397            return;
398        }
399        for v in values {
400            if let Some(f) = v.as_f64() {
401                if f.is_finite() {
402                    if f < self.min {
403                        self.min = f;
404                    }
405                    if f > self.max {
406                        self.max = f;
407                    }
408                }
409            }
410        }
411    }
412
413    fn map(&self, value: &Value) -> f64 {
414        let f = match value.as_f64() {
415            Some(f) => f,
416            None => return 0.0,
417        };
418        let (lo, hi) = self.limits.unwrap_or((self.min, self.max));
419        let range = hi - lo;
420        if range.abs() < f64::EPSILON {
421            0.5
422        } else {
423            (f - lo) / range
424        }
425    }
426
427    fn breaks(&self) -> Vec<(f64, String)> {
428        let (lo, hi) = self.limits.unwrap_or((self.min, self.max));
429        if lo > hi || !lo.is_finite() || !hi.is_finite() {
430            return vec![];
431        }
432
433        let range = hi - lo;
434        if super::util::is_degenerate_range(lo, hi) {
435            return vec![(
436                0.5,
437                super::util::format_number_opt(lo, self.locale.as_ref()),
438            )];
439        }
440
441        let n_breaks = 5;
442        let raw_step = range / n_breaks as f64;
443        let step = nice_step(raw_step);
444
445        let start = (lo / step).ceil() * step;
446        let values = super::util::stepped_breaks(start, hi, step);
447        let labels = super::util::format_numbers_opt(&values, self.locale.as_ref());
448        values
449            .iter()
450            .zip(labels)
451            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
452            .collect()
453    }
454
455    fn name(&self) -> &str {
456        &self.name
457    }
458
459    fn set_name(&mut self, name: &str) {
460        self.name = name.to_string();
461    }
462
463    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
464        // A genuinely missing value, or a non-finite one (NaN/+-Inf),
465        // renders as `na_value` — ggplot2's own behaviour, and review
466        // PR#141 finding 3. A non-numeric-but-present value (e.g. the
467        // wrong column type) keeps the pre-existing contract of falling
468        // through to `map`'s own 0.0 fallback instead (this type's
469        // `map_to_color` always returned `Some` before `with_limits`/`Oob`
470        // existed, and nothing here can tell "wrong type" apart from
471        // "never had a numeric concept to begin with").
472        if value.is_na() {
473            return Some(self.na_value);
474        }
475        if let Some(f) = value.as_f64() {
476            if !f.is_finite() {
477                return Some(self.na_value);
478            }
479            if super::resolve_oob(f, self.limits, self.oob).is_none() {
480                return Some(self.na_value);
481            }
482        }
483        let t = self.map(value);
484        let c = self.color_at(t);
485        Some((c.r, c.g, c.b))
486    }
487
488    fn domain(&self) -> Option<(f64, f64)> {
489        if let Some((lo, hi)) = self.limits {
490            return Some((lo, hi));
491        }
492        if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
493            Some((self.min, self.max))
494        } else {
495            None
496        }
497    }
498
499    fn clone_box(&self) -> Box<dyn Scale> {
500        Box::new(self.clone())
501    }
502
503    fn reset_training(&mut self) {
504        self.min = f64::INFINITY;
505        self.max = f64::NEG_INFINITY;
506    }
507
508    fn validate(&self, warnings: &mut Vec<String>) {
509        if let Some((lo, hi)) = self.limits {
510            if lo > hi {
511                warnings.push(format!(
512                    "scale_color/fill_gradient: with_limits({lo}, {hi}) has lo > hi; every value maps to na_value"
513                ));
514            }
515        }
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use super::*;
522    use crate::aes::Aesthetic;
523
524    #[test]
525    fn continuous_format_value_rejects_non_finite() {
526        let s = ScaleColorContinuous::new(Aesthetic::Fill);
527        assert_eq!(s.format_value(&Value::Float(3.5)), Some("3.5".to_string()));
528        assert_eq!(s.format_value(&Value::Float(f64::NAN)), None);
529        assert_eq!(s.format_value(&Value::Na), None);
530    }
531
532    #[test]
533    fn continuous_with_value_formatter_overrides_format_value() {
534        let s =
535            ScaleColorContinuous::new(Aesthetic::Fill).with_value_formatter(|v| format!("${v:.2}"));
536        assert_eq!(
537            s.format_value(&Value::Float(1234.5)),
538            Some("$1234.50".to_string())
539        );
540        // Non-finite is still rejected before the formatter ever runs.
541        assert_eq!(s.format_value(&Value::Float(f64::NAN)), None);
542    }
543
544    #[test]
545    fn continuous_debug_reports_whether_a_formatter_is_set() {
546        let without = format!("{:?}", ScaleColorContinuous::new(Aesthetic::Fill));
547        assert!(without.contains("ScaleColorContinuous"));
548        assert!(without.contains("value_formatter: false"));
549        let with = format!(
550            "{:?}",
551            ScaleColorContinuous::new(Aesthetic::Fill).with_value_formatter(|v| v.to_string())
552        );
553        assert!(with.contains("value_formatter: true"));
554    }
555
556    /// #164 review: `value_formatter`'s `Arc<dyn Fn>` must not drop
557    /// `UnwindSafe`/`RefUnwindSafe` (or `Send`/`Sync`/`Clone`) from
558    /// `ScaleColorContinuous` — `cargo-semver-checks` flags losing an
559    /// auto trait as a major-version break, which a patch must not cause.
560    #[test]
561    fn scale_color_continuous_keeps_every_auto_trait() {
562        fn assert_traits<
563            T: Send + Sync + std::panic::UnwindSafe + std::panic::RefUnwindSafe + Clone,
564        >() {
565        }
566        assert_traits::<ScaleColorContinuous>();
567    }
568}