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