[Numpy-discussion] replacing the mechanism for dispatching ufuncs

[Mark Wiebe]

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)
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20110620/c378d06e/attachment.html>