polish quiver, pcolormesh and release

Author: t-makaro

animatplot/_version.py

@@ -1,3 +1,3 @@
 version = (0, 2, 0)
-info = '.dev2'
+info = ''
 __version__ = '.'.join(map(str, version))+info

animatplot/blocks/image_like.py

@@ -8,9 +8,14 @@ class Pcolormesh(Block):
 
     Parameters
     ----------
-    X, Y : 2D np.ndarray, optional
-    C : 3D np.ndarray
-    axis : matplotlib axis
+    X : 2D np.ndarray, optional
+    Y : 2D np.ndarray, optional
+    C : list of 2D np.ndarray or a 3D np.ndarray
+    axis : matplotlib axis, optional
+        an axis to attach the block to.
+    t_axis : int, optional
+        The axis of the array that represents time. Defaults to 0.
+        No effect if C is a list.
 
     Notes
     -----
@@ -29,24 +34,34 @@ def __init__(self, *args, axis=None, t_axis=0, **kwargs):
         else:
             raise TypeError(
                 'Illegal arguments to pcolormesh; see help(pcolormesh)')
-        if len(self.C.shape) != 3:
-            raise TypeError('C must be a 3D array')
 
-        super().__init__(axis)
+        super().__init__(axis, t_axis)
+
+        self._is_list = isinstance(self.C, list)
+        self.C = np.asanyarray(self.C)
 
+        Slice = self._make_slice(0, 3)
         if self._arg_len == 1:
-            self.quad = self.ax.pcolormesh(self.C[:, :, 0], **kwargs)
+            self.quad = self.ax.pcolormesh(self.C[Slice], **kwargs)
         elif self._arg_len == 3:
-            self.quad = self.ax.pcolormesh(self.X, self.Y, self.C[:, :, 0],
+            self.quad = self.ax.pcolormesh(self.X, self.Y, self.C[Slice],
                                            **kwargs)
 
     def _update(self, i):
-        self.quad.set_array(self.C[:-1, :-1, i].ravel())
+        Slice = self._make_pcolormesh_slice(i, 3)
+        self.quad.set_array(self.C[Slice].ravel())
         return self.quad
 
     def __len__(self):
         return self.C.shape[2]
 
+    def _make_pcolormesh_slice(self, i, dim):
+        if self._is_list:
+            return i
+        Slice = [slice(-1)]*3  # weird thing to make animation work
+        Slice[self.t_axis] = i
+        return tuple(Slice)
+
 
 class Imshow(Block):
     """Animates a series of images
@@ -64,6 +79,7 @@ class Imshow(Block):
     t_axis : int, optional
         The axis of the array that represents time. Defaults to 0.
         No effect if images is a list.
+
     Notes
     -----
     This block accepts additional keyword arguments to be passed to

animatplot/blocks/vectors.py

@@ -1,4 +1,5 @@
 from .base import Block
+import numpy as np
 
 
 class Quiver(Block):
@@ -7,48 +8,53 @@ class Quiver(Block):
 
     Parameters
     ----------
-    X, Y : 2D or 3D numpy array
-    U, V : 3D numpy array
+    X : 1D or 2D numpy array
+        The x positions of the arrows. Cannot be animated.
+    Y : 1D or 2D numpy array
+        The y positions of the arrows. Cannot be animated.
+    U : 2D or 3D numpy array
+        The U displacement of the arrows. 1 dimension
+        higher than the X, Y arrays.
+    V : 2D or 3D numpy array
+        The V displcement of the arrows. 1 dimension
+        higher than the X, Y arrays.
     axis : matplotlib axis, optional
         The axis to the block to
+    t_axis : int, optional
+        The axis of the array that represents time. Defaults to 0.
+        No effect if U, V are lists.
+
     Notes
     -----
     This block accepts additional keyword arguments to be passed to
     :meth:`matplotlib.axes.Axes.quiver`
     """
-    def __init__(self, X, Y, U, V, axis=None, **kwargs):
+    def __init__(self, X, Y, U, V, axis=None, t_axis=0, **kwargs):
+        self.X = X
+        self.Y = Y
+        self.U = np.asanyarray(U)
+        self.V = np.asanyarray(V)
         if X.shape != Y.shape:
             raise ValueError("X, Y must have the same shape")
-        if U.shape != V.shape:
+        if self.U.shape != self.V.shape:
             raise ValueError("U, V must have the same shape")
 
-        # lots of features removed b/c quiver plots can't set_XY
-        # self.animate_XY = len(X.shape) == 3
-        # self.animate_UV = len(U.shape) == 3
-
-        self.x = X
-        self.y = Y
-        self.u = U
-        self.v = V
-        self.ax = axis
+        super().__init__(axis, t_axis)
 
-        xy_slice = [slice(None)]*2 + ([slice(1)] if len(X.shape) == 3 else [])
-        # uv_slice = [slice(None)]*2 +([slice(1)] if len(U.shape) == 3 else [])
+        self._dim = len(self.U.shape)
+        self._is_list = isinstance(U, list)
 
-        self.Q = self.ax.quiver(X[tuple(xy_slice)].squeeze(),
-                                Y[tuple(xy_slice)].squeeze(),
-                                U[:, :, 0],            V[:, :, 0],
+        Slice = self._make_slice(0, self._dim)
+        self.Q = self.ax.quiver(self.X, self.Y,
+                                self.U[Slice], self.V[Slice],
                                 **kwargs)
 
     def _update(self, i):
-        # if self.animate_UV:
-        self.Q.set_UVC(self.u[:, :, i], self.v[:, :, i])
-        # if self.animate_XY:
-        #    pass  # self.Q.set_XYC(self.x[:, :, i], self.y[:, :, i])
-        return self.Q,
+        Slice = self._make_slice(i, self._dim)
+        self.Q.set_UVC(self.U[Slice], self.V[Slice])
+        return self.Q
 
     def __len__(self):
-        # if len(self.x.shape) == 3:
-        #    return self.x.shape[2]
-        # else:
-        return self.u.shape[2]
+        if self._is_list:
+            return self.U.shape[0]
+        return self.U.shape[self.t_axis]

docs/source/gallery/Nuke.ipynb

@@ -38,7 +38,7 @@
     "\n",
     "Z = np.sin(X*X+Y*Y-T)\n",
     "\n",
-    "block = amp.blocks.Pcolormesh(X[:,:,0], Y[:,:,0], Z, cmap='RdBu')\n",
+    "block = amp.blocks.Pcolormesh(X[:,:,0], Y[:,:,0], Z, t_axis=2, cmap='RdBu')\n",
     "plt.colorbar(block.quad)\n",
     "plt.gca().set_aspect('equal')\n",
     "\n",