Merged revisions 7616-7618 via svnmerge from

dopplershift · dopplershift · commit cf3a22e0799e · 2009-08-31T17:42:37.000Z
https://matplotlib.svn.sourceforge.net/svnroot/matplotlib/branches/v0_99_maint ........ r7616 | ryanmay | 2009-08-31 12:27:12 -0500 (Mon, 31 Aug 2009) | 1 line Update barb_demo.py with an example using masked arrays. ........ r7617 | ryanmay | 2009-08-31 12:29:41 -0500 (Mon, 31 Aug 2009) | 1 line Pull _parse_args out of Quiver and Barbs classes, as we have multiple, diverging copies of the same code. Now the shared code will actually be shared. ........ r7618 | ryanmay | 2009-08-31 12:34:50 -0500 (Mon, 31 Aug 2009) | 1 line Use atleast_1d instead of asanyarray. This allows passing in scalars to quiver/barbs. ........ svn path=/trunk/matplotlib/; revision=7619
diff --git a/examples/pylab_examples/barb_demo.py b/examples/pylab_examples/barb_demo.py
@@ -28,7 +28,7 @@
 ax.barbs(data['x'], data['y'], data['u'], data['v'], length=8, pivot='middle')
 
 #Showing colormapping with uniform grid.  Fill the circle for an empty barb,
-#don't round the values, and change some of the size parameters 
+#don't round the values, and change some of the size parameters
 ax = plt.subplot(2,2,3)
 ax.barbs(X, Y, U, V, np.sqrt(U*U + V*V), fill_empty=True, rounding=False,
     sizes=dict(emptybarb=0.25, spacing=0.2, height=0.3))
@@ -39,4 +39,15 @@
     barbcolor=['b','g'], barb_increments=dict(half=10, full=20, flag=100),
     flip_barb=True)
 
+#Masked arrays are also supported
+masked_u = np.ma.masked_array(data['u'])
+masked_u[4] = 1000 #Bad value that should not be plotted when masked
+masked_u[4] = np.ma.masked
+
+#Identical plot to panel 2 in the first figure, but with the point at
+#(0.5, 0.25) missing (masked)
+fig2 = plt.figure()
+ax = fig2.add_subplot(1, 1, 1)
+ax.barbs(data['x'], data['y'], masked_u, data['v'], length=8, pivot='middle')
+
 plt.show()
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -325,6 +325,33 @@ def contains(self, mouseevent):
     quiverkey_doc = _quiverkey_doc
 
 
+# This is a helper function that parses out the various combination of
+# arguments for doing colored vector plots.  Pulling it out here
+# allows both Quiver and Barbs to use it
+def _parse_args(*args):
+    X, Y, U, V, C = [None]*5
+    args = list(args)
+
+    # The use of atleast_1d allows for handling scalar arguments while also
+    # keeping masked arrays
+    if len(args) == 3 or len(args) == 5:
+        C = np.atleast_1d(args.pop(-1))
+    V = np.atleast_1d(args.pop(-1))
+    U = np.atleast_1d(args.pop(-1))
+    if U.ndim == 1:
+        nr, nc = 1, U.shape[0]
+    else:
+        nr, nc = U.shape
+    if len(args) == 2: # remaining after removing U,V,C
+        X, Y = [np.array(a).ravel() for a in args]
+        if len(X) == nc and len(Y) == nr:
+            X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
+    else:
+        indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
+        X, Y = [np.ravel(a) for a in indexgrid]
+    return X, Y, U, V, C
+
+
 class Quiver(collections.PolyCollection):
     """
     Specialized PolyCollection for arrows.
@@ -351,7 +378,7 @@ def __init__(self, ax, *args, **kw):
         by the following pylab interface documentation:
         %s"""
         self.ax = ax
-        X, Y, U, V, C = self._parse_args(*args)
+        X, Y, U, V, C = _parse_args(*args)
         self.X = X
         self.Y = Y
         self.XY = np.hstack((X[:,np.newaxis], Y[:,np.newaxis]))
@@ -390,26 +417,6 @@ def on_dpi_change(fig):
         self.ax.figure.callbacks.connect('dpi_changed', on_dpi_change)
 
 
-    def _parse_args(self, *args):
-        X, Y, U, V, C = [None]*5
-        args = list(args)
-        if len(args) == 3 or len(args) == 5:
-            C = np.asanyarray(args.pop(-1))
-        V = np.asanyarray(args.pop(-1))
-        U = np.asanyarray(args.pop(-1))
-        if U.ndim == 1:
-            nr, nc = 1, U.shape[0]
-        else:
-            nr, nc = U.shape
-        if len(args) == 2: # remaining after removing U,V,C
-            X, Y = [np.array(a).ravel() for a in args]
-            if len(X) == nc and len(Y) == nr:
-                X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
-        else:
-            indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
-            X, Y = [np.ravel(a) for a in indexgrid]
-        return X, Y, U, V, C
-
     def _init(self):
         """initialization delayed until first draw;
         allow time for axes setup.
@@ -762,7 +769,7 @@ def __init__(self, ax, *args, **kw):
             kw['facecolors'] = flagcolor
 
         #Parse out the data arrays from the various configurations supported
-        x, y, u, v, c = self._parse_args(*args)
+        x, y, u, v, c = _parse_args(*args)
         self.x = x
         self.y = y
         xy = np.hstack((x[:,np.newaxis], y[:,np.newaxis]))
@@ -942,28 +949,6 @@ def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
 
         return barb_list
 
-    #Taken shamelessly from Quiver
-    def _parse_args(self, *args):
-        X, Y, U, V, C = [None]*5
-        args = list(args)
-        if len(args) == 3 or len(args) == 5:
-            C = ma.masked_invalid(args.pop(-1), copy=False).ravel()
-        V = ma.masked_invalid(args.pop(-1), copy=False)
-        U = ma.masked_invalid(args.pop(-1), copy=False)
-        nn = np.shape(U)
-        nc = nn[0]
-        nr = 1
-        if len(nn) > 1:
-            nr = nn[1]
-        if len(args) == 2: # remaining after removing U,V,C
-            X, Y = [np.array(a).ravel() for a in args]
-            if len(X) == nc and len(Y) == nr:
-                X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
-        else:
-            indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
-            X, Y = [np.ravel(a) for a in indexgrid]
-        return X, Y, U, V, C
-
     def set_UVC(self, U, V, C=None):
         self.u = ma.masked_invalid(U, copy=False).ravel()
         self.v = ma.masked_invalid(V, copy=False).ravel()
