Support dictionary cmaps for ImageStack (#6025)

Co-authored-by: Andrew Huang <[email protected]>

holoviews/plotting/bokeh/raster.py

@@ -258,13 +258,31 @@ def _get_cmapper_opts(self, low, high, factors, colors):
 
     def _get_colormapper(self, eldim, element, ranges, style, factors=None,
                          colors=None, group=None, name='color_mapper'):
+        indices = None
+        vdims = element.vdims
+        if isinstance(style.get("cmap"), dict):
+            dict_cmap = style["cmap"]
+            missing = [vd.name for vd in vdims if vd.name not in dict_cmap]
+            if missing:
+                missing_str = "', '".join(sorted(missing))
+                raise ValueError(
+                    "The supplied cmap dictionary must have the same "
+                    f"value dimensions as the element. Missing: '{missing_str}'"
+                )
+            keys, values = zip(*dict_cmap.items())
+            style["cmap"] = list(values)
+            indices = [keys.index(vd.name) for vd in vdims]
+
         cmapper = super()._get_colormapper(
             eldim, element, ranges, style, factors=factors,
             colors=colors, group=group, name=name
         )
-        num_elements = len(element.vdims)
-        step_size = len(cmapper.palette) // num_elements
-        indices = np.arange(num_elements) * step_size
+
+        if indices is None:
+            num_elements = len(vdims)
+            step_size = len(cmapper.palette) // num_elements
+            indices = np.arange(num_elements) * step_size
+
         cmapper.palette = np.array(cmapper.palette)[indices].tolist()
         return cmapper
 

holoviews/tests/plotting/bokeh/test_rasterplot.py

@@ -408,6 +408,34 @@ def test_image_stack_tuple_single_3darray(self):
         assert source.data["dh"][0] == self.ysize
         assert isinstance(plot, ImageStackPlot)
 
+    def test_image_stack_dict_cmap(self):
+        x = np.arange(0, 3)
+        y = np.arange(5, 8)
+        a = np.array([[np.nan, np.nan, 1], [np.nan] * 3, [np.nan] * 3])
+        b = np.array([[np.nan] * 3, [1, 1, np.nan], [np.nan] * 3])
+        c = np.array([[np.nan] * 3, [np.nan] * 3, [1, 1, 1]])
+
+        img_stack = ImageStack((x, y, a, b, c), kdims=["x", "y"], vdims=["b", "a", "c"])
+        img_stack.opts(cmap={"c": "yellow", "a": "red", "b": "green"})
+        plot = bokeh_renderer.get_plot(img_stack)
+        source = plot.handles["source"]
+        np.testing.assert_equal(source.data["image"][0][:, :, 0], a)
+        np.testing.assert_equal(source.data["image"][0][:, :, 1], b)
+        np.testing.assert_equal(source.data["image"][0][:, :, 2], c)
+        assert plot.handles["color_mapper"].palette == ["green", "red", "yellow"]
+
+    def test_image_stack_dict_cmap_missing(self):
+        x = np.arange(0, 3)
+        y = np.arange(5, 8)
+        a = np.array([[np.nan, np.nan, 1], [np.nan] * 3, [np.nan] * 3])
+        b = np.array([[np.nan] * 3, [1, 1, np.nan], [np.nan] * 3])
+        c = np.array([[np.nan] * 3, [np.nan] * 3, [1, 1, 1]])
+
+        img_stack = ImageStack((x, y, a, b, c), kdims=["x", "y"], vdims=["b", "a", "c"])
+        with pytest.raises(ValueError, match="must have the same value dimensions"):
+            img_stack.opts(cmap={"c": "yellow", "a": "red"})
+            bokeh_renderer.get_plot(img_stack)
+
 
 class TestImageStackEven(_ImageStackBase):
     __test__ = True