Is there a way to register new ufuncs with astropy?

What is the problem this feature will solve?

I've been using the @numba.vectorize decorator to create new ufuncs, and I've been curious if there's a to register the new ufunc with astropy.units?

Describe the desired outcome

I think it would be nice if there was a decorator, similar to @quantity_input, which registers the ufunc, converts the inputs to the correct units, and then attaches the appropriate units to the output. This would allow our users to parallelize numerically-intensive functions without having to create an additional wrapper for every function.

Additional context

No response

Hum, I'm not sure exactly how general the problem is and wether it's possible at all to write a reusable decorator for this. Our implementation is already not super clean: see for instance how some ufuncs like np.reduce or np.initial are being special-cased in Quantity.__array_ufunc__.
Basically, I would absolutely get on board with a generalized implementation that could be exposed as a decorator, but I think I would need to see a proof of concept or be shown evidence that it's feasible at all first.

@mhvk

maybe @mhvk has more thoughts

Yes, it is possible, at least in principle. It does require that @numba.vectorize creates a proper numpy ufunc (or, more specifically, can be overridden with __array_ufunc__), and that the function that you create takes plain arrays (with assumed units, like np.sin takes arrays assumed to be in radian).

For instance, np.frompyfunc creates a proper ufunc, and the following works:

import numpy as np
import astropy.units as u
from astropy.units.quantity_helper.helpers import get_converter


def mypyfunc(a, b):
    """Like add, but prefer the unit of the second argument."""
    return a + b

myufunc = np.frompyfunc(mypyfunc, 2, 1)

def my_helper(f, unit1, unit2):
    # Converters for each unit, output unit
    return [get_converter(unit1, unit2), None], unit2

u.quantity_helper.UFUNC_HELPERS[myufunc] = my_helper

print(myufunc(1.*u.deg, 20*u.arcmin))
# 80.0 arcmin

(np.frompyfunc really just iterates over the inputs and calls the python function so it is slow; the above also returns a Quantity with dtype=object, so not so useful! But numba.vectorize should be better than that!)

Obviously, this is basically using private API, but have wondered about exposing this in some way; your case may just be the one that makes it happen... Maybe you can give a concrete test case?

@mhvk

@mhvk, thanks for the helpful example. To start, I've adapted your code to use @numba.vectorize, and I was able to confirm it does return an array with the correct dtype.

import numba
import astropy.units as u
from astropy.units.quantity_helper.helpers import get_converter

@numba.vectorize
def add_numba(a: float, b: float) -> float:
    """Numba-accelerated addition"""
    return a + b

def my_helper(f, unit1, unit2):
    # Converters for each unit, output unit
    return [get_converter(unit1, unit2), None], unit2

u.quantity_helper.UFUNC_HELPERS[add_numba] = my_helper

result = add_numba([1., 2] * u.deg, 20 * u.arcmin)
print(result, result.dtype)

[ 80. 140.] arcmin float64

Do you think it makes sense to expose a decorator that defines/registers the helper function, or would it be better to allow the user to define/register the helper directly?

Super that it works! How to proceed depends in part on whether one is thinking of one-offs, or maybe creating a mini package with a whole set of functions. E.g., we register packages as follows:

astropy/astropy/units/quantity_helper/converters.py

Lines 35 to 50 in b3e5d64

    
               def register_module(self, module, names, importer): 
        
                   """Register (but do not import) a set of ufunc helpers. 
        
                   Parameters 
        
                   ---------- 
        
                   module : str 
        
                       Name of the module with the ufuncs (e.g., 'scipy.special'). 
        
                   names : iterable of str 
        
                       Names of the module ufuncs for which helpers are available. 
        
                   importer : callable 
        
                       Function that imports the ufuncs and returns a dict of helpers 
        
                       keyed by those ufuncs.  If the value is `None`, the ufunc is 
        
                       explicitly *not* supported. 
        
                   """ 
        
                   with self._lock: 
        
                       self.modules[module] = {"names": names, "importer": importer}

This is used in scipy_special.py and erfa.py (former ensuring that things do not fail if scipy is not available).

So, I think it depends a bit on you too... One solution would be for now just document an example, and warn that it is private API which may change.

p.s. If we do introduce API, it would be nice if one could do something similar for Masked... Though there things may work more automatically, as no real helpers are needed.

In my particular use case I'm trying to accelerate a subset of my raytracing library, optika.

Defining a bunch of helpers and using register_module wouldn't be too bad, but in my opinion it would be cleaner if I could chain another decorator to @numba.vectorize.

I've written a naive implementation of this idea below, but I can't quite figure out how to get it to work for keyword arguments yet.

import numpy as np
import numba
import astropy.units as u
from astropy.units.quantity_helper.helpers import get_converter

def quantity_ufunc(
    units_out: None | u.UnitBase | tuple[u.UnitBase],
    *units_in: None | u.UnitBase,
):
    """A decorator which converts quantities to numpy arrays for use with ufuncs."""

    def decorator(func):

        def helper(f, *units):
            converters = [
                get_converter(unit, unit_in) if unit_in is not None else None
                for unit, unit_in in zip(units, units_in)
            ]
            return converters, units_out

        u.quantity_helper.UFUNC_HELPERS[func] = helper

        return func

    return decorator

@quantity_ufunc(u.radian, None, u.um, u.um)
@numba.vectorize
def grating_equation(
    m: float,
    wavelength: float,
    d: float,
) -> float:
    """
    Compute the diffraction angle for normally-incident light
    on a diffraction grating.

    Parameters
    ----------
    m
        The order of the diffracted light
    wavelength
        The wavelength of the incident light.
    d
        The spacing between the grooves of the grating.
    """
    return np.arcsin(m * wavelength / d)

m = 1
wavelength = [500, 600] * u.nm
d = 1 * u.um

result = grating_equation(m, wavelength, d)

expect = np.arcsin(m * wavelength / d)

print(f"{result=}")
print(f"{expect=}")

result=<Quantity [0.52359878, 0.64350111] rad>
expect=<Quantity [0.52359878, 0.64350111] rad>

I'm not sure how to extend this to masked quantities yet, but I do think that would be very useful.

@nstarman

Yes, a decorator does seem like the right solution here! I think we can perhaps take a bit of inspiration from quantity_input on how to use inspection. I could imagine something like the default being

u.register_quantity_ufunc(ufunc=None, helper=None, *args, **kwargs)

where ufunc by default is grabbed from whatever is decorated, helper can be given explicitly, but if it is not, units are gotten from *args or **kwargs (keyed by argument name?).

Alternatively, a bit more a la typing, could it be something more like overloading in typing,

@u.quantity_ufunc_overload
def grating_equation(m : Quantity[u.one], wavelength : Quantity[u.m], d : Quantity[u.m]) -> Quantity[u.rad]
@numba.vectorize
def grating_equation(
    m: float,
    wavelength: float,
    d: float,
) -> float:
   ...

I don't actually know if this same pattern is possible outside typing ... probably ... but this seems a particularly nice way to combine both "overloads".

As yet another alternative, perhaps one would just do the regular overload, and then expose a single routine that scans all functions in a module and registers the helpers for it based on the type annotations. Note, though, that we'd still need a bit more freedom, e.g., with the result unit being some function of the input units. I do not know to what extent that can be expressed using typing.

Let me ping @nstarman - he has thought quite a bit about this in the context of quantity_input, and may have suggestions about how to combine this with typing.

I like the idea of using type annotations to determine the units, similar to the behavior of @quantity_input.

Here is my attempt at implementing that idea, which also includes an equivalencies= feature.

from typing import Callable
import numpy as np
import numba
import astropy.units as u
from astropy.units.quantity_helper.helpers import get_converter
from astropy.units.decorators import _parse_annotation

def quantity_ufunc_overload(
    function: None | Callable = None,
    equivalencies: None | list = None,
):
    """A decorator which converts quantities to numpy arrays for use with ufuncs."""

    eq = equivalencies

    def decorator(func):

        signature = inspect.signature(func)

        units_out = _parse_annotation(signature.return_annotation)

        units_in = [
            _parse_annotation(param.annotation)
            for param in signature.parameters.values()
        ]

        units_in = [
            unit if unit else None
            for unit in units_in
        ]

        def helper(f, *units):
            converters = [
                unit.get_converter(unit_in, equivalencies=eq) if unit_in is not None
                else None
                for unit, unit_in in zip(units, units_in)
            ]
            return converters, units_out

        u.quantity_helper.UFUNC_HELPERS[func] = helper

        return func

    if function is not None:
        decorator = decorator(function)

    return decorator

@quantity_ufunc_overload(equivalencies=u.spectral())
@numba.vectorize
def grating_equation(
    m: float | np.ndarray,
    wavelength: u.Quantity[u.um],
    d: u.Quantity[u.um],
) -> u.Quantity[u.rad]:
    """
    Compute the diffraction angle for normally-incident light on a grating.
    """
    return np.arcsin(m * wavelength / d)

m = 1
wavelength = 2 * u.eV
d = 1 * u.um

result = grating_equation(m, wavelength, d)

print(f"{result=}")

result=<Quantity 0.66864201 rad>

I should have time next week to look at this in more detail. Making a reminder...

Navigation Menu

Repository navigation

Is there a way to register new ufuncs with astropy? #19005

What is the problem this feature will solve?

Describe the desired outcome

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	def register_module(self, module, names, importer):
	"""Register (but do not import) a set of ufunc helpers.

	Parameters
	----------
	module : str
	Name of the module with the ufuncs (e.g., 'scipy.special').
	names : iterable of str
	Names of the module ufuncs for which helpers are available.
	importer : callable
	Function that imports the ufuncs and returns a dict of helpers
	keyed by those ufuncs. If the value is `None`, the ufunc is
	explicitly not supported.
	"""
	with self._lock:
	self.modules[module] = {"names": names, "importer": importer}

Is there a way to register new ufuncs with astropy? #19005

Description

What is the problem this feature will solve?

Describe the desired outcome

Additional context

Activity

neutrinoceros commented on Dec 3, 2025

neutrinoceros commented on Dec 3, 2025

mhvk commented on Dec 3, 2025

roytsmart commented on Dec 3, 2025

mhvk commented on Dec 3, 2025

roytsmart commented on Dec 4, 2025

mhvk commented on Dec 4, 2025

roytsmart commented on Dec 4, 2025

nstarman commented on Dec 5, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions