From 6c91f6de6ceab01fbea6e420f8de27544290c98b Mon Sep 17 00:00:00 2001
From: Andrew Huang <ahuang@anaconda.com>
Date: Thu, 14 Mar 2024 18:56:14 -0700
Subject: [PATCH] Cleanup and add test

---
 holoviews/plotting/bokeh/raster.py            | 17 +++++++---
 .../tests/plotting/bokeh/test_rasterplot.py   | 31 +++++++++++++++++++
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/holoviews/plotting/bokeh/raster.py b/holoviews/plotting/bokeh/raster.py
index 9909065cd0..cb0a9f6ea5 100644
--- a/holoviews/plotting/bokeh/raster.py
+++ b/holoviews/plotting/bokeh/raster.py
@@ -259,17 +259,26 @@ 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):
-            keys, values = zip(*style["cmap"].items())
-            missing = [vd.name for vd in element.vdims if vd.name not in keys]
+            dict_cmap = style["cmap"]
+            missing = [vd.name for vd in vdims if vd.name not in dict_cmap]
+            extra = [k for k in dict_cmap if k not in vdims]
             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}'"
                 )
+            elif extra:
+                extra_str = "', '".join(sorted(extra))
+                self.param.warning(
+                    f"The supplied cmap dictionary has extra value dimensions: "
+                    f"'{extra_str}'. Ignoring these value dimensions."
+                )
+            keys, values = zip(*dict_cmap.items())
             style["cmap"] = list(values)
-            indices = [keys.index(vd.name) for vd in element.vdims]
+            indices = [keys.index(vd.name) for vd in vdims]
 
         cmapper = super()._get_colormapper(
             eldim, element, ranges, style, factors=factors,
@@ -277,7 +286,7 @@ def _get_colormapper(self, eldim, element, ranges, style, factors=None,
         )
 
         if indices is None:
-            num_elements = len(element.vdims)
+            num_elements = len(vdims)
             step_size = len(cmapper.palette) // num_elements
             indices = np.arange(num_elements) * step_size
 
diff --git a/holoviews/tests/plotting/bokeh/test_rasterplot.py b/holoviews/tests/plotting/bokeh/test_rasterplot.py
index 1809bacab9..6dcaa8b6e0 100644
--- a/holoviews/tests/plotting/bokeh/test_rasterplot.py
+++ b/holoviews/tests/plotting/bokeh/test_rasterplot.py
@@ -9,6 +9,7 @@
 from holoviews.plotting.bokeh.raster import ImageStackPlot
 from holoviews.plotting.bokeh.util import bokeh34
 
+from ..utils import ParamLogStream
 from .test_plot import TestBokehPlot, bokeh_renderer
 
 
@@ -424,6 +425,36 @@ def test_image_stack_dict_cmap(self):
         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"})
+
+    def test_image_stack_dict_cmap_extra(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=["a", "b", "c"])
+        with ParamLogStream() as log:
+            img_stack.opts(cmap={"c": "yellow", "a": "red", "b": "green", "d": "blue"})
+            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 == ["red", "green", "yellow"]
+        log_msg = log.stream.read()
+        assert "extra value dimensions: 'd'" in log_msg
+
 
 class TestImageStackEven(_ImageStackBase):
     __test__ = True