ENH: add caching logic

Author: tacaswell

Diff for: data_prototype/containers.py

@@ -21,12 +21,16 @@ def __init__(self, xfuncs, yfuncs=None, xyfuncs=None):
         self._xyfuncs = xyfuncs or {}
 
     def query(self, data_bounds, size, xscale=None, yscale=None):
+        hash_key = hash((data_bounds, size, xscale, yscale))
+        if hash_key in self._cache:
+            return self._cache[hash_key], hash_key
         xmin, xmax, ymin, ymax = data_bounds
         xpix, ypix = size
         x_data = np.linspace(xmin, xmax, int(xpix) * 2)
         y_data = np.linspace(ymin, ymax, int(ypix) * 2)
-        return dict(
+        ret = self._cache[hash_key] = dict(
             **{k: f(x_data) for k, f in self._xfuncs.items()},
             **{k: f(y_data) for k, f in self._yfuncs.items()},
             **{k: f(x_data, y_data) for k, f in self._xyfuncs.items()}
         )
+        return ret, hash_key

Diff for: data_prototype/wrappers.py

@@ -1,3 +1,9 @@
+from typing import List, Dict, Any
+
+import numpy as np
+
+from cachetools import LFUCache
+
 from matplotlib.lines import Line2D as _Line2D
 from matplotlib.image import AxesImage as _AxesImage
 
@@ -24,18 +30,54 @@ def __setattr__(self, key, value):
         else:
             super().__setattr__(key, value)
 
-    def _query_and_transform(self, renderer):
+    def _query_and_transform(self, renderer, *, xunits: List[str], yunits: List[str]) -> 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._wrapped_instance.axes
+        # TODO do we want to trust the implicit renderer on the Axes?
         ax_bbox = ax.get_window_extent(renderer)
-        return {
-            # doing this here is nice because we can write it once, but we really want to
-            # push this computation down a layer
-            k: self.nus.get(k, lambda x: x)(v)
-            for k, v in self.data.query([*ax.get_xlim(), *ax.get_ylim()], ax_bbox.size).items()
-        }
+
+        # actually query the underlying data.  This returns both the (raw) data
+        # and key to use for caching.
+        data, cache_key = self.data.query(
+            # TODO do this need to be (de) unitized
+            (*ax.get_xlim(), *ax.get_ylim()),
+            tuple(np.round(ax_bbox.size).astype(int)),
+            # TODO sort out how to spell the x/y scale
+            # TODO is scale enoguh?  What do we have to do about non-trivial projection?
+            xscale=None,
+            yscale=None,
+        )
+        # see if we can short-circuit
+        try:
+            return self._cache[cache_key]
+        except KeyError:
+            ...
+        # TODO decide if units go pre-nu or post-nu?
+        for x_like in xunits:
+            data[x_like] = ax.xaxis.convert_units(data[x_like])
+        for y_like in yunits:
+            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 interaporates with the transform stack
+        data = {k: self.nus.get(k, lambda x: x)(v) for k, v in data.items()}
+        self._cache[cache_key] = data
+        return data
 
     def __init__(self, data, nus):
         self.data = data
+        self._cache = LFUCache(64)
+        # TODO make sure mutating this will invalidate the cache!
         self.nus = nus or {}
 
 
@@ -48,7 +90,7 @@ def __init__(self, data, nus=None, /, **kwargs):
         self._wrapped_instance = self._wrapped_class([], [], **kwargs)
 
     def draw(self, renderer):
-        data = self._query_and_transform(renderer)
+        data = self._query_and_transform(renderer, xunits=["x"], yunits=["y"])
         self._wrapped_instance.set_data(data["x"], data["y"])
         return self._wrapped_instance.draw(renderer)
 
@@ -62,7 +104,7 @@ def __init__(self, data, nus=None, /, **kwargs):
         self._wrapped_instance = self._wrapped_class(None, **kwargs)
 
     def draw(self, renderer):
-        data = self._query_and_transform(renderer)
+        data = self._query_and_transform(renderer, xunits=["xextent"], yunits=["yextent"])
         self._wrapped_instance.set_array(data["image"])
-        self._wrapped_instance.set_extent(data["extent"])
+        self._wrapped_instance.set_extent([*data["xextent"], *data["yextent"]])
         return self._wrapped_instance.draw(renderer)

Diff for: requirements-dev.txt

@@ -16,3 +16,5 @@ numpydoc
 sphinx-copybutton
 mpl-sphinx-theme
 sphinx-gallery
+# typing
+types-cachetools

Diff for: requirements.txt

@@ -1,4 +1,5 @@
 # List required packages in this file, one per line.
+cachetools
 matplotlib
 pandas
 xarray