Expanded delete_masked_points to handle more types of argument

svn path=/trunk/matplotlib/; revision=5797

Author: efiring

CHANGELOG

@@ -1,3 +1,6 @@
+2008-07-20 Rewrote cbook.delete_masked_points and corresponding
+           unit test to support rgb color array inputs, datetime
+           inputs, etc. - EF
 
 2008-07-20 Renamed unit/axes_unit.py to cbook_unit.py and modified
            in accord with Ryan's move of delete_masked_points from

lib/matplotlib/cbook.py

@@ -6,7 +6,7 @@
 import re, os, errno, sys, StringIO, traceback, locale, threading, types
 import time, datetime
 import numpy as np
-from numpy import ma
+import numpy.ma as ma
 from weakref import ref
 
 major, minor1, minor2, s, tmp = sys.version_info
@@ -1168,41 +1168,82 @@ def recursive_remove(path):
 
 def delete_masked_points(*args):
     """
-    Find all masked points in a set of arguments, and return
-    the arguments with only the unmasked points remaining.
+    Find all masked and/or non-finite points in a set of arguments,
+    and return the arguments with only the unmasked points remaining.
 
-    This will also delete any points that are not finite (nan or inf).
+    Arguments can be in any of 5 categories:
 
-    The overall mask is calculated from any masks that are present.
-    If a mask is found, any argument that does not have the same
-    dimensions is left unchanged; therefore the argument list may
-    include arguments that can take string or array values, for
-    example.
+    1) 1-D masked arrays
+    2) 1-D ndarrays
+    3) ndarrays with more than one dimension
+    4) other non-string iterables
+    5) anything else
 
-    Array arguments must have the same length; masked arguments must
-    be one-dimensional.
+    The first argument must be in one of the first four categories;
+    any argument with a length differing from that of the first
+    argument (and hence anything in category 5) then will be
+    passed through unchanged.
+
+    Masks are obtained from all arguments of the correct length
+    in categories 1, 2, and 4; a point is bad if masked in a masked
+    array or if it is a nan or inf.  No attempt is made to
+    extract a mask from categories 2, 3, and 4 if *np.isfinite()*
+    does not yield a Boolean array.
+
+    All input arguments that are not passed unchanged are returned
+    as ndarrays after removing the points or rows corresponding to
+    masks in any of the arguments.
+
+    A vastly simpler version of this function was originally
+    written as a helper for Axes.scatter().
 
-    Written as a helper for scatter, but may be more generally
-    useful.
     """
-    masks = [ma.getmaskarray(x) for x in args if hasattr(x, 'mask')]
-    isfinite = [np.isfinite(x) for x in args]
-    masks.extend( [~x for x in isfinite if not isinstance(x,types.NotImplementedType)] )
-    if len(masks) == 0:
-        return args
-    mask = reduce(np.logical_or, masks)
+    if not len(args):
+        return ()
+    if (is_string_like(args[0]) or not iterable(args[0])):
+        raise ValueError("First argument must be a sequence")
+    nrecs = len(args[0])
     margs = []
-    for x in args:
-        if (not is_string_like(x)
-            and iterable(x)
-            and len(x) == len(mask)):
-            if (hasattr(x, 'get_compressed_copy')):
-                compressed_x = x.get_compressed_copy(mask)
+    seqlist = [False] * len(args)
+    for i, x in enumerate(args):
+        if (not is_string_like(x)) and iterable(x) and len(x) == nrecs:
+            seqlist[i] = True
+            if ma.isMA(x):
+                if x.ndim > 1:
+                    raise ValueError("Masked arrays must be 1-D")
             else:
-                compressed_x = ma.masked_array(x, mask=mask).compressed()
-            margs.append(compressed_x)
-        else:
-            margs.append(x)
+                x = np.asarray(x)
+        margs.append(x)
+    masks = []    # list of masks that are True where good
+    for i, x in enumerate(margs):
+        if seqlist[i]:
+            if x.ndim > 1:
+                continue  # Don't try to get nan locations unless 1-D.
+            if ma.isMA(x):
+                masks.append(~ma.getmaskarray(x))  # invert the mask
+                xd = x.data
+            else:
+                xd = x
+            try:
+                mask = np.isfinite(xd)
+                if isinstance(mask, np.ndarray):
+                    masks.append(mask)
+            except: #Fixme: put in tuple of possible exceptions?
+                pass
+    if len(masks):
+        mask = reduce(np.logical_and, masks)
+        igood = mask.nonzero()[0]
+        if len(igood) < nrecs:
+            for i, x in enumerate(margs):
+                if seqlist[i]:
+                    if (hasattr(x, 'get_compressed_copy')):
+                        compressed_x = x.get_compressed_copy(~mask)
+                    else:
+                        compressed_x = x.take(igood, axis=0)
+                    margs[i] = compressed_x
+    for i, x in enumerate(margs):
+        if seqlist[i] and ma.isMA(x):
+            margs[i] = x.filled()
     return margs
 
 

lib/matplotlib/collections.py

@@ -90,7 +90,7 @@ def __init__(self,
         self._uniform_offsets = None
         self._offsets = np.array([], np.float_)
         if offsets is not None:
-            offsets = np.asarray(offsets, np.float_)
+            offsets = np.asarray(offsets)
             if len(offsets.shape) == 1:
                 offsets = offsets[np.newaxis,:]  # Make it Nx2.
             if transOffset is not None:

unit/cbook_unit.py

@@ -1,61 +1,51 @@
 import unittest
+from datetime import datetime
+
 import numpy as np
 import matplotlib.cbook as cbook
-
-class TestAxes(unittest.TestCase):
-    def test_delete_masked_points_arrays(self):
-        input = (   [1,2,3,np.nan,5],
-                    np.array((1,2,3,4,5)),
-                    )
-        expected = [np.array((1,2,3,5))]*2
-        actual = cbook.delete_masked_points(*input)
-        assert np.allclose(actual, expected)
-
-        input = (   np.ma.array( [1,2,3,4,5], mask=[False,False,False,True,False] ),
-                    np.array((1,2,3,4,5)),
-                    )
-        expected = [np.array((1,2,3,5))]*2
-        actual = cbook.delete_masked_points(*input)
-        assert np.allclose(actual, expected)
-
-        input = (   [1,2,3,np.nan,5],
-                    np.ma.array( [1,2,3,4,5], mask=[False,False,False,True,False] ),
-                    np.array((1,2,3,4,5)),
-                    )
-        expected = [np.array((1,2,3,5))]*3
-        actual = cbook.delete_masked_points(*input)
-        assert np.allclose(actual, expected)
-
-        input = ()
-        expected = ()
-        actual = cbook.delete_masked_points(*input)
-        assert np.allclose(actual, expected)
-
-
-        input = (   [1,2,3,np.nan,5],
-                    )
-        expected = [np.array((1,2,3,5))]*1
-        actual = cbook.delete_masked_points(*input)
-        assert np.allclose(actual, expected)
-
-        input = (   np.array((1,2,3,4,5)),
-                    )
-        expected = [np.array((1,2,3,4,5))]*1
-        actual = cbook.delete_masked_points(*input)
-        assert np.allclose(actual, expected)
-
-    def test_delete_masked_points_strings(self):
-        input = (   'hello',
-                    )
-        expected = ('hello',)
-        actual = cbook.delete_masked_points(*input)
-        assert actual == expected
-
-        input = (   u'hello',
-                    )
-        expected = (u'hello',)
-        actual = cbook.delete_masked_points(*input)
-        assert actual == expected
+import matplotlib.colors as mcolors
+
+from matplotlib.cbook import delete_masked_points as dmp
+
+class Test_delete_masked_points(unittest.TestCase):
+    def setUp(self):
+        self.mask1 = [False, False, True, True, False, False]
+        self.arr0 = np.arange(1.0,7.0)
+        self.arr1 = [1,2,3,np.nan,np.nan,6]
+        self.arr2 = np.array(self.arr1)
+        self.arr3 = np.ma.array(self.arr2, mask=self.mask1)
+        self.arr_s = ['a', 'b', 'c', 'd', 'e', 'f']
+        self.arr_s2 = np.array(self.arr_s)
+        self.arr_dt = [datetime(2008, 1, 1), datetime(2008, 1, 2),
+                       datetime(2008, 1, 3), datetime(2008, 1, 4),
+                       datetime(2008, 1, 5), datetime(2008, 1, 6)]
+        self.arr_dt2 = np.array(self.arr_dt)
+        self.arr_colors = ['r', 'g', 'b', 'c', 'm', 'y']
+        self.arr_rgba = mcolors.colorConverter.to_rgba_array(self.arr_colors)
+
+    def test_bad_first_arg(self):
+        self.assertRaises(ValueError, dmp, 'a string', self.arr0)
+
+    def test_string_seq(self):
+        actual = dmp(self.arr_s, self.arr1)
+        ind = [0, 1, 2, 5]
+        expected = (self.arr_s2.take(ind), self.arr2.take(ind))
+
+    def test_datetime(self):
+        actual = dmp(self.arr_dt, self.arr3)
+        ind = [0, 1,  5]
+        expected = (self.arr_dt2.take(ind),
+                    self.arr3.take(ind).compressed())
+        self.assert_(np.all(actual[0] == expected[0]) and
+                     np.all(actual[1] == expected[1]))
+
+    def test_rgba(self):
+        actual = dmp(self.arr3, self.arr_rgba)
+        ind = [0, 1, 5]
+        expected = (self.arr3.take(ind).compressed(),
+                    self.arr_rgba.take(ind, axis=0))
+        self.assert_(np.all(actual[0] == expected[0]) and
+                     np.all(actual[1] == expected[1]))
 
 
 if __name__=='__main__':