Merge pull request #36 from ksunden/late_conversion

Delayed conversion Node for units

Author: tacaswell

data_prototype/axes.py

@@ -0,0 +1,141 @@
+import numpy as np
+
+
+import matplotlib as mpl
+from matplotlib.axes._axes import Axes as MPLAxes, _preprocess_data
+import matplotlib.collections as mcoll
+import matplotlib.cbook as cbook
+import matplotlib.markers as mmarkers
+import matplotlib.projections as mprojections
+
+from .containers import ArrayContainer, DataUnion
+from .conversion_node import (
+    DelayedConversionNode,
+    FunctionConversionNode,
+    RenameConversionNode,
+)
+from .wrappers import PathCollectionWrapper
+
+
+class Axes(MPLAxes):
+    # Name for registering as a projection so we can experiment with it
+    name = "data-prototype"
+
+    @_preprocess_data(
+        replace_names=[
+            "x",
+            "y",
+            "s",
+            "linewidths",
+            "edgecolors",
+            "c",
+            "facecolor",
+            "facecolors",
+            "color",
+        ],
+        label_namer="y",
+    )
+    def scatter(
+        self,
+        x,
+        y,
+        s=None,
+        c=None,
+        marker=None,
+        cmap=None,
+        norm=None,
+        vmin=None,
+        vmax=None,
+        alpha=None,
+        linewidths=None,
+        *,
+        edgecolors=None,
+        plotnonfinite=False,
+        **kwargs
+    ):
+        # TODO implement normalize kwargs as a pipeline stage
+        # add edgecolors and linewidths to kwargs so they can be processed by
+        # normalize_kwargs
+        if edgecolors is not None:
+            kwargs.update({"edgecolors": edgecolors})
+        if linewidths is not None:
+            kwargs.update({"linewidths": linewidths})
+
+        kwargs = cbook.normalize_kwargs(kwargs, mcoll.Collection)
+        c, colors, edgecolors = self._parse_scatter_color_args(
+            c,
+            edgecolors,
+            kwargs,
+            np.ma.ravel(x).size,
+            get_next_color_func=self._get_patches_for_fill.get_next_color,
+        )
+
+        inputs = ArrayContainer(
+            x=x,
+            y=y,
+            s=s,
+            c=c,
+            marker=marker,
+            cmap=cmap,
+            norm=norm,
+            vmin=vmin,
+            vmax=vmax,
+            alpha=alpha,
+            plotnonfinite=plotnonfinite,
+            facecolors=colors,
+            edgecolors=edgecolors,
+            **kwargs
+        )
+        # TODO should more go in here?
+        # marker/s are always in Container, but require overriding if None
+        # Color handling is odd too
+        defaults = ArrayContainer(
+            linewidths=mpl.rcParams["lines.linewidth"],
+        )
+
+        cont = DataUnion(defaults, inputs)
+
+        pipeline = []
+        xconvert = DelayedConversionNode.from_keys(("x",), converter_key="xunits")
+        yconvert = DelayedConversionNode.from_keys(("y",), converter_key="yunits")
+        pipeline.extend([xconvert, yconvert])
+        pipeline.append(lambda x: np.ma.ravel(x))
+        pipeline.append(lambda y: np.ma.ravel(y))
+        pipeline.append(
+            lambda s: np.ma.ravel(s)
+            if s is not None
+            else [20]
+            if mpl.rcParams["_internal.classic_mode"]
+            else [mpl.rcParams["lines.markersize"] ** 2.0]
+        )
+        # TODO plotnonfinite/mask combining
+        pipeline.append(
+            lambda marker: marker
+            if marker is not None
+            else mpl.rcParams["scatter.marker"]
+        )
+        pipeline.append(
+            lambda marker: marker
+            if isinstance(marker, mmarkers.MarkerStyle)
+            else mmarkers.MarkerStyle(marker)
+        )
+        pipeline.append(
+            FunctionConversionNode.from_funcs(
+                {
+                    "paths": lambda marker: [
+                        marker.get_path().transformed(marker.get_transform())
+                    ]
+                }
+            )
+        )
+        pipeline.append(RenameConversionNode.from_mapping({"s": "sizes"}))
+
+        # TODO classic mode margin override?
+        pcw = PathCollectionWrapper(cont, pipeline, offset_transform=self.transData)
+        self.add_artist(pcw)
+        self._request_autoscale_view()
+        return pcw
+
+
+# This is a handy trick to allow e.g. plt.subplots(subplot_kw={'projection': 'data-prototype'})
+mprojections.register_projection(Axes)

data_prototype/containers.py

@@ -151,7 +151,12 @@ class ArrayContainer:
     def __init__(self, **data):
         self._data = data
         self._cache_key = str(uuid.uuid4())
-        self._desc = {k: Desc(v.shape, v.dtype) for k, v in data.items()}
+        self._desc = {
+            k: Desc(v.shape, v.dtype)
+            if isinstance(v, np.ndarray)
+            else Desc((), type(v))
+            for k, v in data.items()
+        }
 
     def query(
         self,

data_prototype/conversion_node.py

@@ -9,9 +9,19 @@
 from typing import Any
 
 
-def evaluate_pipeline(nodes: Sequence[ConversionNode], input: dict[str, Any]):
+def evaluate_pipeline(
+    nodes: Sequence[ConversionNode],
+    input: dict[str, Any],
+    delayed_converters: dict[str, Callable] | None = None,
+):
     for node in nodes:
-        input = node.evaluate(input)
+        if isinstance(node, Callable):
+            k = list(inspect.signature(node).parameters.keys())[0]
+            node = FunctionConversionNode.from_funcs({k: node})
+        if isinstance(node, DelayedConversionNode):
+            input = node.evaluate(input, delayed_converters)
+        else:
+            input = node.evaluate(input)
     return input
 
 
@@ -113,3 +123,27 @@ def from_keys(cls, keys: Sequence[str]):
 
     def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
         return {k: v for k, v in input.items() if k in self.keys}
+
+
+@dataclass
+class DelayedConversionNode(ConversionNode):
+    converter_key: str
+
+    @classmethod
+    def from_keys(cls, keys: Sequence[str], converter_key: str):
+        return cls(
+            tuple(keys), tuple(keys), trim_keys=False, converter_key=converter_key
+        )
+
+    def evaluate(
+        self, input: dict[str, Any], converters: dict[str, Callable] | None = None
+    ) -> dict[str, Any]:
+        return super().evaluate(
+            {
+                **input,
+                **{
+                    k: converters[self.converter_key](input[k])
+                    for k in self.required_keys
+                },
+            }
+        )

data_prototype/patches.py

@@ -30,8 +30,6 @@ class PatchWrapper(ProxyWrapper):
         "set_joinstyle",
         "set_path",
     )
-    _xunits = ()
-    _yunits = ()
     required_keys = {
         "edgecolor",
         "facecolor",
@@ -50,11 +48,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
 
     @_stale_wrapper
     def draw(self, renderer):
-        self._update_wrapped(
-            self._query_and_transform(
-                renderer, xunits=self._xunits, yunits=self._yunits
-            )
-        )
+        self._update_wrapped(self._query_and_transform(renderer))
         return self._wrapped_instance.draw(renderer)
 
     def _update_wrapped(self, data):
@@ -77,8 +71,6 @@ class RectangleWrapper(PatchWrapper):
         "set_angle",
         "set_rotation_point",
     )
-    _xunits = ("x", "width")
-    _yunits = ("y", "height")
     required_keys = PatchWrapper.required_keys | {
         "x",
         "y",

data_prototype/wrappers.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any, Protocol, Tuple, get_type_hints
+from typing import Dict, Any, Protocol, Tuple, get_type_hints
 import inspect
 
 import numpy as np
@@ -121,17 +121,13 @@ def draw(self, renderer):
     def _update_wrapped(self, data):
         raise NotImplementedError
 
-    def _query_and_transform(
-        self, renderer, *, xunits: List[str], yunits: List[str]
-    ) -> Dict[str, Any]:
+    def _query_and_transform(self, renderer) -> Dict[str, Any]:
         """
         Helper to centralize the data querying and python-side transforms
 
         Parameters
         ----------
         renderer : RendererBase
-        xunits, yunits : List[str]
-            The list of keys that need to be run through the x and y unit machinery.
         """
         # extract what we need to about the axes to query the data
         ax = self.axes
@@ -153,8 +149,11 @@ def _query_and_transform(
             return self._cache[cache_key]
         except KeyError:
             ...
-        # TODO units
-        transformed_data = evaluate_pipeline(self._converters, data)
+        delayed_conversion = {
+            "xunits": ax.xaxis.convert_units,
+            "yunits": ax.yaxis.convert_units,
+        }
+        transformed_data = evaluate_pipeline(self._converters, data, delayed_conversion)
 
         self._cache[cache_key] = transformed_data
         return transformed_data
@@ -232,7 +231,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["x"], yunits=["y"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -265,7 +264,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["x"], yunits=["y"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -304,7 +303,7 @@ def __init__(
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["xextent"], yunits=["yextent"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -325,7 +324,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["edges"], yunits=["density"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -344,7 +343,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=[], yunits=[]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -425,11 +424,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(
-                renderer,
-                xunits=["x", "xupper", "xlower"],
-                yunits=["y", "yupper", "ylower"],
-            ),
+            self._query_and_transform(renderer),
         )
         for k, v in self._wrapped_instances.items():
             v.draw(renderer)

examples/scatter_with_custom_axes.py

@@ -0,0 +1,33 @@
+"""
+=========================================
+An simple scatter plot using `ax.scatter`
+=========================================
+
+This is a quick comparison between the current Matplotlib `scatter` and
+the version in :file:`data_prototype/axes.py`, which uses data containers
+and a conversion pipeline.
+
+This is here to show what does work and what does not work with the current
+implementation of container-based artist drawing.
+"""
+
+
+import data_prototype.axes  # side-effect registers projection # noqa
+
+import matplotlib.pyplot as plt
+
+fig = plt.figure()
+newstyle = fig.add_subplot(2, 1, 1, projection="data-prototype")
+oldstyle = fig.add_subplot(2, 1, 2)
+
+newstyle.scatter([0, 1, 2], [2, 5, 1])
+oldstyle.scatter([0, 1, 2], [2, 5, 1])
+newstyle.scatter([0, 1, 2], [3, 1, 2])
+oldstyle.scatter([0, 1, 2], [3, 1, 2])
+
+
+# Autoscaling not working
+newstyle.set_xlim(oldstyle.get_xlim())
+newstyle.set_ylim(oldstyle.get_ylim())
+
+plt.show()