Use numerix more in collections; quiver handles masked arrays.

efiring · efiring · commit da94b4669014 · 2006-07-26T16:48:39.000Z
svn path=/trunk/matplotlib/; revision=2622
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -9,7 +9,7 @@
 from cbook import is_string_like, enumerate, strip_math, Stack
 from colors import colorConverter
 from numerix import array, sqrt, pi, log, asarray, ones, zeros, Float, Float32
-from numerix import arange, compress, take
+from numerix import arange, compress, take, isnan, any
 from patches import Rectangle
 from transforms import lbwh_to_bbox, identity_transform
 import widgets
@@ -92,7 +92,7 @@ def draw_line_collection(self, segments, transform, clipbox,
         matplotlib.collections for more details.
 
         segments is a sequence of ( line0, line1, line2), where linen =
-        (x0, y0), (x1, y1), ... (xm, ym).  Each line can be a
+        is an Mx2 array with columns x, y.  Each line can be a
         different length
 
         transform is used to Transform the lines
@@ -110,8 +110,9 @@ def draw_line_collection(self, segments, transform, clipbox,
         antialiseds is a tuple of ones or zeros indicating whether the
         segment should be aa or not
 
-        offsets, if not None, is a list of x,y offsets to translate the lines
-        by after transform is used to transform the offset coords
+        offsets, if not None, is an Nx2 array of x,y offsets to
+        translate the lines by after transform is used to transform
+        the offset coords
 
         This function could be overridden in the backend to possibly implement
         faster drawing, but it is already much faster than using draw_lines()
@@ -133,7 +134,8 @@ def draw_line_collection(self, segments, transform, clipbox,
         usingOffsets = offsets is not None
         Noffsets  = 0
         if usingOffsets:
-            Noffsets = len(offsets)
+            Noffsets = offsets.shape[0]
+            offsets = transOffset.numerix_xy(offsets)
 
         for i in xrange(max(Noffsets, Nsegments)):
             color = colors[i % Nc]
@@ -146,16 +148,12 @@ def draw_line_collection(self, segments, transform, clipbox,
             gc.set_antialiased( antialiaseds[i % Naa] )
             seg = segments[i % Nsegments]
             if not len(seg): continue
-            x, y = zip(*seg)
-
-            x, y = transform.numerix_x_y(array(x), array(y))
+            xy = transform.numerix_xy(seg)
             if usingOffsets:
-                xo, yo = transOffset.xy_tup(offsets[i % Noffsets])
-                x += xo
-                y += yo
+                xy = xy + offsets[i % Noffsets]
 
-            if newstyle: self.draw_lines(gc, x, y, identity)
-            else: self.draw_lines(gc, x, y)
+            if newstyle: self.draw_lines(gc, xy[:,0], xy[:,1], identity)
+            else: self.draw_lines(gc, xy[:,0], xy[:,1])
 
     def draw_line(self, gc, x1, y1, x2, y2):
         """
@@ -241,7 +239,9 @@ def draw_poly_collection(
             gc.set_clip_rectangle(clipbox.get_bounds())
 
         for i in xrange(N):
-
+            polyverts = verts[i % Nverts]
+            if any(isnan(polyverts)):
+                continue
             linewidth = linewidths[i % Nlw]
             rf,gf,bf,af = facecolors[i % Nface]
             re,ge,be,ae = edgecolors[i % Nedge]
@@ -265,7 +265,7 @@ def draw_poly_collection(
 
             gc.set_antialiased( antialiaseds[i % Naa] )  # Used for fill only?
             gc.set_alpha( alpha )
-            tverts = transform.seq_xy_tups(verts[i % Nverts])
+            tverts = transform.seq_xy_tups(polyverts)
             if usingOffsets:
                 xo,yo = transOffset.xy_tup(offsets[i % Noffsets])
                 tverts = [(x+xo,y+yo) for x,y in tverts]
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -15,7 +15,7 @@
 from cbook import is_string_like, iterable
 from colors import colorConverter
 from cm import ScalarMappable
-from numerix import arange, sin, cos, pi, asarray, sqrt, array
+from numerix import arange, sin, cos, pi, asarray, sqrt, array, newaxis
 from transforms import identity_transform
 
 class Collection(Artist):
@@ -393,7 +393,8 @@ def __init__(self, segments,     # Can be None.
                  ):
         """
         segments is a sequence of ( line0, line1, line2), where
-        linen = (x0, y0), (x1, y1), ... (xm, ym).
+        linen = (x0, y0), (x1, y1), ... (xm, ym), or the
+        equivalent numerix array with two columns.
         Each line can be a different length.
 
         colors must be a tuple of RGBA tuples (eg arbitrary color
@@ -443,6 +444,10 @@ def __init__(self, segments,     # Can be None.
         self._lw = linewidths
         self.set_linestyle(linestyle)
         self._uniform_offsets = None
+        if offsets is not None:
+            offsets = asarray(offsets)
+            if len(offsets.shape) == 1:
+                offsets = offsets[newaxis,:]  # Make it Nx2.
         if transOffset is None:
             if offsets is not None:
                 self._uniform_offsets = offsets
@@ -454,7 +459,7 @@ def __init__(self, segments,     # Can be None.
 
     def set_segments(self, segments):
         if segments is None: return
-        self._segments = list(segments)
+        self._segments = [asarray(seg) for seg in segments]
         if self._uniform_offsets is not None:
             self._add_offsets()
 
@@ -464,15 +469,14 @@ def _add_offsets(self):
         segs = self._segments
         offsets = self._uniform_offsets
         Nsegs = len(segs)
-        Noffs = len(offsets)
-        if not iterable(offsets[0]):  # i.e., not a tuple but an x-offset
-            xo, yo = offsets
+        Noffs = offsets.shape[0]
+        if Noffs == 1:
             for i in range(Nsegs):
-                segs[i] = [(x+xo*i, y+yo*i) for x,y in segs[i]]
+                segs[i] = segs[i] + i * offsets
         else:
             for i in range(Nsegs):
-                xo, yo = offsets[i%Noffs]
-                segs[i] = [(x+xo, y+yo) for x,y in segs[i]]
+                io = i%Noffs
+                segs[i] = segs[i] + offsets[io:io+1]
 
 
     def draw(self, renderer):
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -296,8 +296,8 @@ def _parse_args(self, *args):
         if len(args) == 3 or len(args) == 5:
             C = nx.ravel(args.pop(-1))
             #print 'in parse_args, C:', C
-        V = nx.asarray(args.pop(-1))
-        U = nx.asarray(args.pop(-1))
+        V = nx.ma.asarray(args.pop(-1))
+        U = nx.ma.asarray(args.pop(-1))
         nn = nx.shape(U)
         nc = nn[0]
         nr = 1
@@ -333,8 +333,8 @@ def draw(self, renderer):
         PolyCollection.draw(self, renderer)
 
     def set_UVC(self, U, V, C=None):
-        self.U = nx.ravel(U)
-        self.V = nx.ravel(V)
+        self.U = nx.ma.ravel(U)
+        self.V = nx.ma.ravel(V)
         if C is not None:
             self.set_array(nx.ravel(C))
         self._new_UV = True
@@ -367,6 +367,7 @@ def _set_transform(self):
 
     def _make_verts(self, U, V):
         uv = U+V*1j
+        uv = nx.ravel(nx.ma.filled(uv, nx.nan))
         a = nx.absolute(uv)
         if self.scale is None:
             sn = max(10, math.sqrt(self.N))
@@ -379,7 +380,6 @@ def _make_verts(self, U, V):
         xy = xy[:,:,nx.newaxis]
         XY = nx.concatenate((xy.real, xy.imag), axis=2)
         return XY
-        #return [zip(xyrow.real, xyrow.imag) for xyrow in xy]
 
 
     def _h_arrows(self, length):
diff --git a/src/_transforms.cpp b/src/_transforms.cpp
@@ -961,14 +961,17 @@ Transformation::numerix_xy(const Py::Tuple & args) {
   size_t Nxy = xyin->dimensions[0];
   size_t N2 = xyin->dimensions[1];
 
-  if (N2!=2)
+  if (N2!=2) {
+    Py_XDECREF(xyin);
     throw Py::ValueError("xy must have shape (N,2)");
+  }
 
   // evaluate the lazy objects
   try {
     if (!_frozen) eval_scalars();
   }
   catch(...) {
+    Py_XDECREF(xyin);
     throw Py::ValueError("Domain error on Transformation::numerix_xy");
   }
 
