Merge pull request numpy#21262 from ganesh-k13/kron_21257_ma

ENH: Masked Array support for `np.kron`

Author: seberg

doc/release/upcoming_changes/21262.improvement.rst

@@ -0,0 +1,34 @@
+``np.kron`` now maintains subclass information
+----------------------------------------------
+``np.kron`` maintains subclass information now such as masked arrays
+while computing the Kronecker product of the inputs
+
+.. code-block:: python
+
+    >>> x = ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]])
+    >>> np.kron(x,x)
+    masked_array(
+      data=[[1, --, --, --],
+            [--, 4, --, --],
+            [--, --, 4, --],
+            [--, --, --, 16]],
+      mask=[[False,  True,  True,  True],
+            [ True, False,  True,  True],
+            [ True,  True, False,  True],
+            [ True,  True,  True, False]],
+      fill_value=999999)
+
+.. warning::
+    ``np.kron`` output now follows ``ufunc`` ordering (``multiply``)
+    to determine the output class type
+
+    .. code-block:: python
+
+        >>> class myarr(np.ndarray):
+        >>>    __array_priority__ = -1
+        >>> a = np.ones([2, 2])
+        >>> ma = myarray(a.shape, a.dtype, a.data)
+        >>> type(np.kron(a, ma)) == np.ndarray
+        False # Before it was True
+        >>> type(np.kron(a, ma)) == myarr
+        True

numpy/lib/shape_base.py

@@ -1159,23 +1159,19 @@ def kron(a, b):
     bs = (1,)*max(0, nda-ndb) + bs
 
     # Compute the product
-    a_arr = _nx.asarray(a).reshape(a.size, 1)
-    b_arr = _nx.asarray(b).reshape(1, b.size)
-    result = a_arr * b_arr
+    a_arr = a.reshape(a.size, 1)
+    b_arr = b.reshape(1, b.size)
+    is_any_mat = isinstance(a_arr, matrix) or isinstance(b_arr, matrix)
+    # In case of `mat`, convert result to `array`
+    result = _nx.multiply(a_arr, b_arr, subok=(not is_any_mat))
 
     # Reshape back
     result = result.reshape(as_+bs)
     transposer = _nx.arange(nd*2).reshape([2, nd]).ravel(order='f')
     result = result.transpose(transposer)
     result = result.reshape(_nx.multiply(as_, bs))
 
-    wrapper = get_array_prepare(a, b)
-    if wrapper is not None:
-        result = wrapper(result)
-    wrapper = get_array_wrap(a, b)
-    if wrapper is not None:
-        result = wrapper(result)
-    return result
+    return result if not is_any_mat else matrix(result, copy=False)
 
 
 def _tile_dispatcher(A, reps):

numpy/lib/tests/test_shape_base.py

@@ -646,21 +646,31 @@ def test_basic(self):
 class TestKron:
     def test_return_type(self):
         class myarray(np.ndarray):
-            __array_priority__ = 0.0
+            __array_priority__ = 1.0
 
         a = np.ones([2, 2])
         ma = myarray(a.shape, a.dtype, a.data)
         assert_equal(type(kron(a, a)), np.ndarray)
         assert_equal(type(kron(ma, ma)), myarray)
-        assert_equal(type(kron(a, ma)), np.ndarray)
+        assert_equal(type(kron(a, ma)), myarray)
         assert_equal(type(kron(ma, a)), myarray)
 
-    def test_kron_smoke(self):
-        a = np.ones([3, 3])
-        b = np.ones([3, 3])
-        k = np.ones([9, 9])
+    @pytest.mark.parametrize(
+        "array_class", [np.asarray, np.mat]
+    )
+    def test_kron_smoke(self, array_class):
+        a = array_class(np.ones([3, 3]))
+        b = array_class(np.ones([3, 3]))
+        k = array_class(np.ones([9, 9]))
+
+        assert_array_equal(np.kron(a, b), k)
+
+    def test_kron_ma(self):
+        x = np.ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]])
+        k = np.ma.array(np.diag([1, 4, 4, 16]),
+                mask=~np.array(np.identity(4), dtype=bool))
 
-        assert np.array_equal(np.kron(a, b), k), "Smoke test for kron failed"
+        assert_array_equal(k, np.kron(x, x))
 
     @pytest.mark.parametrize(
         "shape_a,shape_b", [