Check for nan and inf in axes.delete_masked_points().

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

Author: astraw

CHANGELOG

@@ -1,3 +1,6 @@
+2008-07-18 Check for nan and inf in axes.delete_masked_points().
+	   This should help hexbin and scatter deal with nans. - ADS
+
 2008-07-17 Added ability to manually select contour label locations.
 	   Also added a waitforbuttonpress function. - DMK
 

lib/matplotlib/axes.py

@@ -1,5 +1,5 @@
 from __future__ import division, generators
-import math, sys, warnings, datetime, new
+import math, sys, warnings, datetime, new, types
 
 import numpy as np
 from numpy import ma
@@ -36,6 +36,8 @@ def delete_masked_points(*args):
     Find all masked 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).
+
     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
@@ -49,9 +51,11 @@ def delete_masked_points(*args):
     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(ma.mask_or, masks)
+    mask = reduce(np.logical_or, masks)
     margs = []
     for x in args:
         if (not is_string_like(x)

unit/axes_unit.py

@@ -0,0 +1,62 @@
+import unittest
+import numpy as np
+import matplotlib.axes as axes
+
+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 = axes.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 = axes.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 = axes.delete_masked_points(*input)
+        assert np.allclose(actual, expected)
+
+        input = ()
+        expected = ()
+        actual = axes.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 = axes.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 = axes.delete_masked_points(*input)
+        assert np.allclose(actual, expected)
+
+    def test_delete_masked_points_strings(self):
+        input = (   'hello',
+                    )
+        expected = ('hello',)
+        actual = axes.delete_masked_points(*input)
+        assert actual == expected
+
+        input = (   u'hello',
+                    )
+        expected = (u'hello',)
+        actual = axes.delete_masked_points(*input)
+        assert actual == expected
+
+
+if __name__=='__main__':
+    unittest.main()