BUG: MaskedNDArray.__array_ufunc__ breaks for erfa.rxp

I stumbled upon this while working on #16070. This looks like a bug in MaskedNDArray.__array_ufunc__, but maybe it's really a bug in pyerfa ?
I'm opening this report as soon as I managed to isolate the problem, but I suspect the issue is more general than it looks at first glance.

import erfa
import numpy as np
from astropy.utils.masked import Masked
import pytest


@pytest.mark.parametrize(
    "m, v",
    [
        (np.eye(3), np.arange(3)), # this works
        (Masked(np.eye(3)), Masked(np.arange(3))), # this doesn't
    ],
)
def test_reg(m, v):
    erfa.rxp(m, v)

relevant pytest output

___________________________________________ test_reg[m1-v1] ___________________________________________

m = MaskedNDArray([[1., 0., 0.],
               [0., 1., 0.],
               [0., 0., 1.]])
v = MaskedNDArray([0, 1, 2])

    @pytest.mark.parametrize(
        "m, v",
        [
            (np.eye(3), np.arange(3)),
            (Masked(np.eye(3)), Masked(np.arange(3))),
        ],
    )
    def test_reg(m, v):
>       erfa.rxp(m, v)

g.py:15: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/Users/clm/.pyenv/versions/astropy.dev/lib/python3.12/site-packages/erfa/core.py:19906: in rxp
    rp = ufunc.rxp(r, p)
../../astropy/utils/masked/core.py:872: in __array_ufunc__
    return self._masked_result(result, mask, out)
../../astropy/utils/masked/core.py:952: in _masked_result
    return Masked(result, mask)
../../astropy/utils/masked/core.py:77: in __new__
    return cls._get_masked_instance(*args, **kwargs)
../../astropy/utils/masked/core.py:125: in _get_masked_instance
    return masked_cls.from_unmasked(data, mask, copy)
../../astropy/utils/masked/core.py:518: in from_unmasked
    self._set_mask(mask, copy=copy)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = MaskedNDArray([0., 1., 2.]), mask = array([[False],
       [False],
       [False]])
copy = False

    def _set_mask(self, mask, copy=False):
        self_dtype = getattr(self, "dtype", None)
        mask_dtype = (
            np.ma.make_mask_descr(self_dtype)
            if self_dtype and self_dtype.names
            else np.dtype("?")
        )
        ma = np.asanyarray(mask, dtype=mask_dtype)
        if ma.shape != self.shape:
            # This will fail (correctly) if not broadcastable.
            self._mask = np.empty(self.shape, dtype=mask_dtype)
>           self._mask[...] = ma
E           ValueError: could not broadcast input array from shape (3,1) into shape (3,)

../../astropy/utils/masked/core.py:238: ValueError
======================================= short test summary info =======================================
FAILED g.py::test_reg[m1-v1] - ValueError: could not broadcast input array from shape (3,1) into shape (3,)

@mhvk

I think the problematic bits came from #11127. Specifically, I think this line is important here

astropy/astropy/utils/masked/core.py

Line 797 in 36d88ec

result_mask = np.expand_dims(mask, axis)

Can I ask @mhvk for confirmation that what I'm seeing is indeed a bug ?

Definitely a bug!

@mhvk

Based the following comment (from astropy/utils/masked/tests/test_containers.py::TestRepresentations::test_transformation)

astropy/astropy/utils/masked/tests/test_containers.py

Lines 49 to 50 in 2284523

    
           # Transformation back.  This also tests erfa.ufunc.s2p, which 
        
           # is special in having a core dimension only in the output.

I tried the following patch

diff --git a/astropy/utils/masked/core.py b/astropy/utils/masked/core.py
index 6bf669a976..46ef27c8c7 100644
--- a/astropy/utils/masked/core.py
+++ b/astropy/utils/masked/core.py
@@ -791,7 +791,7 @@ class MaskedNDArray(Masked, np.ndarray, base_cls=np.ndarray, data_cls=np.ndarray
                 mask = self._combine_masks(in_masks)
                 result_masks = []
                 for os in out_sig:
-                    if os:
+                    if os and not any(in_sig):
                         # Output has core dimensions.  Assume all those
                         # get the same mask.
                         result_mask = np.expand_dims(mask, axis)

which passes the test I posted, but breaks other tests. What seems to be key here are signatures of generalised universal functions (gufuncs) defined in pyerfa. Namely, here are the three functions that I'm struggling to all support at once

erfa.rxp: '(3, 3),(3)->(3)'
erfa.s2c: '(),()->(3)'
erfa.s2p: '(),(),()->(3)'
what they have in common is that they all have core dimensions in the output ((3))).

For what it's worth, I had a look at how pyerfa is generated and it seems very unlikely to me that it's where the bug lives.

So I'm 90% sure the problem is in MaskedNDArray.__array_ufunc__, but I can't seem to find a solution that accommodates every possibilities.
@mhvk, given your involvement with Masked, numpy and (py)erfa, I really can't resist asking you for a hand at this point 😅 Do you have any idea what to do or what to try ?

the tricky part is that the gufunc machinery assumes that all inputs have the same number of elements. I think it may be better for me to try to fix it, since I wrote both the Masked code and the gufunc axes/axis stuff for numpy...

See #16120 for a fix!

Navigation Menu

Repository navigation

BUG: `MaskedNDArray.__array_ufunc__` breaks for `erfa.rxp` #16116

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	# Transformation back. This also tests erfa.ufunc.s2p, which
	# is special in having a core dimension only in the output.

BUG: MaskedNDArray.__array_ufunc__ breaks for erfa.rxp #16116

Description

Activity

neutrinoceros commented on Feb 27, 2024

mhvk commented on Feb 27, 2024

neutrinoceros commented on Feb 27, 2024

mhvk commented on Feb 27, 2024

mhvk commented on Feb 27, 2024

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

BUG: `MaskedNDArray.__array_ufunc__` breaks for `erfa.rxp` #16116