Skip to main content

flow_plots/
colormap.rs

1use colorgrad::Gradient;
2use plotters::{
3    prelude::*,
4    style::colors::colormaps::{BlackWhite, Bone, MandelbrotHSL, ViridisRGB, VulcanoHSL},
5};
6
7/// Color map options for density plots
8///
9/// This enum provides access to a wide variety of colormaps suitable for
10/// scientific data visualization. Colormaps are categorized into:
11///
12/// - **Perceptually uniform sequential**: Viridis, Plasma, Inferno, Magma, Turbo
13///   (excellent for continuous data, colorblind-friendly)
14/// - **Traditional**: Rainbow, Jet (colorful but less perceptually uniform)
15/// - **Grayscale**: Bone, BlackWhite (useful for printing)
16/// - **Specialized**: Mandelbrot, Volcano (artistic/experimental)
17///
18/// # Recommendations
19///
20/// - **Default choice**: `Viridis` - perceptually uniform, colorblind-friendly
21/// - **High contrast**: `Plasma`, `Inferno`, `Magma` - good for presentations
22/// - **Traditional**: `Rainbow`, `Jet` - colorful but use with caution
23/// - **Print-friendly**: `Bone`, `BlackWhite` - grayscale options
24pub enum ColorMaps {
25    // Perceptually uniform colormaps (from colorgrad)
26    /// Viridis - perceptually uniform, colorblind-friendly (default)
27    Viridis,
28    /// Plasma - perceptually uniform, high contrast
29    Plasma,
30    /// Inferno - perceptually uniform, dark background friendly
31    Inferno,
32    /// Magma - perceptually uniform, dark to bright
33    Magma,
34    /// Turbo - perceptually uniform, vibrant colors
35    Turbo,
36    /// Cividis - colorblind-friendly, optimized for printing
37    Cividis,
38    /// Warm - warm color palette
39    Warm,
40    /// Cool - cool color palette
41    Cool,
42    /// Cubehelix - perceptually uniform, customizable
43    CubehelixDefault,
44
45    // Traditional colormaps (from colorgrad)
46    /// Rainbow - traditional rainbow colors (use with caution)
47    Rainbow,
48    /// Jet - traditional jet colormap (use with caution)
49    Jet,
50    /// Spectral - diverging colormap
51    Spectral,
52
53    // Plotters built-in colormaps (kept for backward compatibility)
54    /// Bone - grayscale colormap
55    Bone(Bone),
56    /// Mandelbrot - artistic HSL colormap
57    Mandelbrot(MandelbrotHSL),
58    /// BlackWhite - simple grayscale
59    BlackWhite(BlackWhite),
60    /// Volcano - HSL colormap
61    Volcano(VulcanoHSL),
62    /// ViridisRGB - Plotters' Viridis implementation (use Viridis instead)
63    #[deprecated(note = "Use ColorMaps::Viridis instead")]
64    ViridisRGB(ViridisRGB),
65}
66
67impl Clone for ColorMaps {
68    fn clone(&self) -> Self {
69        match self {
70            // colorgrad colormaps are zero-sized, so we can just copy the variant
71            ColorMaps::Viridis => ColorMaps::Viridis,
72            ColorMaps::Plasma => ColorMaps::Plasma,
73            ColorMaps::Inferno => ColorMaps::Inferno,
74            ColorMaps::Magma => ColorMaps::Magma,
75            ColorMaps::Turbo => ColorMaps::Turbo,
76            ColorMaps::Cividis => ColorMaps::Cividis,
77            ColorMaps::Warm => ColorMaps::Warm,
78            ColorMaps::Cool => ColorMaps::Cool,
79            ColorMaps::CubehelixDefault => ColorMaps::CubehelixDefault,
80            ColorMaps::Rainbow => ColorMaps::Rainbow,
81            ColorMaps::Jet => ColorMaps::Jet,
82            ColorMaps::Spectral => ColorMaps::Spectral,
83            // Plotters colormaps
84            ColorMaps::Bone(_) => ColorMaps::Bone(Bone),
85            ColorMaps::Mandelbrot(_) => ColorMaps::Mandelbrot(MandelbrotHSL),
86            ColorMaps::BlackWhite(_) => ColorMaps::BlackWhite(BlackWhite),
87            ColorMaps::Volcano(_) => ColorMaps::Volcano(VulcanoHSL),
88            #[allow(deprecated)]
89            ColorMaps::ViridisRGB(_) => ColorMaps::Viridis,
90        }
91    }
92}
93impl ColorMaps {
94    /// Map a normalized value (0.0 to 1.0) to an RGB color
95    ///
96    /// # Arguments
97    /// * `value` - Normalized density value between 0.0 and 1.0
98    ///
99    /// # Returns
100    /// An RGB color as `RGBColor(r, g, b)` where each component is 0-255
101    pub fn map(&self, value: f32) -> RGBColor {
102        // Clamp value to [0.0, 1.0]
103        let clamped_value = value.max(0.0).min(1.0);
104
105        match self {
106            // colorgrad colormaps (from preset module)
107            // Note: colorgrad Color has r, g, b, a as f32 in range [0.0, 1.0]
108            ColorMaps::Viridis => {
109                let grad = colorgrad::preset::viridis();
110                let color = grad.at(clamped_value);
111                RGBColor(
112                    (color.r * 255.0) as u8,
113                    (color.g * 255.0) as u8,
114                    (color.b * 255.0) as u8,
115                )
116            }
117            ColorMaps::Plasma => {
118                let grad = colorgrad::preset::plasma();
119                let color = grad.at(clamped_value);
120                RGBColor(
121                    (color.r * 255.0) as u8,
122                    (color.g * 255.0) as u8,
123                    (color.b * 255.0) as u8,
124                )
125            }
126            ColorMaps::Inferno => {
127                let grad = colorgrad::preset::inferno();
128                let color = grad.at(clamped_value);
129                RGBColor(
130                    (color.r * 255.0) as u8,
131                    (color.g * 255.0) as u8,
132                    (color.b * 255.0) as u8,
133                )
134            }
135            ColorMaps::Magma => {
136                let grad = colorgrad::preset::magma();
137                let color = grad.at(clamped_value);
138                RGBColor(
139                    (color.r * 255.0) as u8,
140                    (color.g * 255.0) as u8,
141                    (color.b * 255.0) as u8,
142                )
143            }
144            ColorMaps::Turbo => {
145                let grad = colorgrad::preset::turbo();
146                let color = grad.at(clamped_value);
147                RGBColor(
148                    (color.r * 255.0) as u8,
149                    (color.g * 255.0) as u8,
150                    (color.b * 255.0) as u8,
151                )
152            }
153            ColorMaps::Cividis => {
154                let grad = colorgrad::preset::cividis();
155                let color = grad.at(clamped_value);
156                RGBColor(
157                    (color.r * 255.0) as u8,
158                    (color.g * 255.0) as u8,
159                    (color.b * 255.0) as u8,
160                )
161            }
162            ColorMaps::Warm => {
163                let grad = colorgrad::preset::warm();
164                let color = grad.at(clamped_value);
165                RGBColor(
166                    (color.r * 255.0) as u8,
167                    (color.g * 255.0) as u8,
168                    (color.b * 255.0) as u8,
169                )
170            }
171            ColorMaps::Cool => {
172                let grad = colorgrad::preset::cool();
173                let color = grad.at(clamped_value);
174                RGBColor(
175                    (color.r * 255.0) as u8,
176                    (color.g * 255.0) as u8,
177                    (color.b * 255.0) as u8,
178                )
179            }
180            ColorMaps::CubehelixDefault => {
181                let grad = colorgrad::preset::cubehelix_default();
182                let color = grad.at(clamped_value);
183                RGBColor(
184                    (color.r * 255.0) as u8,
185                    (color.g * 255.0) as u8,
186                    (color.b * 255.0) as u8,
187                )
188            }
189            ColorMaps::Rainbow => {
190                let grad = colorgrad::preset::rainbow();
191                let color = grad.at(clamped_value);
192                RGBColor(
193                    (color.r * 255.0) as u8,
194                    (color.g * 255.0) as u8,
195                    (color.b * 255.0) as u8,
196                )
197            }
198            ColorMaps::Jet => {
199                // colorgrad doesn't have Jet, use sinebow as a similar alternative
200                let grad = colorgrad::preset::sinebow();
201                let color = grad.at(clamped_value);
202                RGBColor(
203                    (color.r * 255.0) as u8,
204                    (color.g * 255.0) as u8,
205                    (color.b * 255.0) as u8,
206                )
207            }
208            ColorMaps::Spectral => {
209                let grad = colorgrad::preset::spectral();
210                let color = grad.at(clamped_value);
211                RGBColor(
212                    (color.r * 255.0) as u8,
213                    (color.g * 255.0) as u8,
214                    (color.b * 255.0) as u8,
215                )
216            }
217            // Plotters built-in colormaps (backward compatibility)
218            ColorMaps::Bone(c) => c.get_color(clamped_value),
219            ColorMaps::Mandelbrot(c) => convert_hsl_to_rgb(c.get_color(clamped_value)),
220            ColorMaps::BlackWhite(c) => c.get_color(clamped_value),
221            ColorMaps::Volcano(c) => convert_hsl_to_rgb(c.get_color(clamped_value)),
222            #[allow(deprecated)]
223            ColorMaps::ViridisRGB(c) => c.get_color(clamped_value),
224        }
225    }
226}
227impl std::fmt::Debug for ColorMaps {
228    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
229        match self {
230            ColorMaps::Viridis => write!(f, "Viridis"),
231            ColorMaps::Plasma => write!(f, "Plasma"),
232            ColorMaps::Inferno => write!(f, "Inferno"),
233            ColorMaps::Magma => write!(f, "Magma"),
234            ColorMaps::Turbo => write!(f, "Turbo"),
235            ColorMaps::Cividis => write!(f, "Cividis"),
236            ColorMaps::Warm => write!(f, "Warm"),
237            ColorMaps::Cool => write!(f, "Cool"),
238            ColorMaps::CubehelixDefault => write!(f, "CubehelixDefault"),
239            ColorMaps::Rainbow => write!(f, "Rainbow"),
240            ColorMaps::Jet => write!(f, "Jet"),
241            ColorMaps::Spectral => write!(f, "Spectral"),
242            ColorMaps::Bone(_) => write!(f, "Bone"),
243            ColorMaps::Mandelbrot(_) => write!(f, "Mandelbrot"),
244            ColorMaps::BlackWhite(_) => write!(f, "BlackWhite"),
245            ColorMaps::Volcano(_) => write!(f, "Volcano"),
246            #[allow(deprecated)]
247            ColorMaps::ViridisRGB(_) => write!(f, "ViridisRGB"),
248        }
249    }
250}
251impl Default for ColorMaps {
252    fn default() -> Self {
253        ColorMaps::Viridis
254    }
255}
256
257fn convert_hsl_to_rgb(hsl: HSLColor) -> RGBColor {
258    let (r, g, b) = hsl.rgb();
259    RGBColor(r, g, b)
260}
261
262// Define your custom color map
263pub struct CustomColorMap;
264
265macro_rules! def_linear_colormap{
266    ($color_scale_name:ident, $color_type:ident, $doc:expr, $(($($color_value:expr),+)),*) => {
267        #[doc = $doc]
268        pub struct $color_scale_name;
269
270        impl $color_scale_name {
271            // const COLORS: [$color_type; $number_colors] = [$($color_type($($color_value),+)),+];
272            // const COLORS: [$color_type; $crate::count!($(($($color_value:expr),+))*)] = [$($color_type($($color_value),+)),+];
273            const COLORS: [$color_type; $crate::count!($(($($color_value:expr),+))*)] = $crate::define_colors_from_list_of_values_or_directly!{$color_type, $(($($color_value),+)),*};
274        }
275
276        $crate::implement_linear_interpolation_color_map!{$color_scale_name, $color_type}
277    };
278    ($color_scale_name:ident, $color_type:ident, $doc:expr, $($color_complete:tt),+) => {
279        #[doc = $doc]
280        pub struct $color_scale_name;
281
282        impl $color_scale_name {
283            const COLORS: [$color_type; $crate::count!($($color_complete)*)] = $crate::define_colors_from_list_of_values_or_directly!{$($color_complete),+};
284        }
285
286        $crate::implement_linear_interpolation_color_map!{$color_scale_name, $color_type}
287    }
288}
289
290#[macro_export]
291#[doc(hidden)]
292/// Implements the [ColorMap] trait on a given color scale.
293macro_rules! implement_linear_interpolation_color_map {
294    ($color_scale_name:ident, $color_type:ident) => {
295        impl<FloatType: std::fmt::Debug + num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive>
296            ColorMap<$color_type, FloatType> for $color_scale_name
297        {
298            fn get_color_normalized(
299                &self,
300                h: FloatType,
301                min: FloatType,
302                max: FloatType,
303            ) -> $color_type {
304                let (
305                    relative_difference,
306                    index_lower,
307                    index_upper
308                ) = calculate_relative_difference_index_lower_upper(
309                    h,
310                    min,
311                    max,
312                    Self::COLORS.len()
313                );
314                // Interpolate the final color linearly
315                $crate::calculate_new_color_value!(
316                    relative_difference,
317                    Self::COLORS,
318                    index_upper,
319                    index_lower,
320                    $color_type
321                )
322            }
323        }
324
325        impl $color_scale_name {
326            #[doc = "Get color value from `"]
327            #[doc = stringify!($color_scale_name)]
328            #[doc = "` by supplying a parameter 0.0 <= h <= 1.0"]
329            pub fn get_color<FloatType: std::fmt::Debug + num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive>(
330                h: FloatType,
331            ) -> $color_type {
332                let color_scale = $color_scale_name {};
333                color_scale.get_color(h)
334            }
335
336            #[doc = "Get color value from `"]
337            #[doc = stringify!($color_scale_name)]
338            #[doc = "` by supplying lower and upper bounds min, max and a parameter h where min <= h <= max"]
339            pub fn get_color_normalized<
340                FloatType: std::fmt::Debug + num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive,
341            >(
342                h: FloatType,
343                min: FloatType,
344                max: FloatType,
345            ) -> $color_type {
346                let color_scale = $color_scale_name {};
347                color_scale.get_color_normalized(h, min, max)
348            }
349        }
350    };
351}
352
353pub fn calculate_relative_difference_index_lower_upper<
354    FloatType: num_traits::Float + num_traits::FromPrimitive + num_traits::ToPrimitive,
355>(
356    h: FloatType,
357    min: FloatType,
358    max: FloatType,
359    n_steps: usize,
360) -> (FloatType, usize, usize) {
361    // Ensure that we do have a value in bounds
362    let h = num_traits::clamp(h, min, max);
363    // Next calculate a normalized value between 0.0 and 1.0
364    let t = (h - min) / (max - min);
365    let approximate_index = t
366        * (FloatType::from_usize(n_steps).expect("should be able to get a float type from usize")
367            - FloatType::one())
368        .max(FloatType::zero());
369    // Calculate which index are the two most nearest of the supplied value
370    let index_lower = approximate_index
371        .floor()
372        .to_usize()
373        .expect("should be able to get the lower index");
374    let index_upper = approximate_index
375        .ceil()
376        .to_usize()
377        .expect("should be able to get the upper index");
378    // Calculate the relative difference, ie. is the actual value more towards the color of index_upper or index_lower?
379    let relative_difference = approximate_index.ceil() - approximate_index;
380    (relative_difference, index_lower, index_upper)
381}
382
383/// Converts a given color identifier and a sequence of colors to an array of them.
384macro_rules! define_colors_from_list_of_values_or_directly{
385    ($color_type:ident, $(($($color_value:expr),+)),+) => {
386        [$($color_type($($color_value),+)),+]
387    };
388    ($($color_complete:tt),+) => {
389        [$($color_complete),+]
390    };
391}