From 9f2bb6fef3ac7873b7f0a1f41ce6464b773087bd Mon Sep 17 00:00:00 2001
From: Thomas A Caswell <tcaswell@gmail.com>
Date: Fri, 28 Oct 2022 01:42:25 -0400
Subject: [PATCH 1/3] ENH: use nu mechanism to do color mapping

---
 data_prototype/wrappers.py | 13 ++++++++++++-
 examples/2Dfunc.py         |  6 +++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/data_prototype/wrappers.py b/data_prototype/wrappers.py
index 4d149ba..51708e2 100644
--- a/data_prototype/wrappers.py
+++ b/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)
diff --git a/examples/2Dfunc.py b/examples/2Dfunc.py
index 9c31a86..9bd6dd4 100644
--- a/examples/2Dfunc.py
+++ b/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)

From 0dc2c6ca1d1be1aa4fe0c8a9f72e27a642bd7891 Mon Sep 17 00:00:00 2001
From: Thomas A Caswell <tcaswell@gmail.com>
Date: Fri, 28 Oct 2022 02:17:17 -0400
Subject: [PATCH 2/3] ENH: add an example of a multi-variate color mapping

---
 examples/mulivariate_cmap.py | 47 ++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)
 create mode 100644 examples/mulivariate_cmap.py

diff --git a/examples/mulivariate_cmap.py b/examples/mulivariate_cmap.py
new file mode 100644
index 0000000..f737bbc
--- /dev/null
+++ b/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()

From 980957b51a3ba7812da50eb39eb19409557c2ad7 Mon Sep 17 00:00:00 2001
From: Thomas A Caswell <tcaswell@gmail.com>
Date: Fri, 28 Oct 2022 09:55:38 -0400
Subject: [PATCH 3/3] STY: apply black

---
 data_prototype/wrappers.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/data_prototype/wrappers.py b/data_prototype/wrappers.py
index 51708e2..2d4875b 100644
--- a/data_prototype/wrappers.py
+++ b/data_prototype/wrappers.py
@@ -197,13 +197,13 @@ def __init__(self, data: DataContainer, nus=None, /, cmap=None, norm=None, **kwa
         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:
+            if nus is not None and "image" in nus:
                 raise ValueError("Conflicting input")
             if cmap is None:
-                cmap = mpl.colormaps['viridis']
+                cmap = mpl.colormaps["viridis"]
             if norm is None:
                 raise ValueError("not sure how to do autoscaling yet")
-            nus['image'] = lambda image: cmap(norm(image))
+            nus["image"] = lambda image: cmap(norm(image))
         super().__init__(data, nus)
         kwargs.setdefault("origin", "lower")
         self._wrapped_instance = self._wrapped_class(None, **kwargs)