map_engine/
cmap.rs

1//! Types and functions to style a `Raster`.
2// cmaps from python
3// Eg: `matplotlib.cm.get_cmap('viridis', 7).colors`
4// Potentialy use this data: https://github.com/matplotlib/matplotlib/blob/c06e8709dde6504d396349c0c80ef019c88c3927/lib/matplotlib/_cm_listed.py
5use crate::colour::{Colour, RgbaComponents};
6use ndarray::Array;
7use palette::{
8    encoding::{Linear, Srgb},
9    rgb::Rgb,
10    Alpha, Gradient, LinSrgba,
11};
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14use std::convert::TryInto;
15
16/// A linear RGBA gradient
17pub type GradientLinearRGBA = Gradient<Alpha<Rgb<Linear<Srgb>, f64>, f64>>;
18
19const VIRIDIS7: [Colour; 7] = [
20    Colour::Seq((0.267004, 0.004874, 0.329415, 1.)),
21    Colour::Seq((0.267968, 0.223549, 0.512008, 1.)),
22    Colour::Seq((0.190631, 0.407061, 0.556089, 1.)),
23    Colour::Seq((0.127568, 0.566949, 0.550556, 1.)),
24    Colour::Seq((0.20803, 0.718701, 0.472873, 1.)),
25    Colour::Seq((0.565498, 0.84243, 0.262877, 1.)),
26    Colour::Seq((0.993248, 0.906157, 0.143936, 1.)),
27];
28
29const INFERNO7: [Colour; 7] = [
30    Colour::Seq((0.001462, 0.000466, 0.013866, 1.)),
31    Colour::Seq((0.197297, 0.0384, 0.367535, 1.)),
32    Colour::Seq((0.472328, 0.110547, 0.428334, 1.)),
33    Colour::Seq((0.735683, 0.215906, 0.330245, 1.)),
34    Colour::Seq((0.929644, 0.411479, 0.145367, 1.)),
35    Colour::Seq((0.986175, 0.713153, 0.103863, 1.)),
36    Colour::Seq((0.988362, 0.998364, 0.644924, 1.)),
37];
38
39// TODO: Is there any way to generate the documentation (link) dynamically?
40macro_rules! gen_cmap_fn {
41    ($(#[$attr:meta])* => ($name:ident, $colours:expr)) => {
42        $(#[$attr])*
43        pub fn $name( vmin: f64, vmax: f64) -> GradientLinearRGBA {
44            make_gradient(vmin, vmax, &$colours)
45        }
46    }
47}
48
49gen_cmap_fn! {
50/// Create Gradient: ![viridis](https://gitlab.com/spadarian/map_engine/-/raw/master/assets/docs/cmaps/viridis.png)
51    => (viridis, VIRIDIS7)
52}
53
54gen_cmap_fn! {
55/// Create Gradient: ![inferno](https://gitlab.com/spadarian/map_engine/-/raw/master/assets/docs/cmaps/inferno.png)
56    => (inferno, INFERNO7)
57}
58
59/// Create a colour gradient.
60///
61/// The colour space is partitioned equally.
62///
63/// # Arguments
64///
65/// * `vmin` - Lower limit (pixel value) of the gradient.
66/// * `vmax` - Upper limit (pixel value) of the gradient.
67/// * `rgba` - Sequence of [`Colour`].
68fn make_gradient(vmin: f64, vmax: f64, rgba: &[Colour]) -> GradientLinearRGBA {
69    let nums = Array::linspace(vmin, vmax, rgba.len());
70    let cols = nums
71        .iter()
72        .zip(rgba)
73        .map(|(v, comps)| {
74            (
75                *v,
76                LinSrgba::from_components(Into::<(f64, f64, f64, f64)>::into(comps.clone())),
77            )
78        })
79        .collect();
80    Gradient::with_domain(cols)
81}
82
83fn make_gradient_with_breaks(nums: &[(f64, Colour)]) -> GradientLinearRGBA {
84    let cols = nums
85        .iter()
86        .map(|(v, comps)| {
87            (
88                *v,
89                LinSrgba::from_components(Into::<(f64, f64, f64, f64)>::into(comps.clone())),
90            )
91        })
92        .collect();
93    Gradient::with_domain(cols)
94}
95
96/// Types of palettes supported.
97#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
98#[serde(untagged)]
99pub enum ColourDefinition {
100    /// A discrete palette. See [`Composite::new_discrete_palette`].
101    Discrete(Vec<(isize, Colour)>),
102    /// An equally-spaced gradient. See [`Composite::new_custom_gradient`].
103    Colours(Vec<Colour>),
104    /// A gradient with custom breaks. See [`Composite::new_gradient_with_breaks`].
105    ColoursAndBreaks(Vec<(f64, Colour)>),
106    /// A RGB composite. See [`Composite::new_rgb`].
107    RGB([f64; 3], [f64; 3]),
108}
109
110/// Object to style `RawPixels`.
111#[derive(Debug, Clone)]
112pub struct Composite {
113    vmin: Option<Vec<f64>>,
114    vmax: Option<Vec<f64>>,
115    gradient: Option<GradientLinearRGBA>,
116    hashmap: Option<HashMap<isize, RgbaComponents>>,
117    display: Option<String>,
118    colour_definition: ColourDefinition,
119    len: usize,
120}
121
122impl Default for Composite {
123    fn default() -> Self {
124        let grad: Gradient<Alpha<Rgb<Linear<Srgb>, f64>, f64>> = viridis(0.0, 1.0);
125        Self {
126            vmin: Some(vec![0.0]),
127            vmax: Some(vec![1.0]),
128            gradient: Some(grad),
129            hashmap: None,
130            display: Some("Gradient".to_string()),
131            colour_definition: ColourDefinition::Colours(vec![
132                (0.0, 0.0, 0.0, 0.0).into(),
133                (1.0, 1.0, 1.0, 1.0).into(),
134            ]),
135            len: 1,
136        }
137    }
138}
139
140impl std::fmt::Display for Composite {
141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
142        write!(f, "{}", self.display.as_ref().unwrap_or(&"".to_string()))
143    }
144}
145
146impl Composite {
147    /// Create a RGB `Composite` that maps 3 pixel values (from 3 different bands) into RGBA
148    /// components.
149    ///
150    /// # Example
151    /// ```
152    /// use map_engine::cmap::{Composite, HandleGet, viridis};
153    /// let comp = Composite::new_rgb(vec![0.0, 0.0, 0.0], vec![100.0, 100.0, 100.0]);
154    /// assert_eq!(comp.get(&[0.0, 50.0, 100.0], None), [0, 127, 255, 255]);
155    /// ```
156    pub fn new_rgb(vmin: Vec<f64>, vmax: Vec<f64>) -> Self {
157        Self {
158            vmin: Some(vmin.clone()),
159            vmax: Some(vmax.clone()),
160            display: Some("RGBComposite".to_string()),
161            gradient: None,
162            colour_definition: ColourDefinition::RGB(
163                vmin.try_into().unwrap_or_else(|v: Vec<_>| {
164                    panic!("Expected a Vec of length {} but it was {}", 3, v.len())
165                }),
166                vmax.try_into().unwrap_or_else(|v: Vec<_>| {
167                    panic!("Expected a Vec of length {} but it was {}", 3, v.len())
168                }),
169            ),
170            len: 3,
171            ..Default::default()
172        }
173    }
174
175    /// Create a `Composite` that maps 1 pixel value into RGBA using the provided function.
176    ///
177    /// You can use [one of the functions provided in the cmap module](/.#functions).
178    ///
179    /// # Example
180    /// ```
181    /// use map_engine::cmap::{Composite, HandleGet, viridis};
182    /// let comp = Composite::new_gradient(0.0, 100.0, &viridis);
183    /// assert_eq!(comp.get(&[0.0], None), [68, 1, 84, 255]);
184    /// ```
185    pub fn new_gradient(
186        vmin: f64,
187        vmax: f64,
188        cmap_f: &'static dyn Fn(f64, f64) -> GradientLinearRGBA,
189    ) -> Self {
190        let grad = cmap_f(vmin, vmax);
191        Self {
192            vmin: Some(vec![vmin]),
193            vmax: Some(vec![vmax]),
194            gradient: Some(grad),
195            display: Some("Gradient".to_string()),
196            len: 1,
197            ..Default::default()
198        }
199    }
200
201    /// Create an equally-spaced `Composite` that maps 1 pixel value into RGBA using a sequence of [`Colour`].
202    ///
203    /// # Example
204    /// ```
205    /// use map_engine::{
206    ///     colour::Colour,
207    ///     cmap::{Composite, HandleGet},
208    /// };
209    /// let comp = Composite::new_custom_gradient(0.0, 100.0, vec![
210    ///     Colour::from((255, 0, 0, 255)), // red
211    ///     Colour::from((0, 0, 255, 255)), // blue
212    /// ]);
213    /// assert_eq!(comp.get(&[50.0], None), [127, 0, 127, 255]); // purple
214    /// ```
215    pub fn new_custom_gradient(vmin: f64, vmax: f64, colours: Vec<Colour>) -> Self {
216        let grad = make_gradient(vmin, vmax, &colours);
217        Self {
218            vmin: Some(vec![vmin]),
219            vmax: Some(vec![vmax]),
220            gradient: Some(grad),
221            display: Some("Gradient".to_string()),
222            colour_definition: ColourDefinition::Colours(colours),
223            len: 1,
224            ..Default::default()
225        }
226    }
227
228    /// Create an `Composite` with custom breaks that maps 1 pixel value into RGBA.
229    ///
230    /// # Example
231    /// ```
232    /// use map_engine::{
233    ///     colour::Colour,
234    ///     cmap::{Composite, HandleGet},
235    /// };
236    /// let comp = Composite::new_gradient_with_breaks(vec![
237    ///     (0.0, Colour::from((255, 0, 0, 255))), // red
238    ///     (25.0, Colour::from((127, 0, 127, 255))), // purple shifted to the red
239    ///     (100.0, Colour::from((0, 0, 255, 255))), // blue
240    /// ]);
241    /// assert_eq!(comp.get(&[25.0], None), [127, 0, 127, 255]); // purple
242    /// ```
243    pub fn new_gradient_with_breaks(cols_and_breaks: Vec<(f64, Colour)>) -> Self {
244        let grad = make_gradient_with_breaks(&cols_and_breaks);
245        Self {
246            gradient: Some(grad),
247            display: Some("GradientWithBreaks".to_string()),
248            colour_definition: ColourDefinition::ColoursAndBreaks(cols_and_breaks),
249            len: 1,
250            ..Default::default()
251        }
252    }
253
254    /// Create an discrete `Composite` that maps 1 pixel value into RGBA.
255    ///
256    /// # Example
257    /// ```
258    /// use map_engine::{
259    ///     colour::Colour,
260    ///     cmap::{Composite, HandleGet},
261    /// };
262    /// let comp = Composite::new_discrete_palette(vec![
263    ///     (0, Colour::from((255, 0, 0, 255))), // red
264    ///     (1, Colour::from((0, 255, 0, 255))), // green
265    ///     (2, Colour::from((0, 0, 255, 255))), // blue
266    /// ]);
267    /// assert_eq!(comp.get(&[0.0], None), [255, 0, 0, 255]); // red
268    /// assert_eq!(comp.get(&[3.0], None), [0, 0, 0, 0]); // transparent if not defined
269    /// ```
270    pub fn new_discrete_palette(cols_and_breaks: Vec<(isize, Colour)>) -> Self {
271        let hashmap = cols_and_breaks
272            .clone()
273            .into_iter()
274            .map(|(b, c)| (b, c.into()))
275            .collect();
276
277        Self {
278            display: Some("DiscretePalette".to_string()),
279            hashmap: Some(hashmap),
280            colour_definition: ColourDefinition::Discrete(cols_and_breaks),
281            len: 1,
282            ..Default::default()
283        }
284    }
285
286    pub(crate) fn is_contiguous(&self) -> bool {
287        !matches!(self.colour_definition, ColourDefinition::RGB(_, _))
288    }
289
290    /// Number of bands supported by the `Composite`.
291    ///
292    /// ⚠ This will probably be deprecated once we enforce the number of bands using the type
293    /// system ⚠
294    pub fn n_bands(&self) -> usize {
295        self.len
296    }
297}
298
299fn gradient_handle(comp: &Composite, values: &[f64], no_data_values: Option<&[f64]>) -> [u8; 4] {
300    let grad = comp.gradient.as_ref().unwrap();
301    let col = grad.get(values[0]);
302    let (r, g, b, a) = col.into_components();
303    let a = if let Some(ndv) = no_data_values {
304        assert!(
305            ndv.len() == 1,
306            "To use a {} style you need to provide 1 `no_data` value",
307            comp
308        );
309        if (values[0] - ndv[0]).abs() < f64::EPSILON {
310            0u8
311        } else {
312            (a * 255.0) as u8
313        }
314    } else {
315        (a * 255.0) as u8
316    };
317    [(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8, a]
318}
319
320fn rgb_handle(comp: &Composite, values: &[f64], no_data_values: Option<&[f64]>) -> [u8; 4] {
321    let norm: Vec<f64> = values
322        .iter()
323        .enumerate()
324        .map(|(i, v)| {
325            if v > &comp.vmax.as_ref().unwrap()[i] {
326                1.0
327            } else if v < &comp.vmin.as_ref().unwrap()[i] {
328                0.0
329            } else {
330                (v - comp.vmin.as_ref().unwrap()[i])
331                    / (comp.vmax.as_ref().unwrap()[i] - comp.vmin.as_ref().unwrap()[i])
332            }
333        })
334        .collect();
335    let (r, g, b, a) = (norm[0], norm[1], norm[2], 1.0);
336    let a = if let Some(no_data_values) = no_data_values {
337        assert!(
338            no_data_values.len() == 3,
339            "To use a {} style you need to provide 3 `no_data` values",
340            comp
341        );
342        if no_data_values
343            .iter()
344            .zip(values)
345            .any(|(ndv, v)| (v - ndv).abs() < f64::EPSILON)
346        {
347            0u8
348        } else {
349            (a * 255.0) as u8
350        }
351    } else {
352        (a * 255.0) as u8
353    };
354    [(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8, a]
355}
356
357fn hashmap_handle(comp: &Composite, values: &[f64]) -> [u8; 4] {
358    let val = values[0];
359    let hash = comp.hashmap.as_ref().unwrap();
360    let (r, g, b, a) = hash
361        .get(&(val.trunc() as isize))
362        .unwrap_or(&(0.0, 0.0, 0.0, 0.0));
363    [
364        (r * 255.0) as u8,
365        (g * 255.0) as u8,
366        (b * 255.0) as u8,
367        (a * 255.0) as u8,
368    ]
369}
370
371/// Get a RGBA colour given a raw pixel value
372pub trait HandleGet {
373    /// Get a RGBA colour given a raw pixel value.
374    ///
375    /// The length of `values` might vary, usually 1 or 3, depending on the `ColourDefinition`
376    /// contained within [`Composite`].
377    ///
378    /// ⚠ This will probably change once we enforce the length using the type system ⚠
379    fn get(&self, values: &[f64], no_data_values: Option<&[f64]>) -> [u8; 4];
380}
381
382impl HandleGet for Composite {
383    fn get(&self, values: &[f64], no_data_values: std::option::Option<&[f64]>) -> [u8; 4] {
384        match &self.colour_definition {
385            ColourDefinition::Discrete(_) => hashmap_handle(self, values),
386            ColourDefinition::Colours(_) | ColourDefinition::ColoursAndBreaks(_) => {
387                gradient_handle(self, values, no_data_values)
388            }
389            ColourDefinition::RGB(_, _) => rgb_handle(self, values, no_data_values),
390        }
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397
398    macro_rules! test_cmap {
399        ($test_name:ident, $fn:expr, $colours:expr) => {
400            #[test]
401            fn $test_name() {
402                let len = $colours.len();
403                let grad = $fn(0.0, (len - 1) as f64);
404                let expected = $colours;
405                let from_grad: Vec<Colour> = (0..len)
406                    .into_iter()
407                    .map(|i| grad.get(i as f64).into())
408                    .collect();
409                assert_eq!(from_grad, expected)
410            }
411        };
412    }
413
414    test_cmap!(test_viridis, viridis, VIRIDIS7);
415    test_cmap!(test_inferno, inferno, INFERNO7);
416
417    #[test]
418    fn test_colour_definition_is_deserialized() {
419        let expected_col_def = ColourDefinition::Colours(vec![
420            (1., 0., 0., 1.).into(),
421            (0., 1., 0., 1.).into(),
422            (0., 0., 1., 1.).into(),
423        ]);
424
425        let s = r#"[
426        [1.0, 0.0, 0.0, 1.0],
427        [0.0, 1.0, 0.0, 1.0],
428        [0.0, 0.0, 1.0, 1.0]
429        ]"#;
430        let col_def: ColourDefinition = serde_json::from_str(s).unwrap();
431        assert_eq!(col_def, expected_col_def);
432        let s = r#"[
433        "ff0000ff",
434        "00ff00ff",
435        "0000ffff"
436        ]"#;
437        let col_def: ColourDefinition = serde_json::from_str(s).unwrap();
438        assert_eq!(col_def, expected_col_def);
439    }
440}
map_engine/cmap.rs

map_engine/
cmap.rs
