[Numpy-svn] r6299 - in trunk/numpy/ma: . tests

[numpy-svn at scipy.org]

Author: pierregm
Date: 2009-01-07 12:14:12 -0600 (Wed, 07 Jan 2009)
New Revision: 6299

Modified:
   trunk/numpy/ma/core.py
   trunk/numpy/ma/tests/test_core.py
Log:
* Fixed iadd/isub/imul when the base array has no mask but the other array does

Modified: trunk/numpy/ma/core.py
===================================================================
--- trunk/numpy/ma/core.py	2009-01-06 20:35:37 UTC (rev 6298)
+++ trunk/numpy/ma/core.py	2009-01-07 18:14:12 UTC (rev 6299)
@@ -2331,7 +2331,9 @@
         "Add other to self in-place."
         m = getmask(other)
         if self._mask is nomask:
-            self._mask = m
+            if m is not nomask and m.any():
+                self._mask = make_mask_none(self.shape, self.dtype)
+                self._mask += m
         else:
             if m is not nomask:
                 self._mask += m
@@ -2342,7 +2344,9 @@
         "Subtract other from self in-place."
         m = getmask(other)
         if self._mask is nomask:
-            self._mask = m
+            if m is not nomask and m.any():
+                self._mask = make_mask_none(self.shape, self.dtype)
+                self._mask += m
         elif m is not nomask:
             self._mask += m
         ndarray.__isub__(self._data, np.where(self._mask, 0, getdata(other)))
@@ -2352,7 +2356,9 @@
         "Multiply self by other in-place."
         m = getmask(other)
         if self._mask is nomask:
-            self._mask = m
+            if m is not nomask and m.any():
+                self._mask = make_mask_none(self.shape, self.dtype)
+                self._mask += m
         elif m is not nomask:
             self._mask += m
         ndarray.__imul__(self._data, np.where(self._mask, 1, getdata(other)))
@@ -3701,7 +3707,7 @@
 
     """
     _data = ndarray.__new__(baseclass, baseshape, basetype)
-    _mask = ndarray.__new__(ndarray, baseshape, 'b1')
+    _mask = ndarray.__new__(ndarray, baseshape, make_mask_descr(basetype))
     return subtype.__new__(subtype, _data, mask=_mask, dtype=basetype,)
 
 

Modified: trunk/numpy/ma/tests/test_core.py
===================================================================
--- trunk/numpy/ma/tests/test_core.py	2009-01-06 20:35:37 UTC (rev 6298)
+++ trunk/numpy/ma/tests/test_core.py	2009-01-07 18:14:12 UTC (rev 6299)
@@ -1400,6 +1400,51 @@
         assert_equal(x.data, [1., 2.**2.5, 3])
         assert_equal(x.mask, [0, 0, 1])
 
+
+    def test_datafriendly_add_arrays(self):
+        a = array([[1, 1], [3, 3]])
+        b = array([1, 1], mask=[0, 0])
+        a += b
+        assert_equal(a, [[2, 2], [4, 4]])
+        if a.mask is not nomask:
+            assert_equal(a.mask, [[0, 0], [0, 0]])
+        #
+        a = array([[1, 1], [3, 3]])
+        b = array([1, 1], mask=[0, 1])
+        a += b
+        assert_equal(a, [[2, 2], [4, 4]])
+        assert_equal(a.mask, [[0, 1], [0, 1]])
+
+
+    def test_datafriendly_sub_arrays(self):
+        a = array([[1, 1], [3, 3]])
+        b = array([1, 1], mask=[0, 0])
+        a -= b
+        assert_equal(a, [[0, 0], [2, 2]])
+        if a.mask is not nomask:
+            assert_equal(a.mask, [[0, 0], [0, 0]])
+        #
+        a = array([[1, 1], [3, 3]])
+        b = array([1, 1], mask=[0, 1])
+        a -= b
+        assert_equal(a, [[0, 0], [2, 2]])
+        assert_equal(a.mask, [[0, 1], [0, 1]])
+
+
+    def test_datafriendly_mul_arrays(self):
+        a = array([[1, 1], [3, 3]])
+        b = array([1, 1], mask=[0, 0])
+        a *= b
+        assert_equal(a, [[1, 1], [3, 3]])
+        if a.mask is not nomask:
+            assert_equal(a.mask, [[0, 0], [0, 0]])
+        #
+        a = array([[1, 1], [3, 3]])
+        b = array([1, 1], mask=[0, 1])
+        a *= b
+        assert_equal(a, [[1, 1], [3, 3]])
+        assert_equal(a.mask, [[0, 1], [0, 1]])
+
 #------------------------------------------------------------------------------
 
 class TestMaskedArrayMethods(TestCase):