testutils:

assert_array_compare : make sure that the comparison is performed on ndarrays, and make sure we use the np version of the comparison function.
core:
* Try not to touch the data in unary/binary ufuncs, (including inplace)

3 files changed, 218 insertions, 67 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index a3cda941d..6b4dc98e6 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -535,17 +535,20 @@ class _MaskedUnaryOperation:
             # ... but np.putmask looks more efficient, despite the copy.
             np.putmask(d1, dm, self.fill)
         # Take care of the masked singletong first ...
-        if not m.ndim and m:
+        if (not m.ndim) and m:
             return masked
-        # Get the result class .......................
-        if isinstance(a, MaskedArray):
-            subtype = type(a)
+        elif m is nomask:
+            result = self.f(d1, *args, **kwargs)
         else:
-            subtype = MaskedArray
-        # Get the result  as a view of the subtype ...
-        result = self.f(d1, *args, **kwargs).view(subtype)
-        # Fix the mask if we don't have a scalar
-        if result.ndim > 0:
+            result = np.where(m, d1, self.f(d1, *args, **kwargs))
+        # If result is not a scalar
+        if result.ndim:
+            # Get the result subclass:
+            if isinstance(a, MaskedArray):
+                subtype = type(a)
+            else:
+                subtype = MaskedArray
+            result = result.view(subtype)
             result._mask = m
             result._update_from(a)
         return result
@@ -584,19 +587,45 @@ class _MaskedBinaryOperation:
     def __call__ (self, a, b, *args, **kwargs):
         "Execute the call behavior."
         m = mask_or(getmask(a), getmask(b))
-        (d1, d2) = (get_data(a), get_data(b))
-        result = self.f(d1, d2, *args, **kwargs).view(get_masked_subclass(a, b))
-        if len(result.shape):
-            if m is not nomask:
-                result._mask = make_mask_none(result.shape)
-                result._mask.flat = m
+        (da, db) = (getdata(a), getdata(b))
+        # Easy case: there's no mask...
+        if m is nomask:
+            result = self.f(da, db, *args, **kwargs)
+        # There are some masked elements: run only on the unmasked
+        else:
+            result = np.where(m, da, self.f(da, db, *args, **kwargs))
+        # Transforms to a (subclass of) MaskedArray if we don't have a scalar
+        if result.shape:
+            result = result.view(get_masked_subclass(a, b))
+            result._mask = make_mask_none(result.shape)
+            result._mask.flat = m
             if isinstance(a, MaskedArray):
                 result._update_from(a)
             if isinstance(b, MaskedArray):
                 result._update_from(b)
+        # ... or return masked if we have a scalar and the common mask is True
         elif m:
             return masked
         return result
+#        
+#        result = self.f(d1, d2, *args, **kwargs).view(get_masked_subclass(a, b))
+#        if len(result.shape):
+#            if m is not nomask:
+#                result._mask = make_mask_none(result.shape)
+#                result._mask.flat = m
+#                #!!!!!
+#                # Force m to be at least 1D
+#                m.shape = m.shape or (1,)
+#                print "Resetting data"
+#                result.data[m].flat = d1.flat
+#                #!!!!!
+#            if isinstance(a, MaskedArray):
+#                result._update_from(a)
+#            if isinstance(b, MaskedArray):
+#                result._update_from(b)
+#        elif m:
+#            return masked
+#        return result
 
     def reduce(self, target, axis=0, dtype=None):
         """Reduce `target` along the given `axis`."""
@@ -639,11 +668,13 @@ class _MaskedBinaryOperation:
             m = umath.logical_or.outer(ma, mb)
         if (not m.ndim) and m:
             return masked
-        rcls = get_masked_subclass(a, b)
-        # We could fill the arguments first, butis it useful ?
-        # d = self.f.outer(filled(a, self.fillx), filled(b, self.filly)).view(rcls)
-        d = self.f.outer(getdata(a), getdata(b)).view(rcls)
-        if d.ndim > 0:
+        (da, db) = (getdata(a), getdata(b))
+        if m is nomask:
+            d = self.f.outer(da, db)
+        else:
+            d = np.where(m, da, self.f.outer(da, db))
+        if d.shape:
+            d = d.view(get_masked_subclass(a, b))
             d._mask = m
         return d
 
@@ -655,7 +686,7 @@ class _MaskedBinaryOperation:
         if isinstance(target, MaskedArray):
             tclass = type(target)
         else:
-            tclass = masked_array
+            tclass = MaskedArray
         t = filled(target, self.filly)
         return self.f.accumulate(t, axis).view(tclass)
 
@@ -664,7 +695,8 @@ class _MaskedBinaryOperation:
 
 #..............................................................................
 class _DomainedBinaryOperation:
-    """Define binary operations that have a domain, like divide.
+    """
+    Define binary operations that have a domain, like divide.
 
     They have no reduce, outer or accumulate.
 
@@ -689,25 +721,29 @@ class _DomainedBinaryOperation:
         ufunc_domain[dbfunc] = domain
         ufunc_fills[dbfunc] = (fillx, filly)
 
-    def __call__(self, a, b):
+    def __call__(self, a, b, *args, **kwargs):
         "Execute the call behavior."
         ma = getmask(a)
         mb = getmask(b)
-        d1 = getdata(a)
-        d2 = get_data(b)
-        t = narray(self.domain(d1, d2), copy=False)
+        da = getdata(a)
+        db = getdata(b)
+        t = narray(self.domain(da, db), copy=False)
         if t.any(None):
             mb = mask_or(mb, t)
             # The following line controls the domain filling
-            if t.size == d2.size:
-                d2 = np.where(t, self.filly, d2)
+            if t.size == db.size:
+                db = np.where(t, self.filly, db)
             else:
-                d2 = np.where(np.resize(t, d2.shape), self.filly, d2)
+                db = np.where(np.resize(t, db.shape), self.filly, db)
         m = mask_or(ma, mb)
         if (not m.ndim) and m:
             return masked
-        result =  self.f(d1, d2).view(get_masked_subclass(a, b))
-        if result.ndim > 0:
+        elif (m is nomask):
+            result = self.f(da, db, *args, **kwargs)
+        else:
+            result = np.where(m, da, self.f(da, db, *args, **kwargs))
+        if result.shape:
+            result = result.view(get_masked_subclass(a, b))
             result._mask = m
             if isinstance(a, MaskedArray):
                 result._update_from(a)
@@ -2243,32 +2279,33 @@ masked_%(name)s(data = %(data)s,
     #............................................
     def __iadd__(self, other):
         "Add other to self in-place."
-        ndarray.__iadd__(self._data, getdata(other))
         m = getmask(other)
         if self._mask is nomask:
             self._mask = m
-        elif m is not nomask:
-            self._mask += m
+        else:
+            if m is not nomask:
+                self._mask += m
+        ndarray.__iadd__(self._data, np.where(self._mask, 0, getdata(other)))
         return self
     #....
     def __isub__(self, other):
         "Subtract other from self in-place."
-        ndarray.__isub__(self._data, getdata(other))
         m = getmask(other)
         if self._mask is nomask:
             self._mask = m
         elif m is not nomask:
             self._mask += m
+        ndarray.__isub__(self._data, np.where(self._mask, 0, getdata(other)))
         return self
     #....
     def __imul__(self, other):
         "Multiply self by other in-place."
-        ndarray.__imul__(self._data, getdata(other))
         m = getmask(other)
         if self._mask is nomask:
             self._mask = m
         elif m is not nomask:
             self._mask += m
+        ndarray.__imul__(self._data, np.where(self._mask, 1, getdata(other)))
         return self
     #....
     def __idiv__(self, other):
@@ -2281,21 +2318,25 @@ masked_%(name)s(data = %(data)s,
         if dom_mask.any():
             (_, fval) = ufunc_fills[np.divide]
             other_data = np.where(dom_mask, fval, other_data)
-        ndarray.__idiv__(self._data, other_data)
-        self._mask = mask_or(self._mask, new_mask)
+#        self._mask = mask_or(self._mask, new_mask)
+        self._mask |= new_mask
+        ndarray.__idiv__(self._data, np.where(self._mask, 1, other_data))
         return self
     #...
     def __ipow__(self, other):
         "Raise self to the power other, in place"
-        _data = self._data
         other_data = getdata(other)
         other_mask = getmask(other)
-        ndarray.__ipow__(_data, other_data)
-        invalid = np.logical_not(np.isfinite(_data))
+        ndarray.__ipow__(self._data, np.where(self._mask, 1, other_data))
+        invalid = np.logical_not(np.isfinite(self._data))
+        if invalid.any():
+            if self._mask is not nomask:
+                self._mask |= invalid
+            else:
+                self._mask = invalid
+            np.putmask(self._data, invalid, self.fill_value)
         new_mask = mask_or(other_mask, invalid)
         self._mask = mask_or(self._mask, new_mask)
-        # The following line is potentially problematic, as we change _data...
-        np.putmask(self._data, invalid, self.fill_value)
         return self
     #............................................
     def __float__(self):
@@ -3848,22 +3889,22 @@ def power(a, b, third=None):
     else:
         basetype = MaskedArray
     # Get the result and view it as a (subclass of) MaskedArray
-    result = umath.power(fa, fb).view(basetype)
+    result = np.where(m, fa, umath.power(fa, fb)).view(basetype)
+    result._update_from(a)
     # Find where we're in trouble w/ NaNs and Infs
     invalid = np.logical_not(np.isfinite(result.view(ndarray)))
-    # Retrieve some extra attributes if needed
-    if isinstance(result, MaskedArray):
-        result._update_from(a)
     # Add the initial mask
     if m is not nomask:
-        if np.isscalar(result):
+        if not (result.ndim):
             return masked
+        m |= invalid
         result._mask = m
     # Fix the invalid parts
     if invalid.any():
         if not result.ndim:
             return masked
-        result[invalid] = masked
+        elif result._mask is nomask:
+            result._mask = invalid
         result._data[invalid] = result.fill_value
     return result
 
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index f9f46e563..8a0ca08ac 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -825,6 +825,7 @@ class TestMaskedArrayArithmetic(TestCase):
             self.failUnless(result is output)
             self.failUnless(output[0] is masked)
 
+
 #------------------------------------------------------------------------------
 
 class TestMaskedArrayAttributes(TestCase):
@@ -922,8 +923,6 @@ class TestMaskedArrayAttributes(TestCase):
         a[1] = 1
         assert_equal(a._mask, zeros(10))
 
-    def _wtv(self):
-        int(np.nan)
 
 #------------------------------------------------------------------------------
 
@@ -1242,22 +1241,131 @@ class TestMaskedArrayInPlaceArithmetics(TestCase):
 
     def test_inplace_division_misc(self):
         #
-        x = np.array([1.,1.,1.,-2., pi/2.0, 4., 5., -10., 10., 1., 2., 3.])
-        y = np.array([5.,0.,3., 2., -1., -4., 0., -10., 10., 1., 0., 3.])
-        m1 = [1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
-        m2 = [0, 0, 1, 0, 0, 1, 1, 0, 0, 0 ,0, 1]
+        x = [1., 1., 1.,-2., pi/2.,  4., 5., -10., 10., 1., 2., 3.]
+        y = [5., 0., 3., 2.,   -1., -4., 0., -10., 10., 1., 0., 3.]
+        m1 = [1,  0,  0,  0,     0,   0, 1,     0,   0,  0,  0, 0]
+        m2 = [0,  0,  1,  0,     0,   1, 1,     0,   0,  0 , 0, 1]
         xm = masked_array(x, mask=m1)
         ym = masked_array(y, mask=m2)
         #
         z = xm/ym
         assert_equal(z._mask, [1,1,1,0,0,1,1,0,0,0,1,1])
-        assert_equal(z._data, [0.2,1.,1./3.,-1.,-pi/2.,-1.,5.,1.,1.,1.,2.,1.])
+        assert_equal(z._data, [1.,1.,1.,-1.,-pi/2.,4.,5.,1.,1.,1.,2.,3.])
+        #assert_equal(z._data, [0.2,1.,1./3.,-1.,-pi/2.,-1.,5.,1.,1.,1.,2.,1.])
         #
         xm = xm.copy()
         xm /= ym
         assert_equal(xm._mask, [1,1,1,0,0,1,1,0,0,0,1,1])
-        assert_equal(xm._data, [1/5.,1.,1./3.,-1.,-pi/2.,-1.,5.,1.,1.,1.,2.,1.])
-
+        assert_equal(z._data, [1.,1.,1.,-1.,-pi/2.,4.,5.,1.,1.,1.,2.,3.])
+        #assert_equal(xm._data, [1/5.,1.,1./3.,-1.,-pi/2.,-1.,5.,1.,1.,1.,2.,1.])
+
+
+    def test_datafriendly_add(self):
+        "Test keeping data w/ (inplace) addition"
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        # Test add w/ scalar
+        xx = x + 1
+        assert_equal(xx.data, [2, 3, 3])
+        assert_equal(xx.mask, [0, 0, 1])
+        # Test iadd w/ scalar
+        x += 1
+        assert_equal(x.data, [2, 3, 3])
+        assert_equal(x.mask, [0, 0, 1])
+        # Test add w/ array
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        xx = x + array([1, 2, 3], mask=[1, 0, 0])
+        assert_equal(xx.data, [1, 4, 3])
+        assert_equal(xx.mask, [1, 0, 1])
+        # Test iadd w/ array
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        x += array([1, 2, 3], mask=[1, 0, 0])
+        assert_equal(x.data, [1, 4, 3])
+        assert_equal(x.mask, [1, 0, 1])
+
+
+    def test_datafriendly_sub(self):
+        "Test keeping data w/ (inplace) subtraction"
+        # Test sub w/ scalar
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        xx = x - 1
+        assert_equal(xx.data, [0, 1, 3])
+        assert_equal(xx.mask, [0, 0, 1])
+        # Test isub w/ scalar
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        x -= 1
+        assert_equal(x.data, [0, 1, 3])
+        assert_equal(x.mask, [0, 0, 1])
+        # Test sub w/ array
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        xx = x - array([1, 2, 3], mask=[1, 0, 0])
+        assert_equal(xx.data, [1, 0, 3])
+        assert_equal(xx.mask, [1, 0, 1])
+        # Test isub w/ array
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        x -= array([1, 2, 3], mask=[1, 0, 0])
+        assert_equal(x.data, [1, 0, 3])
+        assert_equal(x.mask, [1, 0, 1])
+
+
+    def test_datafriendly_mul(self):
+        "Test keeping data w/ (inplace) multiplication"
+        # Test mul w/ scalar
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        xx = x * 2
+        assert_equal(xx.data, [2, 4, 3])
+        assert_equal(xx.mask, [0, 0, 1])
+        # Test imul w/ scalar
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        x *= 2
+        assert_equal(x.data, [2, 4, 3])
+        assert_equal(x.mask, [0, 0, 1])
+        # Test mul w/ array
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        xx = x * array([10, 20, 30], mask=[1, 0, 0])
+        assert_equal(xx.data, [1, 40, 3])
+        assert_equal(xx.mask, [1, 0, 1])
+        # Test imul w/ array
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        x *= array([10, 20, 30], mask=[1, 0, 0])
+        assert_equal(x.data, [1, 40, 3])
+        assert_equal(x.mask, [1, 0, 1])
+
+
+    def test_datafriendly_div(self):
+        "Test keeping data w/ (inplace) division"
+        # Test div on scalar
+        x = array([1, 2, 3], mask=[0, 0, 1])
+        xx = x / 2.
+        assert_equal(xx.data, [1/2., 2/2., 3])
+        assert_equal(xx.mask, [0, 0, 1])
+        # Test idiv on scalar
+        x = array([1., 2., 3.], mask=[0, 0, 1])
+        x /= 2.
+        assert_equal(x.data, [1/2., 2/2., 3])
+        assert_equal(x.mask, [0, 0, 1])
+        # Test div on array
+        x = array([1., 2., 3.], mask=[0, 0, 1])
+        xx = x / array([10., 20., 30.], mask=[1, 0, 0])
+        assert_equal(xx.data, [1., 2./20., 3.])
+        assert_equal(xx.mask, [1, 0, 1])
+        # Test idiv on array
+        x = array([1., 2., 3.], mask=[0, 0, 1])
+        x /= array([10., 20., 30.], mask=[1, 0, 0])
+        assert_equal(x.data, [1., 2/20., 3.])
+        assert_equal(x.mask, [1, 0, 1])
+
+
+    def test_datafriendly_pow(self):
+        "Test keeping data w/ (inplace) power"
+        # Test pow on scalar
+        x = array([1., 2., 3.], mask=[0, 0, 1])
+        xx = x ** 2.5
+        assert_equal(xx.data, [1., 2.**2.5, 3.])
+        assert_equal(xx.mask, [0, 0, 1])
+        # Test ipow on scalar
+        x **= 2.5
+        assert_equal(x.data, [1., 2.**2.5, 3])
+        assert_equal(x.mask, [0, 0, 1])
 
 #------------------------------------------------------------------------------
 
@@ -2155,8 +2263,8 @@ class TestMaskedArrayFunctions(TestCase):
 
     def test_power(self):
         x = -1.1
-        assert_almost_equal(power(x,2.), 1.21)
-        self.failUnless(power(x,masked) is masked)
+        assert_almost_equal(power(x, 2.), 1.21)
+        self.failUnless(power(x, masked) is masked)
         x = array([-1.1,-1.1,1.1,1.1,0.])
         b = array([0.5,2.,0.5,2.,-1.], mask=[0,0,0,0,1])
         y = power(x,b)
diff --git a/numpy/ma/testutils.py b/numpy/ma/testutils.py
index 28754bccc..df7484f0a 100644
--- a/numpy/ma/testutils.py
+++ b/numpy/ma/testutils.py
@@ -167,22 +167,24 @@ def assert_array_compare(comparison, x, y, err_msg='', verbose=True, header='',
     """Asserts that a comparison relation between two masked arrays is satisfied
     elementwise."""
     # Fill the data first
-    xf = filled(x)
-    yf = filled(y)
+#    xf = filled(x)
+#    yf = filled(y)
     # Allocate a common mask and refill
     m = mask_or(getmask(x), getmask(y))
-    x = masked_array(xf, copy=False, mask=m)
-    y = masked_array(yf, copy=False, mask=m)
+    x = masked_array(x, copy=False, mask=m, subok=False)
+    y = masked_array(y, copy=False, mask=m, subok=False)
     if ((x is masked) and not (y is masked)) or \
         ((y is masked) and not (x is masked)):
         msg = build_err_msg([x, y], err_msg=err_msg, verbose=verbose,
                             header=header, names=('x', 'y'))
         raise ValueError(msg)
     # OK, now run the basic tests on filled versions
+    comparison = getattr(np, comparison.__name__, lambda x,y: True)
     return utils.assert_array_compare(comparison,
-                                x.filled(fill_value), y.filled(fill_value),
-                                err_msg=err_msg,
-                                verbose=verbose, header=header)
+                                      x.filled(fill_value),
+                                      y.filled(fill_value),
+                                      err_msg=err_msg,
+                                      verbose=verbose, header=header)
 
 
 def assert_array_equal(x, y, err_msg='', verbose=True):