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}
265
266impl ScaleColorContinuous {
267    pub fn new(aesthetic: Aesthetic) -> Self {
268        ScaleColorContinuous {
269            aesthetic,
270            name: String::new(),
271            low: RGBAColor::new(0, 0, 255),
272            high: RGBAColor::new(255, 0, 0),
273            min: f64::INFINITY,
274            max: f64::NEG_INFINITY,
275            locale: None,
276        }
277    }
278
279    pub fn with_colors(mut self, low: RGBAColor, high: RGBAColor) -> Self {
280        self.low = low;
281        self.high = high;
282        self
283    }
284
285    pub fn color_at(&self, t: f64) -> RGBAColor {
286        self.low.lerp(&self.high, t)
287    }
288}
289
290impl Scale for ScaleColorContinuous {
291    fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
292        if self.locale.is_none() {
293            self.locale = Some(locale.clone());
294        }
295    }
296
297    fn locale(&self) -> Option<&crate::locale::Locale> {
298        self.locale.as_ref()
299    }
300
301    fn aesthetic(&self) -> Aesthetic {
302        self.aesthetic.clone()
303    }
304
305    fn train(&mut self, values: &[Value]) {
306        for v in values {
307            if let Some(f) = v.as_f64() {
308                if f.is_finite() {
309                    if f < self.min {
310                        self.min = f;
311                    }
312                    if f > self.max {
313                        self.max = f;
314                    }
315                }
316            }
317        }
318    }
319
320    fn map(&self, value: &Value) -> f64 {
321        let f = match value.as_f64() {
322            Some(f) => f,
323            None => return 0.0,
324        };
325        let range = self.max - self.min;
326        if range.abs() < f64::EPSILON {
327            0.5
328        } else {
329            (f - self.min) / range
330        }
331    }
332
333    fn breaks(&self) -> Vec<(f64, String)> {
334        if self.min > self.max || !self.min.is_finite() || !self.max.is_finite() {
335            return vec![];
336        }
337
338        let range = self.max - self.min;
339        if super::util::is_degenerate_range(self.min, self.max) {
340            return vec![(
341                0.5,
342                super::util::format_number_opt(self.min, self.locale.as_ref()),
343            )];
344        }
345
346        let n_breaks = 5;
347        let raw_step = range / n_breaks as f64;
348        let step = nice_step(raw_step);
349
350        let start = (self.min / step).ceil() * step;
351        let values = super::util::stepped_breaks(start, self.max, step);
352        let labels = super::util::format_numbers_opt(&values, self.locale.as_ref());
353        values
354            .iter()
355            .zip(labels)
356            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
357            .collect()
358    }
359
360    fn name(&self) -> &str {
361        &self.name
362    }
363
364    fn set_name(&mut self, name: &str) {
365        self.name = name.to_string();
366    }
367
368    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
369        let t = self.map(value);
370        let c = self.color_at(t);
371        Some((c.r, c.g, c.b))
372    }
373
374    fn domain(&self) -> Option<(f64, f64)> {
375        if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
376            Some((self.min, self.max))
377        } else {
378            None
379        }
380    }
381
382    fn clone_box(&self) -> Box<dyn Scale> {
383        Box::new(self.clone())
384    }
385
386    fn reset_training(&mut self) {
387        self.min = f64::INFINITY;
388        self.max = f64::NEG_INFINITY;
389    }
390}