From a8f91a65a7155699373f2f5a65963e4b6ec1f19e Mon Sep 17 00:00:00 2001
From: 4bb4 <67376761+4bb4@users.noreply.github.com>
Date: Sun, 27 Sep 2020 20:21:30 +0200
Subject: [PATCH] Added colormap support

---
 README.md                              |  4 +-
 implot-examples/examples/line_plots.rs | 59 +++++++++++++++++++++++++-
 src/lib.rs                             | 47 +++++++++++++++++++-
 3 files changed, 105 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index c2d342e..b035523 100644
--- a/README.md
+++ b/README.md
@@ -67,11 +67,11 @@ See below for an overview of the progress.
   - [ ] Heatmap
   - [ ] Pie chart
   - [ ] Digital data
-- [ ] Plot customization
+- [x] Plot customization
   - [x] Axis flags
   - [x] Styling colors
   - [x] Styling variables
-  - [ ] Colormaps
+  - [x] Colormaps
 - [ ] Plot querying 
   - [x] is hovered
   - [x] mouse position in plot
diff --git a/implot-examples/examples/line_plots.rs b/implot-examples/examples/line_plots.rs
index a126250..ce4808b 100644
--- a/implot-examples/examples/line_plots.rs
+++ b/implot-examples/examples/line_plots.rs
@@ -4,8 +4,9 @@
 use imgui::{im_str, CollapsingHeader, Condition, Ui, Window};
 use implot::{
     get_plot_limits, get_plot_mouse_position, get_plot_query, is_plot_hovered, is_plot_queried,
-    push_style_color, push_style_var_f32, push_style_var_u32, AxisFlags, ImPlotLimits, ImPlotPoint,
-    ImPlotRange, Marker, Plot, PlotColorElement, PlotFlags, PlotLine, StyleVar,
+    push_style_color, push_style_var_f32, push_style_var_u32, set_colormap_from_preset,
+    set_colormap_from_vec, AxisFlags, Colormap, ImPlotLimits, ImPlotPoint, ImPlotRange, ImVec4,
+    Marker, Plot, PlotColorElement, PlotFlags, PlotLine, StyleVar,
 };
 
 mod support;
@@ -177,6 +178,57 @@ fn show_style_plot(ui: &Ui) {
     style.pop();
 }
 
+fn show_colormaps_plot(ui: &Ui) {
+    ui.text(im_str!("This header demos how to select colormaps."));
+    let content_width = ui.window_content_region_width();
+
+    // Select a colormap from the presets. The presets are listed in the Colormap enum
+    // and usually have something from 9 to 11 colors in them, with the second number
+    // being the option to resample the colormap to a custom number of colors if picked
+    // higher than 1.
+    set_colormap_from_preset(Colormap::Plasma, 1);
+
+    Plot::new("Colormap demo plot")
+        .size(content_width, 300.0)
+        .build(|| {
+            (1..10)
+                .map(|x| x as f64 * 0.1)
+                .map(|x| PlotLine::new(&format!("{:3.3}", x)).plot(&vec![0.1, 0.9], &vec![x, x]))
+                .count();
+        });
+
+    // One can also specify a colormap as a vector of RGBA colors. ImPlot uses ImVec4 for this,
+    // so we follow suit. Make sure to set the last number (w in ImVec4) to 1.0 to see anything -
+    // it's the alpha channel.
+    set_colormap_from_vec(vec![
+        ImVec4 {
+            x: 0.9,
+            y: 0.9,
+            z: 0.0,
+            w: 1.0,
+        },
+        ImVec4 {
+            x: 0.0,
+            y: 0.9,
+            z: 0.9,
+            w: 1.0,
+        },
+    ]);
+
+    Plot::new("Colormap demo plot #2")
+        .size(content_width, 300.0)
+        .build(|| {
+            (1..10)
+                .map(|x| x as f64 * 0.1)
+                .map(|x| PlotLine::new(&format!("{:3.3}", x)).plot(&vec![0.1, 0.9], &vec![x, x]))
+                .count();
+        });
+
+    // Colormaps are not pushed, they are simply set, because they don't stack or anything.
+    // We can reset to the default by just setting the "Standard" preset.
+    set_colormap_from_preset(Colormap::Standard, 0);
+}
+
 fn main() {
     let system = support::init(file!());
     let mut showing_demo = false;
@@ -208,6 +260,9 @@ fn main() {
                 if CollapsingHeader::new(im_str!("Styling a plot")).build(&ui) {
                     show_style_plot(&ui);
                 }
+                if CollapsingHeader::new(im_str!("Colormap selection")).build(&ui) {
+                    show_colormaps_plot(&ui);
+                }
             });
 
         if showing_demo {
diff --git a/src/lib.rs b/src/lib.rs
index 6ec468c..97aadb1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -13,7 +13,7 @@ use std::convert::TryFrom;
 pub use sys::imgui::Condition;
 use sys::imgui::{im_str, ImString};
 // TODO(4bb4) facade-wrap these
-pub use sys::{ImPlotLimits, ImPlotPoint, ImPlotRange};
+pub use sys::{ImPlotLimits, ImPlotPoint, ImPlotRange, ImVec4};
 
 const DEFAULT_PLOT_SIZE_X: f32 = 400.0;
 const DEFAULT_PLOT_SIZE_Y: f32 = 400.0;
@@ -150,6 +150,32 @@ pub enum PlotColorElement {
     Query = sys::ImPlotCol__ImPlotCol_Query,
 }
 
+/// Colormap choice. Documentation copied from implot.h for convenience.
+#[repr(u32)]
+#[derive(Copy, Clone, Debug)]
+pub enum Colormap {
+    /// ImPlot default colormap (n=10). Called "Standard" here because Default is reserved.
+    Standard = sys::ImPlotColormap__ImPlotColormap_Default,
+    /// a.k.a. matplotlib "Set1" (n=9)
+    Dark = sys::ImPlotColormap__ImPlotColormap_Dark,
+    /// a.k.a. matplotlib "Pastel1" (n=9)
+    Pastel = sys::ImPlotColormap__ImPlotColormap_Pastel,
+    /// a.k.a. matplotlib "Paired" (n=12)
+    Paired = sys::ImPlotColormap__ImPlotColormap_Paired,
+    /// a.k.a. matplotlib "viridis" (n=11)
+    Viridis = sys::ImPlotColormap__ImPlotColormap_Viridis,
+    /// a.k.a. matplotlib "plasma" (n=11)
+    Plasma = sys::ImPlotColormap__ImPlotColormap_Plasma,
+    /// a.k.a. matplotlib/MATLAB "hot" (n=11)
+    Hot = sys::ImPlotColormap__ImPlotColormap_Hot,
+    /// a.k.a. matplotlib/MATLAB "cool" (n=11)
+    Cool = sys::ImPlotColormap__ImPlotColormap_Cool,
+    /// a.k.a. matplotlib/MATLAB "pink" (n=11)
+    Pink = sys::ImPlotColormap__ImPlotColormap_Pink,
+    /// a.k.a. MATLAB "jet" (n=11)
+    Jet = sys::ImPlotColormap__ImPlotColormap_Jet,
+}
+
 #[repr(u32)]
 #[derive(Copy, Clone, Debug)]
 pub enum StyleVar {
@@ -746,6 +772,25 @@ impl PlotText {
     }
 }
 
+// --- Color maps -----------------------------------------------------------------------------
+/// Switch to one of the built-in preset colormaps. If samples is greater than 1, the map will be
+/// linearly resampled.
+pub fn set_colormap_from_preset(preset: Colormap, samples: u32) {
+    unsafe {
+        // "as" casts saturate as of Rust 1.45. This is safe here, and at least the enum
+        // values are not expected to go outside the range of an i32 anyway, so there is no
+        // risk of changed values.
+        sys::ImPlot_SetColormapPlotColormap(preset as i32, samples as i32);
+    }
+}
+
+/// Set a custom colormap in the form of a vector of colors.
+pub fn set_colormap_from_vec(colors: Vec<ImVec4>) {
+    unsafe {
+        sys::ImPlot_SetColormapVec4Ptr(colors.as_ptr(), colors.len() as i32);
+    }
+}
+
 // --- Push/pop utils -------------------------------------------------------------------------
 // Currently not in a struct yet. imgui-rs has some smarts about dealing with stacks, in particular
 // leak detection, which I'd like to replicate here at some point.