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.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 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 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
175fn 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 let c0 = g.color_at(0.0);
300 assert_eq!((c0.r, c0.g, c0.b), (0, 0, 0));
301 let c5 = g.color_at(0.5);
303 assert_eq!((c5.r, c5.g, c5.b), (255, 0, 0));
304 let c1 = g.color_at(1.0);
306 assert_eq!((c1.r, c1.g, c1.b), (255, 255, 255));
307 let c25 = g.color_at(0.25);
309 assert_eq!(c25.r, 127); }
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}