ENH: add an example of a multi-variate color mapping

Author: tacaswell

examples/mulivariate_cmap.py

@@ -0,0 +1,48 @@
+"""
+=====================
+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)
+    print(hue.min(), hue.max())
+    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()