fix colorbar for scatter plots

Author: nvaytet

src/matplotgl/image.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import matplotlib as mpl
+import matplotlib.colors as cm
 import numpy as np
 import pythreejs as p3
 
@@ -13,6 +14,7 @@ def __init__(
         array: np.ndarray,
         extent: list[float] | None = None,
         cmap: str = "viridis",
+        norm: str = "linear",
         zorder: float = 0,
     ):
         self.axes = None
@@ -22,7 +24,9 @@ def __init__(
             extent if extent is not None else [0, array.shape[1], 0, array.shape[0]]
         )
         self._zorder = zorder
-        self._norm = Normalizer(vmin=np.min(self._array), vmax=np.max(self._array))
+        self._norm = Normalizer(
+            vmin=np.min(self._array), vmax=np.max(self._array), norm=norm
+        )
         self._cmap = mpl.colormaps[cmap].copy()
         self._texture = p3.DataTexture(
             data=self._make_colors(), format="RGBFormat", type="FloatType"
@@ -97,12 +101,13 @@ def set_extent(self, extent: list[float]) -> None:
 
     def set_cmap(self, cmap: str) -> None:
         self._cmap = mpl.colormaps[cmap].copy()
-        self._texture.data = self._make_colors()
+        # self._texture.data = self._make_colors()
+        self._update_colors()
         if self._colorbar is not None:
             self._colorbar.update()
 
     @property
-    def cmap(self) -> mpl.colormaps:
+    def cmap(self) -> cm.Colormap:
         return self._cmap
 
     @cmap.setter
@@ -121,6 +126,7 @@ def norm(self, norm: Normalizer | str) -> None:
             )
         else:
             self._norm = norm
-        self._texture.data = self._make_colors()
+        # self._texture.data = self._make_colors()
+        self._update_colors()
         if self._colorbar is not None:
             self._colorbar.update()

src/matplotgl/points.py

@@ -3,10 +3,11 @@
 import warnings
 
 import matplotlib as mpl
-import matplotlib.colors as mplc
+import matplotlib.colors as cm
 import numpy as np
 import pythreejs as p3
 
+from .norm import Normalizer
 from .utils import find_limits, fix_empty_range
 
 SHADER_LIBRARY = {
@@ -51,7 +52,17 @@
 
 
 class Points:
-    def __init__(self, x, y, c="C0", s=3, marker="s", zorder=0, cmap="viridis") -> None:
+    def __init__(
+        self,
+        x,
+        y,
+        c="C0",
+        s=3,
+        marker="s",
+        zorder=0,
+        cmap="viridis",
+        norm: str = "linear",
+    ) -> None:
         self.axes = None
         self._x = np.asarray(x)
         self._y = np.asarray(y)
@@ -61,18 +72,21 @@ def __init__(self, x, y, c="C0", s=3, marker="s", zorder=0, cmap="viridis") -> N
 
         if not isinstance(c, str) or not np.isscalar(s) or marker != "s":
             if isinstance(c, str):
-                rgba = mplc.LinearSegmentedColormap.from_list("tmp", [c, c])(
-                    np.ones_like(self._x)
-                )
+                self._c = np.ones_like(self._x)
+                self._norm = Normalizer(vmin=1, vmax=1)
+                self._cmap = cm.LinearSegmentedColormap.from_list("tmp", [c, c])
+                # (
+                #     np.ones_like(self._x)
+                # )
             else:
                 self._c = np.asarray(c)
-                self.norm = mpl.colors.Normalize(
-                    vmin=np.min(self._c), vmax=np.max(self._c)
+                self._norm = Normalizer(
+                    vmin=np.min(self._c), vmax=np.max(self._c), norm=norm
                 )
-                self.cmap = mpl.colormaps[cmap].copy()
-                rgba = self.cmap(self.norm(self._c))
+                self._cmap = mpl.colormaps[cmap].copy()
+                # rgba = self.cmap(self.norm(self._c))
 
-            colors = rgba[:, :3].astype(np.float32)  # Take only RGB, drop alpha
+            colors = self._make_colors()
 
             if np.isscalar(s):
                 sizes = np.full_like(self._x, s, dtype=np.float32)
@@ -123,10 +137,16 @@ def __init__(self, x, y, c="C0", s=3, marker="s", zorder=0, cmap="viridis") -> N
                 }
             )
 
-            self._material = p3.PointsMaterial(color=mplc.to_hex(c), size=s)
+            self._material = p3.PointsMaterial(color=cm.to_hex(c), size=s)
 
         self._points = p3.Points(geometry=self._geometry, material=self._material)
 
+    def _make_colors(self) -> np.ndarray:
+        return self._cmap(self.norm(self._c))[..., :3].astype("float32")
+
+    def _update_colors(self) -> None:
+        self._geometry.attributes["customColor"].array = self._make_colors()
+
     def get_bbox(self):
         pad = 0.03
         left, right = fix_empty_range(find_limits(self._x, scale=self._xscale, pad=pad))
@@ -170,3 +190,33 @@ def _set_xscale(self, scale):
     def _set_yscale(self, scale):
         self._yscale = scale
         self._update()
+
+    def set_cmap(self, cmap: str) -> None:
+        self._cmap = mpl.colormaps[cmap].copy()
+        self._update_colors()
+        if self._colorbar is not None:
+            self._colorbar.update()
+
+    @property
+    def cmap(self) -> cm.Colormap:
+        return self._cmap
+
+    @cmap.setter
+    def cmap(self, cmap: str) -> None:
+        self.set_cmap(cmap)
+
+    @property
+    def norm(self) -> Normalizer:
+        return self._norm
+
+    @norm.setter
+    def norm(self, norm: Normalizer | str) -> None:
+        if isinstance(norm, str):
+            self._norm = Normalizer(
+                vmin=np.min(self._c), vmax=np.max(self._c), norm=norm
+            )
+        else:
+            self._norm = norm
+        self._update_colors()
+        if self._colorbar is not None:
+            self._colorbar.update()