Merge pull request #6653 from charris/fix-ma-dot

Fix ma dot

4 files changed, 262 insertions, 184 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 7d9acbd1c..87082d139 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -32,9 +32,12 @@ import numpy.core.numerictypes as ntypes
 from numpy import ndarray, amax, amin, iscomplexobj, bool_, _NoValue
 from numpy import array as narray
 from numpy.lib.function_base import angle
-from numpy.compat import getargspec, formatargspec, long, basestring, unicode, bytes, sixu
+from numpy.compat import (
+    getargspec, formatargspec, long, basestring, unicode, bytes, sixu
+    )
 from numpy import expand_dims as n_expand_dims
 
+
 if sys.version_info[0] >= 3:
     import pickle
 else:
@@ -4651,24 +4654,44 @@ class MaskedArray(ndarray):
             return D.astype(dtype).filled(0).sum(axis=None, out=out)
     trace.__doc__ = ndarray.trace.__doc__
 
-    def dot(self, other, out=None):
-        am = ~getmaskarray(self)
-        bm = ~getmaskarray(other)
-        if out is None:
-            d = np.dot(filled(self, 0), filled(other, 0))
-            m = ~np.dot(am, bm)
-            if d.ndim == 0:
-                d = np.asarray(d)
-            r = d.view(get_masked_subclass(self, other))
-            r.__setmask__(m)
-            return r
-        d = self.filled(0).dot(other.filled(0), out._data)
-        if out.mask.shape != d.shape:
-            out._mask = np.empty(d.shape, MaskType)
-        np.dot(am, bm, out._mask)
-        np.logical_not(out._mask, out._mask)
-        return out
-    dot.__doc__ = ndarray.dot.__doc__
+    def dot(self, b, out=None, strict=False):
+        """
+        a.dot(b, out=None)
+
+        Masked dot product of two arrays. Note that `out` and `strict` are
+        located in different positions than in `ma.dot`. In order to
+        maintain compatibility with the functional version, it is
+        recommended that the optional arguments be treated as keyword only.
+        At some point that may be mandatory.
+
+        .. versionadded:: 1.10.0
+
+        Parameters
+        ----------
+        b : masked_array_like
+            Inputs array.
+        out : masked_array, optional
+            Output argument. This must have the exact kind that would be
+            returned if it was not used. In particular, it must have the
+            right type, must be C-contiguous, and its dtype must be the
+            dtype that would be returned for `ma.dot(a,b)`. This is a
+            performance feature. Therefore, if these conditions are not
+            met, an exception is raised, instead of attempting to be
+            flexible.
+        strict : bool, optional
+            Whether masked data are propagated (True) or set to 0 (False)
+            for the computation. Default is False.  Propagating the mask
+            means that if a masked value appears in a row or column, the
+            whole row or column is considered masked.
+
+            .. versionadded:: 1.10.2
+
+        See Also
+        --------
+        numpy.ma.dot : equivalent function
+    
+        """
+        return dot(self, b, out=out, strict=strict)
 
     def sum(self, axis=None, dtype=None, out=None):
         """
@@ -5884,7 +5907,7 @@ class mvoid(MaskedArray):
         --------
         MaskedArray.filled
 
-    """
+        """
         return asarray(self).filled(fill_value)[()]
 
     def tolist(self):
@@ -7021,6 +7044,186 @@ def round_(a, decimals=0, out=None):
 round = round_
 
 
+# Needed by dot, so move here from extras.py. It will still be exported
+# from extras.py for compatibility.
+def mask_rowcols(a, axis=None):
+    """
+    Mask rows and/or columns of a 2D array that contain masked values.
+
+    Mask whole rows and/or columns of a 2D array that contain
+    masked values.  The masking behavior is selected using the
+    `axis` parameter.
+
+      - If `axis` is None, rows *and* columns are masked.
+      - If `axis` is 0, only rows are masked.
+      - If `axis` is 1 or -1, only columns are masked.
+
+    Parameters
+    ----------
+    a : array_like, MaskedArray
+        The array to mask.  If not a MaskedArray instance (or if no array
+        elements are masked).  The result is a MaskedArray with `mask` set
+        to `nomask` (False). Must be a 2D array.
+    axis : int, optional
+        Axis along which to perform the operation. If None, applies to a
+        flattened version of the array.
+
+    Returns
+    -------
+    a : MaskedArray
+        A modified version of the input array, masked depending on the value
+        of the `axis` parameter.
+
+    Raises
+    ------
+    NotImplementedError
+        If input array `a` is not 2D.
+
+    See Also
+    --------
+    mask_rows : Mask rows of a 2D array that contain masked values.
+    mask_cols : Mask cols of a 2D array that contain masked values.
+    masked_where : Mask where a condition is met.
+
+    Notes
+    -----
+    The input array's mask is modified by this function.
+
+    Examples
+    --------
+    >>> import numpy.ma as ma
+    >>> a = np.zeros((3, 3), dtype=np.int)
+    >>> a[1, 1] = 1
+    >>> a
+    array([[0, 0, 0],
+           [0, 1, 0],
+           [0, 0, 0]])
+    >>> a = ma.masked_equal(a, 1)
+    >>> a
+    masked_array(data =
+     [[0 0 0]
+     [0 -- 0]
+     [0 0 0]],
+          mask =
+     [[False False False]
+     [False  True False]
+     [False False False]],
+          fill_value=999999)
+    >>> ma.mask_rowcols(a)
+    masked_array(data =
+     [[0 -- 0]
+     [-- -- --]
+     [0 -- 0]],
+          mask =
+     [[False  True False]
+     [ True  True  True]
+     [False  True False]],
+          fill_value=999999)
+
+    """
+    a = array(a, subok=False)
+    if a.ndim != 2:
+        raise NotImplementedError("mask_rowcols works for 2D arrays only.")
+    m = getmask(a)
+    # Nothing is masked: return a
+    if m is nomask or not m.any():
+        return a
+    maskedval = m.nonzero()
+    a._mask = a._mask.copy()
+    if not axis:
+        a[np.unique(maskedval[0])] = masked
+    if axis in [None, 1, -1]:
+        a[:, np.unique(maskedval[1])] = masked
+    return a
+
+
+# Include masked dot here to avoid import problems in getting it from
+# extras.py. Note that it is not included in __all__, but rather exported
+# from extras in order to avoid backward compatibility problems.
+def dot(a, b, strict=False, out=None):
+    """
+    Return the dot product of two arrays.
+
+    This function is the equivalent of `numpy.dot` that takes masked values
+    into account. Note that `strict` and `out` are in different position
+    than in the method version. In order to maintain compatibility with the
+    corresponding method, it is recommended that the optional arguments be
+    treated as keyword only.  At some point that may be mandatory.
+
+    .. note::
+      Works only with 2-D arrays at the moment.
+
+
+    Parameters
+    ----------
+    a, b : masked_array_like
+        Inputs arrays.
+    strict : bool, optional
+        Whether masked data are propagated (True) or set to 0 (False) for
+        the computation. Default is False.  Propagating the mask means that
+        if a masked value appears in a row or column, the whole row or
+        column is considered masked.
+    out : masked_array, optional
+        Output argument. This must have the exact kind that would be returned
+        if it was not used. In particular, it must have the right type, must be
+        C-contiguous, and its dtype must be the dtype that would be returned
+        for `dot(a,b)`. This is a performance feature. Therefore, if these
+        conditions are not met, an exception is raised, instead of attempting
+        to be flexible.
+
+        .. versionadded:: 1.10.2
+
+    See Also
+    --------
+    numpy.dot : Equivalent function for ndarrays.
+
+    Examples
+    --------
+    >>> a = ma.array([[1, 2, 3], [4, 5, 6]], mask=[[1, 0, 0], [0, 0, 0]])
+    >>> b = ma.array([[1, 2], [3, 4], [5, 6]], mask=[[1, 0], [0, 0], [0, 0]])
+    >>> np.ma.dot(a, b)
+    masked_array(data =
+     [[21 26]
+     [45 64]],
+                 mask =
+     [[False False]
+     [False False]],
+           fill_value = 999999)
+    >>> np.ma.dot(a, b, strict=True)
+    masked_array(data =
+     [[-- --]
+     [-- 64]],
+                 mask =
+     [[ True  True]
+     [ True False]],
+           fill_value = 999999)
+
+    """
+    # !!!: Works only with 2D arrays. There should be a way to get it to run
+    # with higher dimension
+    if strict and (a.ndim == 2) and (b.ndim == 2):
+        a = mask_rowcols(a, 0)
+        b = mask_rowcols(b, 1)
+    am = ~getmaskarray(a)
+    bm = ~getmaskarray(b)
+
+    if out is None:
+        d = np.dot(filled(a, 0), filled(b, 0))
+        m = ~np.dot(am, bm)
+        if d.ndim == 0:
+            d = np.asarray(d)
+        r = d.view(get_masked_subclass(a, b))
+        r.__setmask__(m)
+        return r
+    else:
+        d = np.dot(filled(a, 0), filled(b, 0), out._data)
+        if out.mask.shape != d.shape:
+            out._mask = np.empty(d.shape, MaskType)
+        np.dot(am, bm, out._mask)
+        np.logical_not(out._mask, out._mask)
+        return out
+
+
 def inner(a, b):
     """
     Returns the inner product of a and b for arrays of floating point types.
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index ae4e0cee5..e1d228e73 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -29,7 +29,8 @@ from . import core as ma
 from .core import (
     MaskedArray, MAError, add, array, asarray, concatenate, filled,
     getmask, getmaskarray, make_mask_descr, masked, masked_array, mask_or,
-    nomask, ones, sort, zeros, getdata
+    nomask, ones, sort, zeros, getdata, get_masked_subclass, dot,
+    mask_rowcols
     )
 
 import numpy as np
@@ -846,96 +847,6 @@ def compress_cols(a):
         raise NotImplementedError("compress_cols works for 2D arrays only.")
     return compress_rowcols(a, 1)
 
-def mask_rowcols(a, axis=None):
-    """
-    Mask rows and/or columns of a 2D array that contain masked values.
-
-    Mask whole rows and/or columns of a 2D array that contain
-    masked values.  The masking behavior is selected using the
-    `axis` parameter.
-
-      - If `axis` is None, rows *and* columns are masked.
-      - If `axis` is 0, only rows are masked.
-      - If `axis` is 1 or -1, only columns are masked.
-
-    Parameters
-    ----------
-    a : array_like, MaskedArray
-        The array to mask.  If not a MaskedArray instance (or if no array
-        elements are masked).  The result is a MaskedArray with `mask` set
-        to `nomask` (False). Must be a 2D array.
-    axis : int, optional
-        Axis along which to perform the operation. If None, applies to a
-        flattened version of the array.
-
-    Returns
-    -------
-    a : MaskedArray
-        A modified version of the input array, masked depending on the value
-        of the `axis` parameter.
-
-    Raises
-    ------
-    NotImplementedError
-        If input array `a` is not 2D.
-
-    See Also
-    --------
-    mask_rows : Mask rows of a 2D array that contain masked values.
-    mask_cols : Mask cols of a 2D array that contain masked values.
-    masked_where : Mask where a condition is met.
-
-    Notes
-    -----
-    The input array's mask is modified by this function.
-
-    Examples
-    --------
-    >>> import numpy.ma as ma
-    >>> a = np.zeros((3, 3), dtype=np.int)
-    >>> a[1, 1] = 1
-    >>> a
-    array([[0, 0, 0],
-           [0, 1, 0],
-           [0, 0, 0]])
-    >>> a = ma.masked_equal(a, 1)
-    >>> a
-    masked_array(data =
-     [[0 0 0]
-     [0 -- 0]
-     [0 0 0]],
-          mask =
-     [[False False False]
-     [False  True False]
-     [False False False]],
-          fill_value=999999)
-    >>> ma.mask_rowcols(a)
-    masked_array(data =
-     [[0 -- 0]
-     [-- -- --]
-     [0 -- 0]],
-          mask =
-     [[False  True False]
-     [ True  True  True]
-     [False  True False]],
-          fill_value=999999)
-
-    """
-    a = array(a, subok=False)
-    if a.ndim != 2:
-        raise NotImplementedError("mask_rowcols works for 2D arrays only.")
-    m = getmask(a)
-    # Nothing is masked: return a
-    if m is nomask or not m.any():
-        return a
-    maskedval = m.nonzero()
-    a._mask = a._mask.copy()
-    if not axis:
-        a[np.unique(maskedval[0])] = masked
-    if axis in [None, 1, -1]:
-        a[:, np.unique(maskedval[1])] = masked
-    return a
-
 def mask_rows(a, axis=None):
     """
     Mask rows of a 2D array that contain masked values.
@@ -1027,58 +938,6 @@ def mask_cols(a, axis=None):
     return mask_rowcols(a, 1)
 
 
-def dot(a, b, strict=False):
-    """
-    Return the dot product of two arrays.
-
-    .. note::
-      Works only with 2-D arrays at the moment.
-
-    This function is the equivalent of `numpy.dot` that takes masked values
-    into account, see `numpy.dot` for details.
-
-    Parameters
-    ----------
-    a, b : ndarray
-        Inputs arrays.
-    strict : bool, optional
-        Whether masked data are propagated (True) or set to 0 (False) for the
-        computation. Default is False.
-        Propagating the mask means that if a masked value appears in a row or
-        column, the whole row or column is considered masked.
-
-    See Also
-    --------
-    numpy.dot : Equivalent function for ndarrays.
-
-    Examples
-    --------
-    >>> a = ma.array([[1, 2, 3], [4, 5, 6]], mask=[[1, 0, 0], [0, 0, 0]])
-    >>> b = ma.array([[1, 2], [3, 4], [5, 6]], mask=[[1, 0], [0, 0], [0, 0]])
-    >>> np.ma.dot(a, b)
-    masked_array(data =
-     [[21 26]
-     [45 64]],
-                 mask =
-     [[False False]
-     [False False]],
-           fill_value = 999999)
-    >>> np.ma.dot(a, b, strict=True)
-    masked_array(data =
-     [[-- --]
-     [-- 64]],
-                 mask =
-     [[ True  True]
-     [ True False]],
-           fill_value = 999999)
-
-    """
-    #!!!: Works only with 2D arrays. There should be a way to get it to run with higher dimension
-    if strict and (a.ndim == 2) and (b.ndim == 2):
-        a = mask_rows(a)
-        b = mask_cols(b)
-    return a.dot(b)
-
 #####--------------------------------------------------------------------------
 #---- --- arraysetops ---
 #####--------------------------------------------------------------------------
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 70c1ee12c..61fd77bda 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -3192,7 +3192,7 @@ class TestMaskedArrayMathMethods(TestCase):
         assert_almost_equal(r.filled(0), fX.dot(fX))
         assert_(r.mask[1,3])
         r1 = empty_like(r)
-        mX.dot(mX, r1)
+        mX.dot(mX, out=r1)
         assert_almost_equal(r, r1)
 
         mYY = mXX.swapaxes(-1, -2)
@@ -3200,7 +3200,7 @@ class TestMaskedArrayMathMethods(TestCase):
         r = mXX.dot(mYY)
         assert_almost_equal(r.filled(0), fXX.dot(fYY))
         r1 = empty_like(r)
-        mXX.dot(mYY, r1)
+        mXX.dot(mYY, out=r1)
         assert_almost_equal(r, r1)
 
     def test_dot_shape_mismatch(self):
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index c41c629fc..6138d0573 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -538,26 +538,26 @@ class TestCompressFunctions(TestCase):
         m = [1, 0, 0, 0, 0, 0]
         a = masked_array(n, mask=m).reshape(2, 3)
         b = masked_array(n, mask=m).reshape(3, 2)
-        c = dot(a, b, True)
+        c = dot(a, b, strict=True)
         assert_equal(c.mask, [[1, 1], [1, 0]])
-        c = dot(b, a, True)
+        c = dot(b, a, strict=True)
         assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
-        c = dot(a, b, False)
+        c = dot(a, b, strict=False)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
-        c = dot(b, a, False)
+        c = dot(b, a, strict=False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
         m = [0, 0, 0, 0, 0, 1]
         a = masked_array(n, mask=m).reshape(2, 3)
         b = masked_array(n, mask=m).reshape(3, 2)
-        c = dot(a, b, True)
+        c = dot(a, b, strict=True)
         assert_equal(c.mask, [[0, 1], [1, 1]])
-        c = dot(b, a, True)
+        c = dot(b, a, strict=True)
         assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]])
-        c = dot(a, b, False)
+        c = dot(a, b, strict=False)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
         assert_equal(c, dot(a, b))
-        c = dot(b, a, False)
+        c = dot(b, a, strict=False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
         m = [0, 0, 0, 0, 0, 0]
@@ -570,37 +570,53 @@ class TestCompressFunctions(TestCase):
         #
         a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3)
         b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
-        c = dot(a, b, True)
+        c = dot(a, b, strict=True)
         assert_equal(c.mask, [[1, 1], [0, 0]])
-        c = dot(a, b, False)
+        c = dot(a, b, strict=False)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
-        c = dot(b, a, True)
+        c = dot(b, a, strict=True)
         assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
-        c = dot(b, a, False)
+        c = dot(b, a, strict=False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
         a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
         b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
-        c = dot(a, b, True)
+        c = dot(a, b, strict=True)
         assert_equal(c.mask, [[0, 0], [1, 1]])
         c = dot(a, b)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
-        c = dot(b, a, True)
+        c = dot(b, a, strict=True)
         assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]])
-        c = dot(b, a, False)
+        c = dot(b, a, strict=False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
         a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
         b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2)
-        c = dot(a, b, True)
+        c = dot(a, b, strict=True)
         assert_equal(c.mask, [[1, 0], [1, 1]])
-        c = dot(a, b, False)
+        c = dot(a, b, strict=False)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
-        c = dot(b, a, True)
+        c = dot(b, a, strict=True)
         assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]])
-        c = dot(b, a, False)
+        c = dot(b, a, strict=False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
 
+    def test_dot_returns_maskedarray(self):
+        # See gh-6611
+        a = np.eye(3)
+        b = array(a)
+        assert_(type(dot(a, a)) is MaskedArray)
+        assert_(type(dot(a, b)) is MaskedArray)
+        assert_(type(dot(b, a)) is MaskedArray)
+        assert_(type(dot(b, b)) is MaskedArray)
+
+    def test_dot_out(self):
+        a = array(np.eye(3))
+        out = array(np.zeros((3, 3)))
+        res = dot(a, a, out=out)
+        assert_(res is out)
+        assert_equal(a, res)
+
 
 class TestApplyAlongAxis(TestCase):
     # Tests 2D functions