From f627f362bf4e396d4f5bf8d159e875ffc8012896 Mon Sep 17 00:00:00 2001
From: Marten van Kerkwijk <mhvk@astro.utoronto.ca>
Date: Tue, 27 Feb 2024 17:19:00 -0500
Subject: [PATCH] BUG: fix support for gufuncs with more than one core
 dimension

In particular, `erfa.rxp` and `erfa.rxr`.
---
 astropy/utils/masked/core.py                 | 14 +++++++--
 astropy/utils/masked/tests/test_functions.py | 30 ++++++++++++++++++++
 docs/changes/utils/16120.bugfix.rst          |  2 ++
 3 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 docs/changes/utils/16120.bugfix.rst

diff --git a/astropy/utils/masked/core.py b/astropy/utils/masked/core.py
index 6bf669a9760..ca9d76254cf 100644
--- a/astropy/utils/masked/core.py
+++ b/astropy/utils/masked/core.py
@@ -773,7 +773,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                     ) = np.lib._function_base_impl._parse_gufunc_signature(
                         ufunc.signature.replace(" ", "")
                     )
-                axis = kwargs.get("axis", -1)
+                axis = kwargs.get("axis")
                 keepdims = kwargs.get("keepdims", False)
                 in_masks = []
                 for sig, mask in zip(in_sig, masks):
@@ -783,8 +783,13 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                             # value in those is masked, the output will be
                             # masked too (TODO: for multiple core dimensions
                             # this may be too strong).
+                            in_axis = (
+                                tuple(range(-1, -1 - len(sig), -1))
+                                if axis is None
+                                else axis
+                            )
                             mask = np.logical_or.reduce(
-                                mask, axis=axis, keepdims=keepdims
+                                mask, axis=in_axis, keepdims=keepdims
                             )
                         in_masks.append(mask)
 
@@ -794,7 +799,10 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                     if os:
                         # Output has core dimensions.  Assume all those
                         # get the same mask.
-                        result_mask = np.expand_dims(mask, axis)
+                        out_axis = (
+                            tuple(range(-1, -1 - len(os), -1)) if axis is None else axis
+                        )
+                        result_mask = np.expand_dims(mask, out_axis)
                     else:
                         result_mask = mask
                     result_masks.append(result_mask)
diff --git a/astropy/utils/masked/tests/test_functions.py b/astropy/utils/masked/tests/test_functions.py
index 9ace08958f6..828595e1fe4 100644
--- a/astropy/utils/masked/tests/test_functions.py
+++ b/astropy/utils/masked/tests/test_functions.py
@@ -173,6 +173,36 @@ def test_multi_op_ufunc(self):
             assert res0.unmasked == exp[0]
             assert res0.mask == expected_mask[0]
 
+    def test_erfa_rxp(self):
+        # Regression tests for gh-16116
+        m = Masked(np.eye(3))
+        v = Masked(np.arange(6).reshape(2, 3))
+        rxp1 = erfa_ufunc.rxp(m, v)
+        exp = erfa_ufunc.rxp(m.unmasked, v.unmasked)
+        assert_array_equal(rxp1.unmasked, exp)
+        assert_array_equal(rxp1.mask, False)
+        v.mask[0, 0] = True
+        rxp2 = erfa_ufunc.rxp(m, v)
+        assert_array_equal(rxp2.unmasked, exp)
+        assert_array_equal(rxp2.mask, [[True] * 3, [False] * 3])
+        m.mask[1, 1] = True
+        v.mask[...] = False
+        rxp3 = erfa_ufunc.rxp(m, v)
+        assert_array_equal(rxp3.unmasked, exp)
+        assert_array_equal(rxp3.mask, True)
+
+    def test_erfa_rxr(self):
+        m1 = Masked(np.arange(27.0).reshape(3, 3, 3))
+        m2 = Masked(np.arange(-27.0, 0.0).reshape(3, 3, 3))
+        rxr1 = erfa_ufunc.rxr(m1, m2)
+        exp = erfa_ufunc.rxr(m1.unmasked, m2.unmasked)
+        assert_array_equal(rxr1.unmasked, exp)
+        assert_array_equal(rxr1.mask, False)
+        m1.mask[0, 1, 2] = True
+        rxr2 = erfa_ufunc.rxr(m1, m2)
+        assert_array_equal(rxr2.unmasked, exp)
+        assert np.all(rxr2.mask == [[[True]], [[False]], [[False]]])
+
     @pytest.mark.parametrize("axis", (0, 1, None))
     def test_add_reduce(self, axis):
         ma_reduce = np.add.reduce(self.ma, axis=axis)
diff --git a/docs/changes/utils/16120.bugfix.rst b/docs/changes/utils/16120.bugfix.rst
new file mode 100644
index 00000000000..0f22acb2000
--- /dev/null
+++ b/docs/changes/utils/16120.bugfix.rst
@@ -0,0 +1,2 @@
+Fix support in ``Masked`` for generalized ufuncs with more than a
+single core dimension (such as ``erfa.rxp``).