1use crate::aes::Aesthetic;
2use crate::data::Value;
3use std::collections::HashMap;
4
5use super::util::{format_number, nice_step};
6use super::Scale;
7
8#[derive(Clone, Debug, Copy)]
10pub struct RGBAColor {
11 pub r: u8,
12 pub g: u8,
13 pub b: u8,
14 pub a: f64,
15}
16
17impl RGBAColor {
18 pub fn new(r: u8, g: u8, b: u8) -> Self {
19 RGBAColor { r, g, b, a: 1.0 }
20 }
21
22 pub fn with_alpha(mut self, a: f64) -> Self {
23 self.a = a;
24 self
25 }
26
27 pub fn lerp(&self, other: &RGBAColor, t: f64) -> RGBAColor {
29 let t = t.clamp(0.0, 1.0);
30 RGBAColor {
31 r: (self.r as f64 * (1.0 - t) + other.r as f64 * t) as u8,
32 g: (self.g as f64 * (1.0 - t) + other.g as f64 * t) as u8,
33 b: (self.b as f64 * (1.0 - t) + other.b as f64 * t) as u8,
34 a: self.a * (1.0 - t) + other.a * t,
35 }
36 }
37}
38
39pub const DEFAULT_PALETTE: &[RGBAColor] = &[
41 RGBAColor {
42 r: 248,
43 g: 118,
44 b: 109,
45 a: 1.0,
46 }, RGBAColor {
48 r: 0,
49 g: 186,
50 b: 56,
51 a: 1.0,
52 }, RGBAColor {
54 r: 97,
55 g: 156,
56 b: 255,
57 a: 1.0,
58 }, RGBAColor {
60 r: 163,
61 g: 103,
62 b: 203,
63 a: 1.0,
64 }, RGBAColor {
66 r: 231,
67 g: 138,
68 b: 0,
69 a: 1.0,
70 }, RGBAColor {
72 r: 0,
73 g: 191,
74 b: 196,
75 a: 1.0,
76 }, RGBAColor {
78 r: 199,
79 g: 124,
80 b: 255,
81 a: 1.0,
82 }, RGBAColor {
84 r: 127,
85 g: 127,
86 b: 127,
87 a: 1.0,
88 }, ];
90
91#[derive(Clone, Debug)]
93pub struct ScaleColorDiscrete {
94 aesthetic: Aesthetic,
95 name: String,
96 levels: Vec<String>,
97 level_index: HashMap<String, usize>,
100 palette: Vec<RGBAColor>,
101 sorted: bool,
104}
105
106impl ScaleColorDiscrete {
107 pub fn new(aesthetic: Aesthetic) -> Self {
108 ScaleColorDiscrete {
109 aesthetic,
110 name: String::new(),
111 levels: Vec::new(),
112 level_index: HashMap::new(),
113 palette: DEFAULT_PALETTE.to_vec(),
114 sorted: false,
115 }
116 }
117
118 pub fn sorted(mut self) -> Self {
123 self.sorted = true;
124 self.resort();
125 self
126 }
127
128 fn resort(&mut self) {
129 if !self.sorted {
130 return;
131 }
132 self.levels.sort();
133 self.level_index = self
134 .levels
135 .iter()
136 .enumerate()
137 .map(|(i, l)| (l.clone(), i))
138 .collect();
139 }
140
141 pub fn with_palette(mut self, colors: Vec<RGBAColor>) -> Self {
142 self.palette = colors;
143 self
144 }
145
146 pub fn with_named_palette(mut self, name: &super::palettes::PaletteName) -> Self {
147 self.palette = super::palettes::palette(name).to_vec();
148 self
149 }
150
151 pub fn with_levels(mut self, levels: Vec<String>) -> Self {
155 self.levels.clear();
156 self.level_index.clear();
157 for l in levels {
158 self.push_level(l);
159 }
160 self.resort();
161 self
162 }
163
164 fn push_level(&mut self, key: String) {
165 if !self.level_index.contains_key(&key) {
166 self.level_index.insert(key.clone(), self.levels.len());
167 self.levels.push(key);
168 }
169 }
170
171 fn level_position(&self, key: &str) -> Option<usize> {
172 self.level_index.get(key).copied()
173 }
174
175 pub fn color_for_index(&self, idx: usize) -> RGBAColor {
177 if self.palette.is_empty() {
178 return DEFAULT_PALETTE[idx % DEFAULT_PALETTE.len()];
179 }
180 self.palette[idx % self.palette.len()]
181 }
182
183 pub fn color_for_value(&self, value: &Value) -> RGBAColor {
185 let key = value.key_str();
186 let idx = self.level_position(&key).unwrap_or(0);
187 self.color_for_index(idx)
188 }
189
190 pub fn levels(&self) -> &[String] {
191 &self.levels
192 }
193}
194
195impl Scale for ScaleColorDiscrete {
196 fn aesthetic(&self) -> Aesthetic {
197 self.aesthetic.clone()
198 }
199
200 fn train(&mut self, values: &[Value]) {
201 let before = self.levels.len();
202 for v in values {
203 let key = v.key_str();
204 if !self.level_index.contains_key(key.as_ref()) {
205 self.push_level(key.into_owned());
206 }
207 }
208 if self.levels.len() != before {
209 self.resort();
210 }
211 }
212
213 fn map(&self, value: &Value) -> f64 {
214 let key = value.key_str();
215 self.level_position(&key).map(|i| i as f64).unwrap_or(0.0)
216 }
217
218 fn breaks(&self) -> Vec<(f64, String)> {
219 self.levels
220 .iter()
221 .enumerate()
222 .map(|(i, label)| (i as f64, label.clone()))
223 .collect()
224 }
225
226 fn name(&self) -> &str {
227 &self.name
228 }
229
230 fn set_name(&mut self, name: &str) {
231 self.name = name.to_string();
232 }
233
234 fn is_discrete(&self) -> bool {
235 true
236 }
237
238 fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
239 let c = self.color_for_value(value);
240 Some((c.r, c.g, c.b))
241 }
242
243 fn clone_box(&self) -> Box<dyn Scale> {
244 Box::new(self.clone())
245 }
246
247 fn reset_training(&mut self) {
248 self.levels.clear();
249 self.level_index.clear();
250 }
251}
252
253#[derive(Clone, Debug)]
255pub struct ScaleColorContinuous {
256 aesthetic: Aesthetic,
257 name: String,
258 low: RGBAColor,
259 high: RGBAColor,
260 min: f64,
261 max: f64,
262}
263
264impl ScaleColorContinuous {
265 pub fn new(aesthetic: Aesthetic) -> Self {
266 ScaleColorContinuous {
267 aesthetic,
268 name: String::new(),
269 low: RGBAColor::new(0, 0, 255),
270 high: RGBAColor::new(255, 0, 0),
271 min: f64::INFINITY,
272 max: f64::NEG_INFINITY,
273 }
274 }
275
276 pub fn with_colors(mut self, low: RGBAColor, high: RGBAColor) -> Self {
277 self.low = low;
278 self.high = high;
279 self
280 }
281
282 pub fn color_at(&self, t: f64) -> RGBAColor {
283 self.low.lerp(&self.high, t)
284 }
285}
286
287impl Scale for ScaleColorContinuous {
288 fn aesthetic(&self) -> Aesthetic {
289 self.aesthetic.clone()
290 }
291
292 fn train(&mut self, values: &[Value]) {
293 for v in values {
294 if let Some(f) = v.as_f64() {
295 if f.is_finite() {
296 if f < self.min {
297 self.min = f;
298 }
299 if f > self.max {
300 self.max = f;
301 }
302 }
303 }
304 }
305 }
306
307 fn map(&self, value: &Value) -> f64 {
308 let f = match value.as_f64() {
309 Some(f) => f,
310 None => return 0.0,
311 };
312 let range = self.max - self.min;
313 if range.abs() < f64::EPSILON {
314 0.5
315 } else {
316 (f - self.min) / range
317 }
318 }
319
320 fn breaks(&self) -> Vec<(f64, String)> {
321 if self.min > self.max || !self.min.is_finite() || !self.max.is_finite() {
322 return vec![];
323 }
324
325 let range = self.max - self.min;
326 if super::util::is_degenerate_range(self.min, self.max) {
327 return vec![(0.5, format_number(self.min))];
328 }
329
330 let n_breaks = 5;
331 let raw_step = range / n_breaks as f64;
332 let step = nice_step(raw_step);
333
334 let start = (self.min / step).ceil() * step;
335 super::util::stepped_breaks(start, self.max, step)
336 .into_iter()
337 .map(|v| (self.map(&Value::Float(v)), format_number(v)))
338 .collect()
339 }
340
341 fn name(&self) -> &str {
342 &self.name
343 }
344
345 fn set_name(&mut self, name: &str) {
346 self.name = name.to_string();
347 }
348
349 fn map_to_color(&self, value: &Value) -> Option<(u8, u8, u8)> {
350 let t = self.map(value);
351 let c = self.color_at(t);
352 Some((c.r, c.g, c.b))
353 }
354
355 fn domain(&self) -> Option<(f64, f64)> {
356 if self.min.is_finite() && self.max.is_finite() && self.min <= self.max {
357 Some((self.min, self.max))
358 } else {
359 None
360 }
361 }
362
363 fn clone_box(&self) -> Box<dyn Scale> {
364 Box::new(self.clone())
365 }
366
367 fn reset_training(&mut self) {
368 self.min = f64::INFINITY;
369 self.max = f64::NEG_INFINITY;
370 }
371}