use crate::colour::{Colour, RgbaComponents};
use ndarray::Array;
use palette::{
encoding::{Linear, Srgb},
rgb::Rgb,
Alpha, Gradient, LinSrgba,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::convert::TryInto;
pub type GradientLinearRGBA = Gradient<Alpha<Rgb<Linear<Srgb>, f64>, f64>>;
const VIRIDIS7: [Colour; 7] = [
Colour::Seq((0.267004, 0.004874, 0.329415, 1.)),
Colour::Seq((0.267968, 0.223549, 0.512008, 1.)),
Colour::Seq((0.190631, 0.407061, 0.556089, 1.)),
Colour::Seq((0.127568, 0.566949, 0.550556, 1.)),
Colour::Seq((0.20803, 0.718701, 0.472873, 1.)),
Colour::Seq((0.565498, 0.84243, 0.262877, 1.)),
Colour::Seq((0.993248, 0.906157, 0.143936, 1.)),
];
const INFERNO7: [Colour; 7] = [
Colour::Seq((0.001462, 0.000466, 0.013866, 1.)),
Colour::Seq((0.197297, 0.0384, 0.367535, 1.)),
Colour::Seq((0.472328, 0.110547, 0.428334, 1.)),
Colour::Seq((0.735683, 0.215906, 0.330245, 1.)),
Colour::Seq((0.929644, 0.411479, 0.145367, 1.)),
Colour::Seq((0.986175, 0.713153, 0.103863, 1.)),
Colour::Seq((0.988362, 0.998364, 0.644924, 1.)),
];
macro_rules! gen_cmap_fn {
($(#[$attr:meta])* => ($name:ident, $colours:expr)) => {
$(#[$attr])*
pub fn $name( vmin: f64, vmax: f64) -> GradientLinearRGBA {
make_gradient(vmin, vmax, &$colours)
}
}
}
gen_cmap_fn! {
=> (viridis, VIRIDIS7)
}
gen_cmap_fn! {
=> (inferno, INFERNO7)
}
fn make_gradient(vmin: f64, vmax: f64, rgba: &[Colour]) -> GradientLinearRGBA {
let nums = Array::linspace(vmin, vmax, rgba.len());
let cols = nums
.iter()
.zip(rgba)
.map(|(v, comps)| {
(
*v,
LinSrgba::from_components(Into::<(f64, f64, f64, f64)>::into(comps.clone())),
)
})
.collect();
Gradient::with_domain(cols)
}
fn make_gradient_with_breaks(nums: &[(f64, Colour)]) -> GradientLinearRGBA {
let cols = nums
.iter()
.map(|(v, comps)| {
(
*v,
LinSrgba::from_components(Into::<(f64, f64, f64, f64)>::into(comps.clone())),
)
})
.collect();
Gradient::with_domain(cols)
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ColourDefinition {
Discrete(Vec<(isize, Colour)>),
Colours(Vec<Colour>),
ColoursAndBreaks(Vec<(f64, Colour)>),
RGB([f64; 3], [f64; 3]),
}
#[derive(Debug, Clone)]
pub struct Composite {
vmin: Option<Vec<f64>>,
vmax: Option<Vec<f64>>,
gradient: Option<GradientLinearRGBA>,
hashmap: Option<HashMap<isize, RgbaComponents>>,
display: Option<String>,
colour_definition: ColourDefinition,
len: usize,
}
impl Default for Composite {
fn default() -> Self {
let grad: Gradient<Alpha<Rgb<Linear<Srgb>, f64>, f64>> = viridis(0.0, 1.0);
Self {
vmin: Some(vec![0.0]),
vmax: Some(vec![1.0]),
gradient: Some(grad),
hashmap: None,
display: Some("Gradient".to_string()),
colour_definition: ColourDefinition::Colours(vec![
(0.0, 0.0, 0.0, 0.0).into(),
(1.0, 1.0, 1.0, 1.0).into(),
]),
len: 1,
}
}
}
impl std::fmt::Display for Composite {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
write!(f, "{}", self.display.as_ref().unwrap_or(&"".to_string()))
}
}
impl Composite {
pub fn new_rgb(vmin: Vec<f64>, vmax: Vec<f64>) -> Self {
Self {
vmin: Some(vmin.clone()),
vmax: Some(vmax.clone()),
display: Some("RGBComposite".to_string()),
gradient: None,
colour_definition: ColourDefinition::RGB(
vmin.try_into().unwrap_or_else(|v: Vec<_>| {
panic!("Expected a Vec of length {} but it was {}", 3, v.len())
}),
vmax.try_into().unwrap_or_else(|v: Vec<_>| {
panic!("Expected a Vec of length {} but it was {}", 3, v.len())
}),
),
len: 3,
..Default::default()
}
}
pub fn new_gradient(
vmin: f64,
vmax: f64,
cmap_f: &'static dyn Fn(f64, f64) -> GradientLinearRGBA,
) -> Self {
let grad = cmap_f(vmin, vmax);
Self {
vmin: Some(vec![vmin]),
vmax: Some(vec![vmax]),
gradient: Some(grad),
display: Some("Gradient".to_string()),
len: 1,
..Default::default()
}
}
pub fn new_custom_gradient(vmin: f64, vmax: f64, colours: Vec<Colour>) -> Self {
let grad = make_gradient(vmin, vmax, &colours);
Self {
vmin: Some(vec![vmin]),
vmax: Some(vec![vmax]),
gradient: Some(grad),
display: Some("Gradient".to_string()),
colour_definition: ColourDefinition::Colours(colours),
len: 1,
..Default::default()
}
}
pub fn new_gradient_with_breaks(cols_and_breaks: Vec<(f64, Colour)>) -> Self {
let grad = make_gradient_with_breaks(&cols_and_breaks);
Self {
gradient: Some(grad),
display: Some("GradientWithBreaks".to_string()),
colour_definition: ColourDefinition::ColoursAndBreaks(cols_and_breaks),
len: 1,
..Default::default()
}
}
pub fn new_discrete_palette(cols_and_breaks: Vec<(isize, Colour)>) -> Self {
let hashmap = cols_and_breaks
.clone()
.into_iter()
.map(|(b, c)| (b, c.into()))
.collect();
Self {
display: Some("DiscretePalette".to_string()),
hashmap: Some(hashmap),
colour_definition: ColourDefinition::Discrete(cols_and_breaks),
len: 1,
..Default::default()
}
}
pub(crate) fn is_contiguous(&self) -> bool {
!matches!(self.colour_definition, ColourDefinition::RGB(_, _))
}
pub fn n_bands(&self) -> usize {
self.len
}
}
fn gradient_handle(comp: &Composite, values: &[f64], no_data_values: Option<&[f64]>) -> [u8; 4] {
let grad = comp.gradient.as_ref().unwrap();
let col = grad.get(values[0]);
let (r, g, b, a) = col.into_components();
let a = if let Some(ndv) = no_data_values {
assert!(
ndv.len() == 1,
"To use a {} style you need to provide 1 `no_data` value",
comp
);
if (values[0] - ndv[0]).abs() < f64::EPSILON {
0u8
} else {
(a * 255.0) as u8
}
} else {
(a * 255.0) as u8
};
[(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8, a]
}
fn rgb_handle(comp: &Composite, values: &[f64], no_data_values: Option<&[f64]>) -> [u8; 4] {
let norm: Vec<f64> = values
.iter()
.enumerate()
.map(|(i, v)| {
if v > &comp.vmax.as_ref().unwrap()[i] {
1.0
} else if v < &comp.vmin.as_ref().unwrap()[i] {
0.0
} else {
(v - comp.vmin.as_ref().unwrap()[i])
/ (comp.vmax.as_ref().unwrap()[i] - comp.vmin.as_ref().unwrap()[i])
}
})
.collect();
let (r, g, b, a) = (norm[0], norm[1], norm[2], 1.0);
let a = if let Some(no_data_values) = no_data_values {
assert!(
no_data_values.len() == 3,
"To use a {} style you need to provide 3 `no_data` values",
comp
);
if no_data_values
.iter()
.zip(values)
.any(|(ndv, v)| (v - ndv).abs() < f64::EPSILON)
{
0u8
} else {
(a * 255.0) as u8
}
} else {
(a * 255.0) as u8
};
[(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8, a]
}
fn hashmap_handle(comp: &Composite, values: &[f64]) -> [u8; 4] {
let val = values[0];
let hash = comp.hashmap.as_ref().unwrap();
let (r, g, b, a) = hash
.get(&(val.trunc() as isize))
.unwrap_or(&(0.0, 0.0, 0.0, 0.0));
[
(r * 255.0) as u8,
(g * 255.0) as u8,
(b * 255.0) as u8,
(a * 255.0) as u8,
]
}
pub trait HandleGet {
fn get(&self, values: &[f64], no_data_values: Option<&[f64]>) -> [u8; 4];
}
impl HandleGet for Composite {
fn get(&self, values: &[f64], no_data_values: std::option::Option<&[f64]>) -> [u8; 4] {
match &self.colour_definition {
ColourDefinition::Discrete(_) => hashmap_handle(self, values),
ColourDefinition::Colours(_) | ColourDefinition::ColoursAndBreaks(_) => {
gradient_handle(self, values, no_data_values)
}
ColourDefinition::RGB(_, _) => rgb_handle(self, values, no_data_values),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
macro_rules! test_cmap {
($test_name:ident, $fn:expr, $colours:expr) => {
#[test]
fn $test_name() {
let len = $colours.len();
let grad = $fn(0.0, (len - 1) as f64);
let expected = $colours;
let from_grad: Vec<Colour> = (0..len)
.into_iter()
.map(|i| grad.get(i as f64).into())
.collect();
assert_eq!(from_grad, expected)
}
};
}
test_cmap!(test_viridis, viridis, VIRIDIS7);
test_cmap!(test_inferno, inferno, INFERNO7);
#[test]
fn test_colour_definition_is_deserialized() {
let expected_col_def = ColourDefinition::Colours(vec![
(1., 0., 0., 1.).into(),
(0., 1., 0., 1.).into(),
(0., 0., 1., 1.).into(),
]);
let s = r#"[
[1.0, 0.0, 0.0, 1.0],
[0.0, 1.0, 0.0, 1.0],
[0.0, 0.0, 1.0, 1.0]
]"#;
let col_def: ColourDefinition = serde_json::from_str(s).unwrap();
assert_eq!(col_def, expected_col_def);
let s = r#"[
"ff0000ff",
"00ff00ff",
"0000ffff"
]"#;
let col_def: ColourDefinition = serde_json::from_str(s).unwrap();
assert_eq!(col_def, expected_col_def);
}
}