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.total_cmp(&b.0));
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 super::util::is_degenerate_range(self.min, self.max) {
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 values = super::util::stepped_breaks(start, self.max, step);
134        let labels = super::util::format_numbers(&values);
135        values
136            .iter()
137            .zip(labels)
138            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
139            .collect()
140    }
141
142    fn name(&self) -> &str {
143        &self.name
144    }
145
146    fn set_name(&mut self, name: &str) {
147        self.name = name.to_string();
148    }
149
150    fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
151        let t = self.map(value);
152        let c = self.color_at(t);
153        Some((c.r, c.g, c.b))
154    }
155
156    fn domain(&self) -> Option<(f64, f64)> {
157        if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
158            Some((self.min, self.max))
159        } else {
160            None
161        }
162    }
163
164    fn clone_box(&self) -> Box<dyn Scale> {
165        Box::new(self.clone())
166    }
167
168    fn reset_training(&mut self) {
169        self.min = f64::INFINITY;
170        self.max = f64::NEG_INFINITY;
171    }
172}
173
174// ─── Continuous palette color stops ──────────────────────────────
175
176fn c(r: u8, g: u8, b: u8) -> RGBAColor {
177    RGBAColor::new(r, g, b)
178}
179
180fn viridis_stops() -> Vec<(f64, RGBAColor)> {
181    let colors = [
182        c(68, 1, 84),
183        c(72, 26, 108),
184        c(71, 47, 126),
185        c(65, 68, 135),
186        c(57, 86, 140),
187        c(47, 104, 142),
188        c(38, 121, 142),
189        c(31, 138, 141),
190        c(30, 155, 138),
191        c(42, 172, 130),
192        c(70, 188, 115),
193        c(109, 202, 93),
194        c(155, 213, 67),
195        c(200, 222, 39),
196        c(240, 229, 30),
197        c(253, 231, 37),
198    ];
199    evenly_spaced_stops(&colors)
200}
201
202fn magma_stops() -> Vec<(f64, RGBAColor)> {
203    let colors = [
204        c(0, 0, 4),
205        c(16, 12, 50),
206        c(41, 17, 90),
207        c(72, 12, 110),
208        c(101, 19, 110),
209        c(131, 29, 103),
210        c(160, 42, 93),
211        c(187, 55, 84),
212        c(213, 72, 72),
213        c(232, 99, 62),
214        c(247, 131, 57),
215        c(254, 167, 69),
216        c(254, 203, 99),
217        c(252, 235, 141),
218        c(252, 254, 188),
219        c(252, 253, 191),
220    ];
221    evenly_spaced_stops(&colors)
222}
223
224fn plasma_stops() -> Vec<(f64, RGBAColor)> {
225    let colors = [
226        c(13, 8, 135),
227        c(53, 5, 157),
228        c(82, 1, 163),
229        c(109, 1, 159),
230        c(133, 7, 147),
231        c(156, 23, 127),
232        c(175, 42, 106),
233        c(192, 61, 85),
234        c(206, 82, 66),
235        c(218, 105, 46),
236        c(228, 130, 24),
237        c(236, 157, 6),
238        c(240, 185, 11),
239        c(239, 213, 38),
240        c(232, 240, 73),
241        c(240, 249, 33),
242    ];
243    evenly_spaced_stops(&colors)
244}
245
246fn inferno_stops() -> Vec<(f64, RGBAColor)> {
247    let colors = [
248        c(0, 0, 4),
249        c(14, 11, 49),
250        c(39, 15, 90),
251        c(67, 10, 107),
252        c(95, 13, 106),
253        c(122, 21, 97),
254        c(149, 33, 81),
255        c(174, 49, 60),
256        c(196, 69, 38),
257        c(215, 95, 15),
258        c(231, 124, 3),
259        c(243, 155, 7),
260        c(250, 189, 28),
261        c(252, 222, 67),
262        c(247, 252, 118),
263        c(252, 255, 164),
264    ];
265    evenly_spaced_stops(&colors)
266}
267
268fn evenly_spaced_stops(colors: &[RGBAColor]) -> Vec<(f64, RGBAColor)> {
269    let n = colors.len();
270    if n == 0 {
271        return vec![];
272    }
273    if n == 1 {
274        return vec![(0.0, colors[0])];
275    }
276    colors
277        .iter()
278        .enumerate()
279        .map(|(i, c)| (i as f64 / (n - 1) as f64, *c))
280        .collect()
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    #[test]
288    fn test_gradient_n_interpolation() {
289        let g = ScaleColorGradientN::new(
290            Aesthetic::Color,
291            vec![
292                (0.0, RGBAColor::new(0, 0, 0)),
293                (0.5, RGBAColor::new(255, 0, 0)),
294                (1.0, RGBAColor::new(255, 255, 255)),
295            ],
296        );
297        // At t=0 should be black
298        let c0 = g.color_at(0.0);
299        assert_eq!((c0.r, c0.g, c0.b), (0, 0, 0));
300        // At t=0.5 should be red
301        let c5 = g.color_at(0.5);
302        assert_eq!((c5.r, c5.g, c5.b), (255, 0, 0));
303        // At t=1.0 should be white
304        let c1 = g.color_at(1.0);
305        assert_eq!((c1.r, c1.g, c1.b), (255, 255, 255));
306        // At t=0.25 should be ~midpoint between black and red
307        let c25 = g.color_at(0.25);
308        assert_eq!(c25.r, 127); // roughly half of 255
309    }
310
311    #[test]
312    fn test_viridis_continuous_endpoints() {
313        let g = ScaleColorGradientN::viridis(Aesthetic::Fill);
314        let c0 = g.color_at(0.0);
315        assert_eq!((c0.r, c0.g, c0.b), (68, 1, 84));
316        let c1 = g.color_at(1.0);
317        assert_eq!((c1.r, c1.g, c1.b), (253, 231, 37));
318    }
319}