Get examples working with conversion node infrastructure

Author: ksunden

data_prototype/conversion_node.py

@@ -95,3 +95,15 @@ def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
                 **{k: func(**{p: input[p] for p in sig.parameters}) for (k, (func, sig)) in self._sigs.items()},
             }
         )
+
+
+@dataclass
+class LimitKeysConversionNode(ConversionNode):
+    keys: set[str]
+
+    @classmethod
+    def from_keys(cls, name: str, keys: Sequence[str]):
+        return cls(name, (), tuple(keys), trim_keys=True, keys=set(keys))
+
+    def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
+        return {k: v for k, v in input.items() if k in self.keys}

data_prototype/wrappers.py

@@ -16,6 +16,13 @@
 from matplotlib.artist import Artist as _Artist
 
 from data_prototype.containers import DataContainer, _MatplotlibTransform
+from data_prototype.conversion_node import (
+    ConversionNode,
+    RenameConversionNode,
+    evaluate_pipeline,
+    FunctionConversionNode,
+    LimitKeysConversionNode,
+)
 
 
 class _BBox(Protocol):
@@ -139,45 +146,26 @@ def _query_and_transform(self, renderer, *, xunits: List[str], yunits: List[str]
             return self._cache[cache_key]
         except KeyError:
             ...
-        # TODO decide if units go pre-nu or post-nu?
-        for x_like in xunits:
-            if x_like in data:
-                data[x_like] = ax.xaxis.convert_units(data[x_like])
-        for y_like in yunits:
-            if y_like in data:
-                data[y_like] = ax.xaxis.convert_units(data[y_like])
-
-        # doing the nu work here is nice because we can write it once, but we
-        # really want to push this computation down a layer
-        # TODO sort out how this interoperates with the transform stack
-        transformed_data = {}
-        for k, (nu, sig) in self._sigs.items():
-            to_pass = set(sig.parameters)
-            transformed_data[k] = nu(**{k: data[k] for k in to_pass})
+        # TODO units
+        transformed_data = evaluate_pipeline(self._converters, data)
 
         self._cache[cache_key] = transformed_data
         return transformed_data
 
-    def __init__(self, data, converters, **kwargs):
+    def __init__(self, data, converters: ConversionNode | list[ConversionNode] | None, **kwargs):
         super().__init__(**kwargs)
         self.data = data
         self._cache = LFUCache(64)
         # TODO make sure mutating this will invalidate the cache!
-        self._converters = converters or {}
-        for k in self.required_keys:
-            self._converters.setdefault(k, _make_identity(k))
-        desc = data.describe()
-        for k in self.expected_keys:
-            if k in desc:
-                self._converters.setdefault(k, _make_identity(k))
-        self._sigs = {k: (nu, inspect.signature(nu)) for k, nu in self._converters.items()}
+        if isinstance(converters, ConversionNode):
+            converters = [converters]
+        self._converters: list[ConversionNode] = converters or []
+        setters = list(self.expected_keys | self.required_keys)
+        if hasattr(self, "_wrapped_class"):
+            setters += [f[4:] for f in dir(self._wrapped_class) if f.startswith("set_")]
+        self._converters.append(LimitKeysConversionNode.from_keys("setters", setters))
         self.stale = True
 
-    # TODO add a setter
-    @property
-    def converters(self):
-        return dict(self._converters)
-
 
 class ProxyWrapper(ProxyWrapperBase):
     _privtized_methods: Tuple[str, ...] = ()
@@ -208,6 +196,9 @@ class LineWrapper(ProxyWrapper):
     def __init__(self, data: DataContainer, converters=None, /, **kwargs):
         super().__init__(data, converters)
         self._wrapped_instance = self._wrapped_class(np.array([]), np.array([]), **kwargs)
+        self._converters.insert(-1, RenameConversionNode.from_mapping("xydata", {"x": "xdata", "y": "ydata"}))
+        setters = [f[4:] for f in dir(self._wrapped_class) if f.startswith("set_")]
+        self._converters[-1] = LimitKeysConversionNode.from_keys("setters", setters)
 
     @_stale_wrapper
     def draw(self, renderer):
@@ -218,7 +209,6 @@ def draw(self, renderer):
 
     def _update_wrapped(self, data):
         for k, v in data.items():
-            k = {"x": "xdata", "y": "ydata"}.get(k, k)
             getattr(self._wrapped_instance, f"set_{k}")(v)
 
 
@@ -263,15 +253,17 @@ class ImageWrapper(ProxyWrapper):
     required_keys = {"xextent", "yextent", "image"}
 
     def __init__(self, data: DataContainer, converters=None, /, cmap=None, norm=None, **kwargs):
-        converters = dict(converters or {})
+        converters = converters or []
         if cmap is not None or norm is not None:
             if converters is not None and "image" in converters:
                 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")
-            converters["image"] = lambda image: cmap(norm(image))
+            converters.append(
+                FunctionConversionNode.from_funcs("map colors", {"image": lambda image: cmap(norm(image))})
+            )
         super().__init__(data, converters)
         kwargs.setdefault("origin", "lower")
         self._wrapped_instance = self._wrapped_class(None, **kwargs)

examples/2Dfunc.py

@@ -27,7 +27,7 @@
 )
 cmap = mpl.colormaps["viridis"]
 norm = Normalize(-1, 1)
-im = ImageWrapper(fc, {"image": lambda image: cmap(norm(image))})
+im = ImageWrapper(fc)
 
 fig, ax = plt.subplots()
 ax.add_artist(im)

examples/animation.py

@@ -17,6 +17,7 @@
 from matplotlib.animation import FuncAnimation
 
 from data_prototype.containers import _MatplotlibTransform, Desc
+from data_prototype.conversion_node import FunctionConversionNode
 
 from data_prototype.wrappers import LineWrapper, FormatedText
 
@@ -63,9 +64,9 @@ def update(frame, art):
 lw = LineWrapper(sot_c, lw=5, color="green", label="sin(time)")
 fc = FormatedText(
     sot_c,
-    {"text": lambda phase: f"ϕ={phase:.2f}"},
-    x=2 * np.pi,
-    y=1,
+    FunctionConversionNode.from_funcs(
+        "fmt", {"text": lambda phase: f"ϕ={phase:.2f}", "x": lambda: 2 * np.pi, "y": lambda: 1}
+    ),
     ha="right",
 )
 fig, ax = plt.subplots()

examples/mapped.py

@@ -14,25 +14,33 @@
 
 from data_prototype.wrappers import LineWrapper, FormatedText
 from data_prototype.containers import ArrayContainer
+from data_prototype.conversion_node import FunctionConversionNode
 
 cmap = plt.colormaps["viridis"]
 cmap.set_over("k")
 cmap.set_under("r")
 norm = Normalize(1, 8)
 
-line_converter = {
-    # arbitrary functions
-    "lw": lambda lw: min(1 + lw, 5),
-    # standard color mapping
-    "color": lambda j: cmap(norm(j)),
-    # categorical
-    "ls": lambda cat: {"A": "-", "B": ":", "C": "--"}[cat[()]],
-}
+line_converter = FunctionConversionNode.from_funcs(
+    "line converter",
+    {
+        # arbitrary functions
+        "lw": lambda lw: min(1 + lw, 5),
+        # standard color mapping
+        "color": lambda j: cmap(norm(j)),
+        # categorical
+        "ls": lambda cat: {"A": "-", "B": ":", "C": "--"}[cat[()]],
+    },
+)
 
-text_converter = {
-    "text": lambda j, cat: f"index={j[()]} class={cat[()]!r}",
-    "y": lambda j: j,
-}
+text_converter = FunctionConversionNode.from_funcs(
+    "text converter",
+    {
+        "text": lambda j, cat: f"index={j[()]} class={cat[()]!r}",
+        "y": lambda j: j,
+        "x": lambda x: 2 * np.pi,
+    },
+)
 
 
 th = np.linspace(0, 2 * np.pi, 128)

examples/mulivariate_cmap.py

@@ -13,6 +13,7 @@
 
 from data_prototype.wrappers import ImageWrapper
 from data_prototype.containers import FuncContainer
+from data_prototype.conversion_node import FunctionConversionNode
 
 from matplotlib.colors import hsv_to_rgb
 
@@ -40,7 +41,7 @@ def image_nu(image):
     },
 )
 
-im = ImageWrapper(fc, {"image": image_nu})
+im = ImageWrapper(fc, FunctionConversionNode.from_funcs("bivariate", {"image": image_nu}))
 
 fig, ax = plt.subplots()
 ax.add_artist(im)

examples/subsample.py

@@ -68,7 +68,7 @@ def query(
 sub = Subsample()
 cmap = mpl.colormaps["coolwarm"]
 norm = Normalize(-2.2, 2.2)
-im = ImageWrapper(sub, {"image": lambda image: cmap(norm(image))})
+im = ImageWrapper(sub)
 
 fig, ax = plt.subplots()
 ax.add_artist(im)

examples/widgets.py

@@ -15,6 +15,7 @@
 
 from data_prototype.wrappers import LineWrapper
 from data_prototype.containers import FuncContainer
+from data_prototype.conversion_node import FunctionConversionNode
 
 
 class SliderContainer(FuncContainer):
@@ -117,7 +118,7 @@ def _query_hash(self, coord_transform, size):
 lw = LineWrapper(
     fc,
     # color map phase (scaled to 2pi and wrapped to [0, 1])
-    {"color": lambda color: cmap((color / (2 * np.pi)) % 1)},
+    FunctionConversionNode.from_funcs("cmap", {"color": lambda color: cmap((color / (2 * np.pi)) % 1)}),
     lw=5,
 )
 ax.add_artist(lw)