fix(ColorMap): allow user to adjust number of colors

Author: jourdain

src/e3sm_quickview/components/view.py

@@ -1,4 +1,5 @@
-from trame.widgets import html, vuetify3 as v3
+from trame.widgets import html
+from trame.widgets import vuetify3 as v3
 
 
 def create_size_menu(name, config):
@@ -173,6 +174,7 @@ def create_bottom_bar(config, update_color_preset):
                                 text="Use data range",
                                 variant="text",
                             )
+
                             v3.VTextField(
                                 v_model="config.search",
                                 clearable=True,
@@ -194,7 +196,22 @@ def create_bottom_bar(config, update_color_preset):
                                 click="config.menu=false",
                             )
 
-                    with v3.VCardItem(v_show="config.override_range", classes="py-0"):
+                    with v3.VCardItem(classes="py-0 mb-2"):
+                        v3.VNumberInput(
+                            v_model="config.n_colors",
+                            hide_details=True,
+                            density="compact",
+                            variant="outlined",
+                            flat=True,
+                            label="Number of colors",
+                            classes="mt-2",
+                            step=[1],
+                            min=[2],
+                            max=[255],
+                        )
+                    with v3.VCardItem(
+                        v_show="config.override_range", classes="py-0 mb-2"
+                    ):
                         v3.VTextField(
                             v_model="config.color_value_min",
                             hide_details=True,
@@ -224,7 +241,7 @@ def create_bottom_bar(config, update_color_preset):
                             subtitle=("entry.name",),
                             click=(
                                 update_color_preset,
-                                "[entry.name, config.invert, config.use_log_scale]",
+                                "[entry.name, config.invert, config.use_log_scale, config.n_colors]",
                             ),
                             active=("config.preset === entry.name",),
                         ):
@@ -239,7 +256,7 @@ def create_bottom_bar(config, update_color_preset):
             )
             with html.Div(classes="overflow-hidden rounded w-100", style="height:70%;"):
                 html.Img(
-                    src=("config.preset_img",),
+                    src=("config.lut_img",),
                     style="width:100%;height:2rem;",
                     draggable=False,
                 )

src/e3sm_quickview/utils/color.py

@@ -85,3 +85,36 @@ def get_cached_colorbar_image(colormap_name, inverted=False):
         return COLORBAR_CACHE[colormap_name].get(variant, "")
 
     return ""
+
+
+def lut_to_img(lut_proxy):
+    samples = 255
+    rgb = [0, 0, 0]
+    vtk_lut = lut_proxy.GetClientSideObject()
+    colorArray = vtkUnsignedCharArray()
+    colorArray.SetNumberOfComponents(3)
+    colorArray.SetNumberOfTuples(samples)
+    imgData = vtkImageData()
+    imgData.SetDimensions(samples, 1, 1)
+    imgData.GetPointData().SetScalars(colorArray)
+    writer = vtkPNGWriter()
+    writer.WriteToMemoryOn()
+    writer.SetInputData(imgData)
+    writer.SetCompressionLevel(1)
+
+    v_min = lut_proxy.RGBPoints[0]
+    v_max = lut_proxy.RGBPoints[-4]
+    step = (v_max - v_min) / (samples - 1)
+
+    for i in range(samples):
+        value = v_min + step * float(i)
+        vtk_lut.GetColor(value, rgb)
+        r = int(round(rgb[0] * 255))
+        g = int(round(rgb[1] * 255))
+        b = int(round(rgb[2] * 255))
+        colorArray.SetTuple3(i, r, g, b)
+
+    writer.Write()
+    base64_img = base64.standard_b64encode(writer.GetResult()).decode("utf-8")
+
+    return f"data:image/png;base64,{base64_img}"

src/e3sm_quickview/view_manager.py

@@ -11,7 +11,7 @@
 
 from e3sm_quickview.components import view as tview
 from e3sm_quickview.presets import COLOR_BLIND_SAFE
-from e3sm_quickview.utils.color import COLORBAR_CACHE, get_cached_colorbar_image
+from e3sm_quickview.utils.color import COLORBAR_CACHE, lut_to_img
 
 
 def auto_size_to_col(size):
@@ -49,7 +49,6 @@ def lut_name(element):
 class ViewConfiguration(StateDataModel):
     variable: str
     preset: str = "Inferno (matplotlib)"
-    preset_img: str
     invert: bool = False
     color_blind: bool = False
     use_log_scale: bool = False
@@ -66,6 +65,8 @@ class ViewConfiguration(StateDataModel):
     menu: bool = False
     swap_group: list[str]
     search: str | None
+    n_colors: int = 255
+    lut_img: str
 
 
 class VariableView(TrameComponent):
@@ -127,7 +128,9 @@ def __init__(self, server, source, variable_name, variable_type):
             ["override_range", "color_range"], self.update_color_range, eager=True
         )
         self.config.watch(
-            ["preset", "invert", "use_log_scale"], self.update_color_preset, eager=True
+            ["preset", "invert", "use_log_scale", "n_colors"],
+            self.update_color_preset,
+            eager=True,
         )
 
         # GUI
@@ -151,18 +154,21 @@ def reset_camera(self):
         self.view.ResetCamera(True, 0.9)
         self.ctx[self.name].update()
 
-    def update_color_preset(self, name, invert, log_scale):
+    def update_color_preset(self, name, invert, log_scale, n_colors=255):
         self.config.preset = name
-        self.config.preset_img = get_cached_colorbar_image(
-            self.config.preset,
-            self.config.invert,
-        )
+        self.lut.UseLogScale = 0
         self.lut.ApplyPreset(self.config.preset, True)
         if invert:
             self.lut.InvertTransferFunction()
+
         if log_scale:
             self.lut.MapControlPointsToLogSpace()
             self.lut.UseLogScale = 1
+
+        self.lut.NumberOfTableValues = n_colors
+
+        self.config.lut_img = lut_to_img(self.lut)
+
         self.render()
 
     def color_range_str_to_float(self, color_value_min, color_value_max):