Make safezip accept more args; quadmesh_demo cleanup

efiring · efiring · commit 93676b7ad458 · 2007-11-05T20:13:00.000Z
svn path=/trunk/matplotlib/; revision=4114
diff --git a/examples/quadmesh_demo.py b/examples/quadmesh_demo.py
@@ -2,23 +2,26 @@
 """
 pcolormesh uses a QuadMesh, a faster generalization of pcolor, but
 with some restrictions.
+
+This demo illustrates a bug in quadmesh with masked data.
 """
 
-import numpy as nx
-from pylab import figure,show
+import numpy as npy
+from matplotlib.pyplot import figure, show
 from matplotlib import cm, colors
+from matplotlib.numerix import npyma as ma
 
 n = 56
-x = nx.linspace(-1.5,1.5,n)
-X,Y = nx.meshgrid(x,x);
-Qx = nx.cos(Y) - nx.cos(X)
-Qz = nx.sin(Y) + nx.sin(X)
+x = npy.linspace(-1.5,1.5,n)
+X,Y = npy.meshgrid(x,x);
+Qx = npy.cos(Y) - npy.cos(X)
+Qz = npy.sin(Y) + npy.sin(X)
 Qx = (Qx + 1.1)
-Z = nx.sqrt(X**2 + Y**2)/5;
-Z = (Z - nx.amin(Z)) / (nx.amax(Z) - nx.amin(Z))
+Z = npy.sqrt(X**2 + Y**2)/5;
+Z = (Z - Z.min()) / (Z.max() - Z.min())
 
 # The color array can include masked values:
-Zm = nx.ma.masked_where(nx.fabs(Qz) < 0.5*nx.amax(Qz), Z)
+Zm = ma.masked_where(npy.fabs(Qz) < 0.5*npy.amax(Qz), Z)
 
 
 fig = figure()
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -557,7 +557,7 @@ def dedent(s):
     # expressions.  However, this function accounted for almost 30% of
     # matplotlib startup time, so it is worthy of optimization at all
     # costs.
-    
+
     if not s:      # includes case of s is None
         return ''
 
@@ -576,7 +576,7 @@ def dedent(s):
     if unindent is None:
         unindent = re.compile("\n\r? {0,%d}" % nshift)
         _dedent_regex[nshift] = unindent
-        
+
     result = unindent.sub("\n", s).strip()
     return result
 
@@ -844,15 +844,15 @@ def report_memory(i=0):  # argument may go away
 
     return mem
 
+_safezip_msg = 'In safezip, len(args[0])=%d but len(args[%d])=%d'
+def safezip(*args):
+    'make sure args are equal len before zipping'
+    Nx = len(args[0])
+    for i, arg in enumerate(args[1:]):
+        if len(arg) != Nx:
+            raise ValueError(_safezip_msg % (Nx, i+1, len(arg)))
+    return zip(*args)
 
-def safezip(x, y):
-    'make sure x and y are equal len before zipping'
-    Nx = len(x)
-    Ny = len(y)
-    if Nx!=Ny:
-        raise RuntimeError('x and y must be equal length; found len(x)=%d and len(y)=%d'%
-                           (Nx, Ny))
-    return zip(x, y)
 class MemoryMonitor:
     def __init__(self, nmax=20000):
         self._nmax = nmax
