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