Merge pull request #9 from ksunden/scatter

Early PathCollection support

Author: tacaswell

data_prototype/wrappers.py

@@ -11,7 +11,8 @@
 from matplotlib.image import AxesImage as _AxesImage
 from matplotlib.patches import StepPatch as _StepPatch
 from matplotlib.text import Text as _Text
-from matplotlib.collections import LineCollection as _LineCollection
+import matplotlib.transforms as mtransforms
+from matplotlib.collections import LineCollection as _LineCollection, PathCollection as _PathCollection
 from matplotlib.artist import Artist as _Artist
 
 from data_prototype.containers import DataContainer, _MatplotlibTransform
@@ -206,7 +207,7 @@ class LineWrapper(ProxyWrapper):
 
     def __init__(self, data: DataContainer, nus=None, /, **kwargs):
         super().__init__(data, nus)
-        self._wrapped_instance = self._wrapped_class([], [], **kwargs)
+        self._wrapped_instance = self._wrapped_class(np.array([]), np.array([]), **kwargs)
 
     @_stale_wrapper
     def draw(self, renderer):
@@ -221,6 +222,42 @@ def _update_wrapped(self, data):
             getattr(self._wrapped_instance, f"set_{k}")(v)
 
 
+class PathCollectionWrapper(ProxyWrapper):
+    _wrapped_class = _PathCollection
+    required_keys = {"x", "y", "paths", "facecolors", "edgecolors", "sizes"}
+    _privtized_methods = (
+        "set_facecolors",
+        "set_edgecolors",
+        "set_offsets",
+        "set_sizes",
+        "set_paths",
+        "get_facecolors",
+        "get_edgecolors",
+        "get_offsets",
+        "get_sizes",
+        "get_paths",
+    )
+
+    def __init__(self, data: DataContainer, nus=None, /, **kwargs):
+        super().__init__(data, nus)
+        self._wrapped_instance = self._wrapped_class([], **kwargs)
+        self._wrapped_instance.set_transform(mtransforms.IdentityTransform())
+
+    @_stale_wrapper
+    def draw(self, renderer):
+        self._update_wrapped(
+            self._query_and_transform(renderer, xunits=["x"], yunits=["y"]),
+        )
+        return self._wrapped_instance.draw(renderer)
+
+    def _update_wrapped(self, data):
+        self._wrapped_instance.set_offsets(np.array([data["x"], data["y"]]).T)
+        self._wrapped_instance.set_paths(data["paths"])
+        self._wrapped_instance.set_facecolors(data["facecolors"])
+        self._wrapped_instance.set_edgecolors(data["edgecolors"])
+        self._wrapped_instance.set_sizes(data["sizes"])
+
+
 class ImageWrapper(ProxyWrapper):
     _wrapped_class = _AxesImage
     required_keys = {"xextent", "yextent", "image"}

examples/lissajous.py

@@ -0,0 +1,97 @@
+"""
+==========================
+An animated lissajous ball
+==========================
+
+Inspired by https://twitter.com/_brohrer_/status/1584681864648065027
+
+
+"""
+import time
+from typing import Dict, Tuple, Any, Union
+from functools import partial
+
+import numpy as np
+
+import matplotlib.pyplot as plt
+import matplotlib.markers as mmarkers
+from matplotlib.animation import FuncAnimation
+
+from data_prototype.containers import _MatplotlibTransform, Desc
+
+from data_prototype.wrappers import PathCollectionWrapper, FormatedText
+
+
+class Lissajous:
+    N = 1024
+    # cycles per minutes
+    scale = 2
+
+    def describe(self):
+        return {
+            "x": Desc([self.N], float),
+            "y": Desc([self.N], float),
+            "phase": Desc([], float),
+            "time": Desc([], float),
+            "sizes": Desc([], float),
+            "paths": Desc([], float),
+            "edgecolors": Desc([], str),
+            "facecolors": Desc([self.N], str),
+        }
+
+    def query(
+        self,
+        transform: _MatplotlibTransform,
+        size: Tuple[int, int],
+    ) -> Tuple[Dict[str, Any], Union[str, int]]:
+        def next_time():
+            cur_time = time.time()
+            cur_time = np.array([cur_time, cur_time - 0.1, cur_time - 0.2, cur_time - 0.3])
+
+            phase = 15 * np.pi * (self.scale * cur_time % 60) / 150
+            marker_obj = mmarkers.MarkerStyle("o")
+            return {
+                "x": np.cos(5 * phase),
+                "y": np.sin(3 * phase),
+                "phase": phase[0],
+                "sizes": np.array([256]),
+                "paths": [marker_obj.get_path().transformed(marker_obj.get_transform())],
+                "edgecolors": "k",
+                "facecolors": ["#4682b4ff", "#82b446aa", "#46b48288", "#8246b433"],
+                "time": cur_time[0],
+            }, hash(cur_time[0])
+
+        return next_time()
+
+
+def update(frame, art):
+    return art
+
+
+sot_c = Lissajous()
+
+fc = FormatedText(
+    sot_c,
+    {"text": lambda phase: f"ϕ={phase:.2f}  "},
+    x=1,
+    y=1,
+    ha="right",
+)
+fig, ax = plt.subplots()
+ax.set_xlim(-1.1, 1.1)
+ax.set_ylim(-1.1, 1.1)
+lw = PathCollectionWrapper(sot_c, offset_transform=ax.transData)
+ax.add_artist(lw)
+ax.add_artist(fc)
+# ax.set_xticks([])
+# ax.set_yticks([])
+ax.set_aspect(1)
+ani = FuncAnimation(
+    fig,
+    partial(update, art=(lw, fc)),
+    frames=60 * 15,
+    interval=1000 / 100,
+    # TODO: blitting does not work because wrappers do not inherent from Artist
+    # blit=True,
+)
+plt.show()

examples/simple_scatter.py

@@ -0,0 +1,43 @@
+"""
+==========================
+An animated lissajous ball
+==========================
+
+Inspired by https://twitter.com/_brohrer_/status/1584681864648065027
+
+
+"""
+import numpy as np
+
+import matplotlib.pyplot as plt
+import matplotlib.markers as mmarkers
+
+from data_prototype.containers import ArrayContainer
+
+from data_prototype.wrappers import PathCollectionWrapper
+
+
+def update(frame, art):
+    return art
+
+
+marker_obj = mmarkers.MarkerStyle("o")
+
+cont = ArrayContainer(
+    x=np.array([0, 1, 2]),
+    y=np.array([1, 4, 2]),
+    paths=np.array([marker_obj.get_path()]),
+    sizes=np.array([12]),
+    edgecolors=np.array(["k"]),
+    facecolors=np.array(["C3"]),
+)
+
+fig, ax = plt.subplots()
+ax.set_xlim(-0.5, 2.5)
+ax.set_ylim(0, 5)
+lw = PathCollectionWrapper(cont, offset_transform=ax.transData)
+ax.add_artist(lw)
+# ax.set_xticks([])
+# ax.set_yticks([])
+# ax.set_aspect(1)
+plt.show()