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