Added astropy.units.quantity_overload decorator which allows instances of Quantity to be used with functions designed to accept only instances of ndarray. - #19024

Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

• Do the proposed changes actually accomplish desired goals?
• Do the proposed changes follow the Astropy coding guidelines?
• Are tests added/updated as required? If so, do they follow the Astropy testing guidelines?
• Are docs added/updated as required? If so, do they follow the Astropy documentation guidelines?
• Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see instructions for rebase and squash.
• Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the bot.
• Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
• Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
• At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) before merge.

I really didn't want to add the new numba dependency, but I haven't yet been able to find something which has the same behavior as numba.vectorize. I'll keep thinking about it, but lmk anyone has any ideas.

I think it may be reasonable to add it as a dev-only dependency using a dependency group, but in any case it should come with an explicit lower bound on the requirement.

Adding numba just to test one thing seems a bit of an overkill? I use my dev env for many things. I am not excited to have to install numba just to run this test.

You would only need to install it to run this test, not to run any other part of the test suite. Is this still too much ?
(I do agree it would be nicer not to require it at all)

We also have to think about all the flavors of CI jobs that would be required to pull it in. CI is red anyway.

@mhvk, do you understand why @numba.vectorize works, but @numpy.vectorize does not?

For example, if I do

@u.quantity_ufunc_overload
@np.vectorize
def add2(x: u.Quantity[u.mm], y: u.Quantity[u.mm]) -> u.Quantity[u.mm]:
    return x + y

add2(3 * u.mm, 2 * u.mm)

I get the error

TypeError: unknown ufunc add2 (vectorized).  If you believe this ufunc should be supported, please raise an issue on https://github.com/astropy/astropy

even though the function add2 appears to be a member of UFUNC_HELPERS,

u.quantity_helper.UFUNC_HELPERS[add2]

<function astropy.units.decorators.quantity_ufunc_overload.<locals>.decorator.<locals>.helper(f, *units)>

This obviously all works if I use @numba.vectorize, but it would be nice if we didn't need an additional testing dependency.

np.vectorize does not directly give a ufunc, unfortunately (it is stored somewhere on the class). For testing purposes, maybe better to use np.frompyfunc.

I just took a look at why CI is red and it may be a little confusing, but the short answer is we'll need to somehow isolate the test so it doesn't permanently affect the global state of the registry.

Indeed, I think you need some kind of setup/teardown combination that ensures QUANTITY_HELPERS is restored (there's also a monkeypatch pytest feature, I think).

p.s. Sorry, haven't had a chance to actually look at the code... (end of term time, so lots of exams to give...).

I can probably provide a custom pytest fixture to handle setup and teardown. Any way, just something to keep in mind.

@neutrinoceros and @mhvk, if you changed

assert all_q_ufuncs - all_np_ufuncs - all_erfa_ufuncs == set()

to

assert (all_np_ufuncs | all_erfa_ufuncs).issubset(all_q_ufuncs)

wouldn't that preserve the spirit of the test and allow the tests to pass?

No, definitely not OK. I would want to know if there are "wrong" ufuncs registered. But it is not hard to do the setup/teardown; let's first focus on the actual code (for me, not before the weekend, sorry!)

No worries, enjoy your weekend!

I realized that registering new ufuncs was probably unnecessary. The same behavior can be accomplished by writing a decorator which strips the units off of the arguments and then attaches the units after the numpy function is complete. This fixes the issue around needing setup/teardown methods for the UFUNC_HELPERS dict, as well as making this code more general since it will work for any function that accepts Numpy arrays, not just ufuncs.

@roytsmart - sorry still not had time to look properly... But the idea to make it more general is certainly nice (even if it effectively becomes a different beast). Indeed, with it I think you are going to close a couple of long-standing issues: #7368 and #6007 (latter, indirectly, by having a different way to do it),

@mhvk, that's great that this addresses some long-standing issues. I think this way is going to work much better.

Unfortunately, we won't be able to fix #6007 yet since @numpy.vectorize clobbers the function signature. However, there appears to be an open PR, numpy/numpy#29946 which may fix the issue as soon as it's merged.

Hi humans 👋 - this pull request hasn't had any new commits for approximately 5 months. I plan to close this in a month if the pull request doesn't have any new commits by then.

In lieu of a stalled pull request, please consider closing this and open an issue instead if a reminder is needed to revisit in the future. Maintainers may also choose to add keep-open label to keep this PR open but it is discouraged unless absolutely necessary.

If this PR still needs to be reviewed, as an author, you can rebase it to reset the clock.

If you believe I commented on this pull request incorrectly, please report this here.

Agreed, this would still be useful for me as well. I just rebased onto main so we can see if everything still works.

I think the devdeps failures are related.

units/tests/test_quantity_overload.py::test_grating_equation[...] - TypeError: Only one unit specification allowed, got [].

Hi humans 👋 - this pull request hasn't had any new commits for approximately 5 months. I plan to close this in a month if the pull request doesn't have any new commits by then.

In lieu of a stalled pull request, please consider closing this and open an issue instead if a reminder is needed to revisit in the future.

If this PR still needs to be reviewed, as an author, you can rebase it to reset the clock.

If you believe I commented on this pull request incorrectly, please report this here.