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}
21
22impl ScaleColorGradientN {
23    /// Create a new N-stop gradient for the given aesthetic.
24    /// Stops are `(position, color)` where position is in [0, 1].
25    pub fn new(aesthetic: Aesthetic, stops: Vec<(f64, RGBAColor)>) -> Self {
26        let mut stops = stops;
27        stops.sort_by(|a, b| a.0.total_cmp(&b.0));
28        ScaleColorGradientN {
29            aesthetic,
30            name: String::new(),
31            stops,
32            min: f64::INFINITY,
33            max: f64::NEG_INFINITY,
34            locale: None,
35        }
36    }
37
38    /// Create a continuous viridis palette.
39    pub fn viridis(aesthetic: Aesthetic) -> Self {
40        Self::new(aesthetic, viridis_stops())
41    }
42
43    /// Create a continuous magma palette.
44    pub fn magma(aesthetic: Aesthetic) -> Self {
45        Self::new(aesthetic, magma_stops())
46    }
47
48    /// Create a continuous plasma palette.
49    pub fn plasma(aesthetic: Aesthetic) -> Self {
50        Self::new(aesthetic, plasma_stops())
51    }
52
53    /// Create a continuous inferno palette.
54    pub fn inferno(aesthetic: Aesthetic) -> Self {
55        Self::new(aesthetic, inferno_stops())
56    }
57
58    /// Interpolate the color at a normalized position t in [0, 1].
59    fn color_at(&self, t: f64) -> RGBAColor {
60        let t = t.clamp(0.0, 1.0);
61        if self.stops.is_empty() {
62            return RGBAColor::new(127, 127, 127);
63        }
64        if self.stops.len() == 1 {
65            return self.stops[0].1;
66        }
67        // Find the two surrounding stops
68        if t <= self.stops[0].0 {
69            return self.stops[0].1;
70        }
71        if t >= self.stops[self.stops.len() - 1].0 {
72            return self.stops[self.stops.len() - 1].1;
73        }
74        for i in 0..self.stops.len() - 1 {
75            let (p0, c0) = &self.stops[i];
76            let (p1, c1) = &self.stops[i + 1];
77            if t >= *p0 && t <= *p1 {
78                let range = p1 - p0;
79                let local_t = if range.abs() < f64::EPSILON {
80                    0.0
81                } else {
82                    (t - p0) / range
83                };
84                return c0.lerp(c1, local_t);
85            }
86        }
87        self.stops[self.stops.len() - 1].1
88    }
89}
90
91impl Scale for ScaleColorGradientN {
92    fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
93        if self.locale.is_none() {
94            self.locale = Some(locale.clone());
95        }
96    }
97
98    fn locale(&self) -> Option<&crate::locale::Locale> {
99        self.locale.as_ref()
100    }
101
102    fn aesthetic(&self) -> Aesthetic {
103        self.aesthetic.clone()
104    }
105
106    fn train(&mut self, values: &[Value]) {
107        for v in values {
108            if let Some(f) = v.as_f64() {
109                if f.is_finite() {
110                    if f < self.min {
111                        self.min = f;
112                    }
113                    if f > self.max {
114                        self.max = f;
115                    }
116                }
117            }
118        }
119    }
120
121    fn map(&self, value: &Value) -> f64 {
122        let f = match value.as_f64() {
123            Some(f) => f,
124            None => return 0.0,
125        };
126        let range = self.max - self.min;
127        if range.abs() < f64::EPSILON {
128            0.5
129        } else {
130            (f - self.min) / range
131        }
132    }
133
134    fn breaks(&self) -> Vec<(f64, String)> {
135        if self.min > self.max || !self.min.is_finite() || !self.max.is_finite() {
136            return vec![];
137        }
138        let range = self.max - self.min;
139        if super::util::is_degenerate_range(self.min, self.max) {
140            return vec![(
141                0.5,
142                super::util::format_number_opt(self.min, self.locale.as_ref()),
143            )];
144        }
145        let n_breaks = 5;
146        let raw_step = range / n_breaks as f64;
147        let step = nice_step(raw_step);
148        let start = (self.min / step).ceil() * step;
149        let values = super::util::stepped_breaks(start, self.max, step);
150        let labels = super::util::format_numbers_opt(&values, self.locale.as_ref());
151        values
152            .iter()
153            .zip(labels)
154            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
155            .collect()
156    }
157
158    fn name(&self) -> &str {
159        &self.name
160    }
161
162    fn set_name(&mut self, name: &str) {
163        self.name = name.to_string();
164    }
165
166    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
167        let t = self.map(value);
168        let c = self.color_at(t);
169        Some((c.r, c.g, c.b))
170    }
171
172    fn domain(&self) -> Option<(f64, f64)> {
173        if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
174            Some((self.min, self.max))
175        } else {
176            None
177        }
178    }
179
180    fn clone_box(&self) -> Box<dyn Scale> {
181        Box::new(self.clone())
182    }
183
184    fn reset_training(&mut self) {
185        self.min = f64::INFINITY;
186        self.max = f64::NEG_INFINITY;
187    }
188}
189
190// ─── Continuous palette color stops ──────────────────────────────
191
192fn c(r: u8, g: u8, b: u8) -> RGBAColor {
193    RGBAColor::new(r, g, b)
194}
195
196fn viridis_stops() -> Vec<(f64, RGBAColor)> {
197    let colors = [
198        c(68, 1, 84),
199        c(72, 26, 108),
200        c(71, 47, 126),
201        c(65, 68, 135),
202        c(57, 86, 140),
203        c(47, 104, 142),
204        c(38, 121, 142),
205        c(31, 138, 141),
206        c(30, 155, 138),
207        c(42, 172, 130),
208        c(70, 188, 115),
209        c(109, 202, 93),
210        c(155, 213, 67),
211        c(200, 222, 39),
212        c(240, 229, 30),
213        c(253, 231, 37),
214    ];
215    evenly_spaced_stops(&colors)
216}
217
218fn magma_stops() -> Vec<(f64, RGBAColor)> {
219    let colors = [
220        c(0, 0, 4),
221        c(16, 12, 50),
222        c(41, 17, 90),
223        c(72, 12, 110),
224        c(101, 19, 110),
225        c(131, 29, 103),
226        c(160, 42, 93),
227        c(187, 55, 84),
228        c(213, 72, 72),
229        c(232, 99, 62),
230        c(247, 131, 57),
231        c(254, 167, 69),
232        c(254, 203, 99),
233        c(252, 235, 141),
234        c(252, 254, 188),
235        c(252, 253, 191),
236    ];
237    evenly_spaced_stops(&colors)
238}
239
240fn plasma_stops() -> Vec<(f64, RGBAColor)> {
241    let colors = [
242        c(13, 8, 135),
243        c(53, 5, 157),
244        c(82, 1, 163),
245        c(109, 1, 159),
246        c(133, 7, 147),
247        c(156, 23, 127),
248        c(175, 42, 106),
249        c(192, 61, 85),
250        c(206, 82, 66),
251        c(218, 105, 46),
252        c(228, 130, 24),
253        c(236, 157, 6),
254        c(240, 185, 11),
255        c(239, 213, 38),
256        c(232, 240, 73),
257        c(240, 249, 33),
258    ];
259    evenly_spaced_stops(&colors)
260}
261
262fn inferno_stops() -> Vec<(f64, RGBAColor)> {
263    let colors = [
264        c(0, 0, 4),
265        c(14, 11, 49),
266        c(39, 15, 90),
267        c(67, 10, 107),
268        c(95, 13, 106),
269        c(122, 21, 97),
270        c(149, 33, 81),
271        c(174, 49, 60),
272        c(196, 69, 38),
273        c(215, 95, 15),
274        c(231, 124, 3),
275        c(243, 155, 7),
276        c(250, 189, 28),
277        c(252, 222, 67),
278        c(247, 252, 118),
279        c(252, 255, 164),
280    ];
281    evenly_spaced_stops(&colors)
282}
283
284fn evenly_spaced_stops(colors: &[RGBAColor]) -> Vec<(f64, RGBAColor)> {
285    let n = colors.len();
286    if n == 0 {
287        return vec![];
288    }
289    if n == 1 {
290        return vec![(0.0, colors[0])];
291    }
292    colors
293        .iter()
294        .enumerate()
295        .map(|(i, c)| (i as f64 / (n - 1) as f64, *c))
296        .collect()
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn test_gradient_n_interpolation() {
305        let g = ScaleColorGradientN::new(
306            Aesthetic::Color,
307            vec![
308                (0.0, RGBAColor::new(0, 0, 0)),
309                (0.5, RGBAColor::new(255, 0, 0)),
310                (1.0, RGBAColor::new(255, 255, 255)),
311            ],
312        );
313        // At t=0 should be black
314        let c0 = g.color_at(0.0);
315        assert_eq!((c0.r, c0.g, c0.b), (0, 0, 0));
316        // At t=0.5 should be red
317        let c5 = g.color_at(0.5);
318        assert_eq!((c5.r, c5.g, c5.b), (255, 0, 0));
319        // At t=1.0 should be white
320        let c1 = g.color_at(1.0);
321        assert_eq!((c1.r, c1.g, c1.b), (255, 255, 255));
322        // At t=0.25 should be ~midpoint between black and red
323        let c25 = g.color_at(0.25);
324        assert_eq!(c25.r, 127); // roughly half of 255
325    }
326
327    #[test]
328    fn test_viridis_continuous_endpoints() {
329        let g = ScaleColorGradientN::viridis(Aesthetic::Fill);
330        let c0 = g.color_at(0.0);
331        assert_eq!((c0.r, c0.g, c0.b), (68, 1, 84));
332        let c1 = g.color_at(1.0);
333        assert_eq!((c1.r, c1.g, c1.b), (253, 231, 37));
334    }
335}