Condense argument handling in fill_between, and add masked example

efiring · efiring · commit da5181d15363 · 2009-04-29T19:38:35.000Z
svn path=/trunk/matplotlib/; revision=7071
diff --git a/examples/pylab_examples/fill_between.py b/examples/pylab_examples/fill_between.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 import matplotlib.mlab as mlab
-from pylab import figure, show
+from matplotlib.pyplot import figure, show
 import numpy as np
 
 x = np.arange(0.0, 2, 0.01)
@@ -27,11 +27,24 @@
 #   fill_between(x[where], y1[where],y2[where]
 # because of edge effects over multiple contiguous regions.
 fig = figure()
-ax = fig.add_subplot(111)
+ax = fig.add_subplot(211)
 ax.plot(x, y1, x, y2, color='black')
-ax.fill_between(x, y1, y2, where=y2>y1, facecolor='green')
+ax.fill_between(x, y1, y2, where=y2>=y1, facecolor='green')
 ax.fill_between(x, y1, y2, where=y2<=y1, facecolor='red')
 ax.set_title('fill between where')
 
+# Test support for masked arrays.
+y2 = np.ma.masked_greater(y2, 1.0)
+ax1 = fig.add_subplot(212, sharex=ax)
+ax1.plot(x, y1, x, y2, color='black')
+ax1.fill_between(x, y1, y2, where=y2>=y1, facecolor='green')
+ax1.fill_between(x, y1, y2, where=y2<=y1, facecolor='red')
+ax1.set_title('Now regions with y2>1 are masked')
+
+# This example illustrates a problem; because of the data
+# gridding, there are undesired unfilled triangles at the crossover
+# points.  A brute-force solution would be to interpolate all
+# arrays to a very fine grid before plotting.
+
 show()
 
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -5832,39 +5832,28 @@ def fill_between(self, x, y1, y2=0, where=None, **kwargs):
         self._process_unit_info(xdata=x, ydata=y1, kwargs=kwargs)
         self._process_unit_info(ydata=y2)
 
-        if where is None:
-            where = np.ones(len(x), np.bool)
-        else:
-            where = np.asarray(where)
-
-        maskedx = isinstance(x, np.ma.MaskedArray)
-        maskedy1 = isinstance(y1, np.ma.MaskedArray)
-        maskedy2 = isinstance(y2, np.ma.MaskedArray)
-
-        if (maskedx or maskedy1 or maskedy2):
-            if maskedx:
-                where = where & (~x.mask)
-
-            if maskedy1:
-                where = where & (~y1.mask)
-
-            if maskedy2:
-                where = where & (~y2.mask)
-
-
         # Convert the arrays so we can work with them
-        x = np.asarray(self.convert_xunits(x))
-        y1 = np.asarray(self.convert_yunits(y1))
-        y2 = np.asarray(self.convert_yunits(y2))
+        x = np.asanyarray(self.convert_xunits(x))
+        y1 = np.asanyarray(self.convert_yunits(y1))
+        y2 = np.asanyarray(self.convert_yunits(y2))
 
-        if not cbook.iterable(y1):
+        if y1.ndim == 0:
             y1 = np.ones_like(x)*y1
-
-        if not cbook.iterable(y2):
+        if y2.ndim == 0:
             y2 = np.ones_like(x)*y2
 
+        if where is None:
+            where = np.ones(len(x), np.bool)
+        else:
+            where = np.asarray(where, np.bool)
+
+        if not (x.shape == y1.shape == y2.shape == where.shape):
+            raise ValueError("Argument dimensions are incompatible")
 
-        assert( (len(x)==len(y1)) and (len(x)==len(y2)) and len(x)==len(where))
+        mask = reduce(ma.mask_or,
+                        [ma.getmask(x), ma.getmask(y1), ma.getmask(y2)])
+        if mask is not ma.nomask:
+            where &= ~mask
 
         polys = []
         for ind0, ind1 in mlab.contiguous_regions(where):
