NumPy has a mechanism built in to allow subclasses to adjust or override aspects of the ufunc behavior. While this goal is important, this mechanism only allows for very limited customization, making for instance the masked arrays unable to work with the native ufuncs in a full and proper way. I would like to deprecate the current mechanism, in particular __array_prepare__ and __array_wrap__, and introduce a new method I will describe below. If you've ever used these mechanisms, please review this design to see if it meets your needs.

Any class type which would like to override its behavior in ufuncs would define a method called _numpy_ufunc_, and optionally an attribute __array_priority__ as can already be done. The class which wins the priority battle gets its _numpy_ufunc_ function called as follows:

return arr._numpy_ufunc_(current_ufunc, *args, **kwargs)

To support this overloading, the ufunc would get a new support method, result_type, and there would be a new global function, broadcast_empty_like.

The function ufunc.empty_like behaves like the global np.result_type, but produces the output type or a tuple of output types specific to the ufunc, which may follow a different convention than regular arithmetic type promotion. This allows for a class to create an output array of the correct type to pass to the ufunc if it needs to be different than the default.

The function broadcast_empty_like is just like empty_like, but takes a list or tuple of arrays which are to be broadcast together for producing the output, instead of just one.

Thanks,

Mark

A simple class which overrides the ufuncs might look as follows:

def sin(ufunc, *args, **kwargs):

    # Convert degrees to radians

    args[0] = np.deg2rad(args[0])

    # Return a regular array, since the result is not in degrees

    return ufunc(*args, **kwargs)

class MyDegreesClass:

    """Array-like object with a degrees unit"""

    def __init__(arr):

        self.arr = arr

    def _numpy_ufunc_(ufunc, *args, **kwargs):

        override = globals().get(ufunc.name)

        if override:

            return override(ufunc, *args, **kwargs)

        else:

            raise TypeError, 'ufunc %s incompatible with MyDegreesClass' % ufunc.name

A more complex example will be something like this:

def my_general_ufunc(ufunc, *args, **kwargs):

    # Extract the 'out' argument. This only supports ufuncs with

    # one output, currently.

    out = kwargs.get('out')

    if len(args) > ufunc.nin:

        if out is None:

            out = args[ufunc.nin]

        else:

            raise ValueError, "'out' given as both a position and keyword argument"

    # Just want the inputs from here on

    args = args[:ufunc.nin]

    # Strip out MyArrayClass, but allow operations with regular ndarrays

    raw_in = []

    for a in args:

        if isinstance(a, MyArrayClass):

            raw_in.append(a.arr)

        else:

            raw_in.append(a)

    # Allocate the output array

    if not out is None:

        if isinstance(out, MyArrayClass):

            raise TypeError, "'out' must have type MyArrayClass"

    else:

        # Create the output array, obeying the 'order' parameter,

        # but disallowing subclasses

        out = np.broadcast_empty_like([args,

                                order=kwargs.get('order'),

                                dtype=ufunc.result_type(args),

                                subok=False)

    # Override the output argument

    kwargs['out'] = out.arr

    # Call the ufunc

    ufunc(*args, **kwargs)

    # Return the output

    return out

class MyArrayClass:

    def __init__(arr):

        self.arr = arr

    def _numpy_ufunc_(ufunc, *args, **kwargs):

        override = globals().get(ufunc.name)

        if override:

            return override(ufunc, *args, **kwargs)

        else:

            return my_general_ufunc(ufunc, *args, **kwargs)