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 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
171fn 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 let c0 = g.color_at(0.0);
296 assert_eq!((c0.r, c0.g, c0.b), (0, 0, 0));
297 let c5 = g.color_at(0.5);
299 assert_eq!((c5.r, c5.g, c5.b), (255, 0, 0));
300 let c1 = g.color_at(1.0);
302 assert_eq!((c1.r, c1.g, c1.b), (255, 255, 255));
303 let c25 = g.color_at(0.25);
305 assert_eq!(c25.r, 127); }
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}