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#[derive(Clone, Debug)]
11pub struct ScaleColorGradientN {
12 aesthetic: Aesthetic,
13 name: String,
14 stops: Vec<(f64, RGBAColor)>,
16 min: f64,
17 max: f64,
18}
19
20impl ScaleColorGradientN {
21 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 pub fn viridis(aesthetic: Aesthetic) -> Self {
37 Self::new(aesthetic, viridis_stops())
38 }
39
40 pub fn magma(aesthetic: Aesthetic) -> Self {
42 Self::new(aesthetic, magma_stops())
43 }
44
45 pub fn plasma(aesthetic: Aesthetic) -> Self {
47 Self::new(aesthetic, plasma_stops())
48 }
49
50 pub fn inferno(aesthetic: Aesthetic) -> Self {
52 Self::new(aesthetic, inferno_stops())
53 }
54
55 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 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
174fn 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 let c0 = g.color_at(0.0);
299 assert_eq!((c0.r, c0.g, c0.b), (0, 0, 0));
300 let c5 = g.color_at(0.5);
302 assert_eq!((c5.r, c5.g, c5.b), (255, 0, 0));
303 let c1 = g.color_at(1.0);
305 assert_eq!((c1.r, c1.g, c1.b), (255, 255, 255));
306 let c25 = g.color_at(0.25);
308 assert_eq!(c25.r, 127); }
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}