Merge pull request #5 from tacaswell/multivariate_cmap

Multivariate cmap

Author: ksunden

data_prototype/wrappers.py

@@ -5,6 +5,7 @@
 from cachetools import LFUCache
 from functools import partial, wraps
 
+import matplotlib as mpl
 from matplotlib.lines import Line2D as _Line2D
 from matplotlib.image import AxesImage as _AxesImage
 from matplotlib.patches import StepPatch as _StepPatch
@@ -192,7 +193,17 @@ def _update_wrapped(self, data):
 class ImageWrapper(ProxyWrapper):
     _wrapped_class = _AxesImage
 
-    def __init__(self, data: DataContainer, nus=None, /, **kwargs):
+    def __init__(self, data: DataContainer, nus=None, /, cmap=None, norm=None, **kwargs):
+        print(kwargs, nus)
+        nus = dict(nus or {})
+        if cmap is not None or norm is not None:
+            if nus is not None and "image" in nus:
+                raise ValueError("Conflicting input")
+            if cmap is None:
+                cmap = mpl.colormaps["viridis"]
+            if norm is None:
+                raise ValueError("not sure how to do autoscaling yet")
+            nus["image"] = lambda image: cmap(norm(image))
         super().__init__(data, nus)
         kwargs.setdefault("origin", "lower")
         self._wrapped_instance = self._wrapped_class(None, **kwargs)

examples/2Dfunc.py

@@ -12,8 +12,10 @@
 from data_prototype.wrappers import ImageWrapper
 from data_prototype.containers import FuncContainer
 
+import matplotlib as mpl
 from matplotlib.colors import Normalize
 
+
 fc = FuncContainer(
     {},
     xyfuncs={
@@ -22,7 +24,9 @@
         "image": (("N", "M"), lambda x, y: np.sin(x).reshape(1, -1) * np.cos(y).reshape(-1, 1)),
     },
 )
-im = ImageWrapper(fc, norm=Normalize(-1, 1))
+cmap = mpl.colormaps["viridis"]
+norm = Normalize(-1, 1)
+im = ImageWrapper(fc, {"image": lambda image: cmap(norm(image))})
 
 fig, ax = plt.subplots()
 ax.add_artist(im)

examples/mulivariate_cmap.py

@@ -0,0 +1,47 @@
+"""
+=====================
+A functional 2D image
+=====================
+
+
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from data_prototype.wrappers import ImageWrapper
+from data_prototype.containers import FuncContainer
+
+from matplotlib.colors import hsv_to_rgb
+
+
+def func(x, y):
+    return (
+        (np.sin(x).reshape(1, -1) * np.cos(y).reshape(-1, 1)) ** 2,
+        np.arctan2(np.cos(y).reshape(-1, 1), np.sin(x).reshape(1, -1)),
+    )
+
+
+def image_nu(image):
+    saturation, angle = image
+    hue = (angle + np.pi) / (2 * np.pi)
+    value = np.ones_like(hue)
+    return np.clip(hsv_to_rgb(np.stack([hue, saturation, value], axis=2)), 0, 1)
+
+
+fc = FuncContainer(
+    {},
+    xyfuncs={
+        "xextent": ((2,), lambda x, y: [x[0], x[-1]]),
+        "yextent": ((2,), lambda x, y: [y[0], y[-1]]),
+        "image": (("N", "M", 2), func),
+    },
+)
+
+im = ImageWrapper(fc, {"image": image_nu})
+
+fig, ax = plt.subplots()
+ax.add_artist(im)
+ax.set_xlim(-5, 5)
+ax.set_ylim(-5, 5)
+plt.show()