Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow wrapping astropy.units.Quantity #9705

Closed
wants to merge 2 commits into from

Conversation

tien-vo
Copy link
Contributor

@tien-vo tien-vo commented Nov 3, 2024

Ticks off first item in #9704 and also astropy #14454 .

The conditions added catch astropy.units.Quantity type while banning numpy.matrix and numpy.ma.masked_array, as suggested in #525 , i.e.,

>>> _is_array_like = lambda data: isinstance(data, np.ndarray | np.generic)
>>> _is_nep18 = lambda data: hasattr(data, "__array_function__")
>>> _has_array_api = lambda data: hasattr(data, "__array_namespace__")
>>> _has_unit = lambda data: hasattr(data, "_unit")
>>> catch_astropy = (
...    lambda data: _is_array_like(data)
...    and (_is_nep18(data) or _has_array_api(data))
...    and _has_unit(data)
... )

>>> catch_non_numpy_array = lambda data: not _is_array_like(data) and (
...    _is_nep18(data) or _has_array_api(data)
... )

>>> narray = np.arange(100)
>>> matrix = np.matrix(np.identity(3))
>>> marray = np.ma.masked_array(narray)
>>> uarray = u.Quantity(narray)
>>> parray = pint.Quantity(narray)
>>> darray = da.array(narray)

>>> catch_astropy(narray)  # False
>>> catch_astropy(matrix)  # False
>>> catch_astropy(marray)  # False
>>> catch_astropy(uarray)  # True
>>> catch_astropy(parray)  # False
>>> catch_astropy(darray)  # False

>>> catch_non_numpy_array(narray)  # False
>>> catch_non_numpy_array(matrix)  # False
>>> catch_non_numpy_array(marray)  # False
>>> catch_non_numpy_array(uarray)  # False
>>> catch_non_numpy_array(parray)  # True
>>> catch_non_numpy_array(darray)  # True

@keewis The above can be put into tests. I'm just not sure where. xarray/tests/test_units.py seems busy with pint-related stuff.

Copy link

welcome bot commented Nov 3, 2024

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@keewis
Copy link
Collaborator

keewis commented Nov 3, 2024

Thanks for the PR, @tien-vo.

This is good a start, but I believe we should try to avoid special-casing astropy.units.Quantity as much as possible. Instead, what I had envisioned is passing through numpy.ndarray subclasses (except np.matrix and np.ma.MaskedArray), but not numpy.ndarray itself. The condition would then be something like

if isinstance(data, np.matrix):
    data = np.asarray(data)

# immediately return array-like types except `numpy.ndarray` and `numpy` scalars
# compare types with `is` instead of `isinstance` to allow `numpy.ndarray` subclasses
is_numpy = type(data) is np.ndarray or isinstance(data, np.generic)
if not is_numpy and (hasattr(data, "__array_function__") or hasattr(data, "__array_namespace__")):
    return cast("T_DuckArray", data)

and any tests can be added to xarray/tests/test_variable.py::TestAsCompatibleData. You could call it test_numpy_subclass, and similarly to test_unsupported_type it may be enough to create a very simple numpy.ndarray subclass and pass that to the constructor of xr.Variable.

@shoyer, do you have any opinions on allowing numpy.ndarray subclasses? This is very similar to what I tried to do back in #2956.

@shoyer
Copy link
Member

shoyer commented Nov 4, 2024 via email

@slevang
Copy link
Contributor

slevang commented Nov 5, 2024

Instead, what I had envisioned is passing through numpy.ndarray subclasses

+1 to this! Was just playing around with a very simple numpy subclass and was surprised to see that this is not currently supported.

@slevang
Copy link
Contributor

slevang commented Nov 10, 2024

I started #9760 just because I would love to get this feature added soon.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants