Getting started#
Mixins for adding __array_ufunc__
&/or __array_function__
methods.
Examples#
First, some imports:
>>> from dataclasses import dataclass, fields
>>> from typing import ClassVar
>>> import numpy as np
>>> from overload_numpy import NumPyOverloader, NPArrayOverloadMixin
Now we can define a NumPyOverloader
instance:
>>> W_FUNCS = NumPyOverloader()
The overloads apply to an array wrapping class. Let’s define one:
>>> @dataclass
... class Wrap1D(NPArrayOverloadMixin):
... '''A simple array wrapper.'''
... x: np.ndarray
... NP_OVERLOADS: ClassVar[NumPyOverloader] = W_FUNCS
>>> w1d = Wrap1D(np.arange(3))
Implementing an Overload#
Now both numpy.ufunc
(e.g. numpy.add
) and numpy
functions
(e.g. numpy.concatenate()
) can be overloaded and registered for
Wrap1D
.
>>> @W_FUNCS.implements(np.add, Wrap1D)
... def add(w1, w2):
... return Wrap1D(np.add(w1.x, w2.x))
>>> @W_FUNCS.implements(np.concatenate, Wrap1D)
... def concatenate(w1ds):
... return Wrap1D(np.concatenate(tuple(w.x for w in w1ds)))
Time to check these work:
>>> np.add(w1d, w1d)
Wrap1D(x=array([0, 2, 4]))
>>> np.concatenate((w1d, w1d))
Wrap1D(x=array([0, 1, 2, 0, 1, 2]))
ufunc
also have a number of methods: ‘at’, ‘accumulate’, ‘outer’, etc. The
function dispatch mechanism in NEP13 says that “If one of
the input or output arguments implements __array_ufunc__, it is executed instead
of the ufunc.” Currently the overloaded numpy.add
does not work for any of the
ufunc
methods.
>>> try: np.add.accumulate(w1d)
... except Exception: print("failed")
failed
ufunc
method overloads can be registered on the wrapped add
implementation:
>>> @add.register('accumulate')
... def add_accumulate(w1):
... return Wrap1D(np.add.accumulate(w1.x))
>>> np.add.accumulate(w1d)
Wrap1D(x=array([0, 1, 3]))
Dispatching Overloads for Subclasses#
What if we defined a subclass of Wrap1D
?
>>> @dataclass
... class Wrap2D(Wrap1D):
... '''A simple 2-array wrapper.'''
... y: np.ndarray
The overload for numpy.concatenate()
registered on Wrap1D
will not
work correctly for Wrap2D
. However, NumPyOverloader
supports
single-dispatch on the calling type for the overload, so overloads can be
customized for subclasses.
>>> @W_FUNCS.implements(np.add, Wrap2D)
... def add(w1, w2):
... print("using Wrap2D implementation...")
... return Wrap2D(np.add(w1.x, w2.x),
... np.add(w1.y, w2.y))
>>> @W_FUNCS.implements(np.concatenate, Wrap2D)
... def concatenate2(w2ds):
... print("using Wrap2D implementation...")
... return Wrap2D(np.concatenate(tuple(w.x for w in w2ds)),
... np.concatenate(tuple(w.y for w in w2ds)))
Checking these work:
>>> w2d = Wrap2D(np.arange(3), np.arange(3, 6))
>>> np.add(w2d, w2d)
using Wrap2D implementation...
Wrap2D(x=array([0, 2, 4]), y=array([ 6, 8, 10]))
>>> np.concatenate((w2d, w2d))
using Wrap2D implementation...
Wrap2D(x=array([0, 1, 2, 0, 1, 2]), y=array([3, 4, 5, 3, 4, 5]))
Great! But rather than defining a new implementation for each subclass, let’s see how we could write a more broadly applicable overload:
>>> @W_FUNCS.implements(np.add, Wrap1D) # overriding both
... @W_FUNCS.implements(np.add, Wrap2D) # overriding both
... def add_general(w1, w2):
... WT = type(w1)
... return WT(*(np.add(getattr(w1, f.name), getattr(w2, f.name))
... for f in fields(WT)))
>>> @W_FUNCS.implements(np.concatenate, Wrap1D) # overriding both
... @W_FUNCS.implements(np.concatenate, Wrap2D) # overriding both
... def concatenate_general(ws):
... WT = type(ws[0])
... return WT(*(np.concatenate(tuple(getattr(w, f.name) for w in ws))
... for f in fields(WT)))
Checking these work:
>>> np.add(w2d, w2d)
Wrap2D(x=array([0, 2, 4]), y=array([ 6, 8, 10]))
>>> np.concatenate((w2d, w2d))
Wrap2D(x=array([0, 1, 2, 0, 1, 2]), y=array([3, 4, 5, 3, 4, 5]))
>>> @dataclass
... class Wrap3D(Wrap2D):
... '''A simple 3-array wrapper.'''
... z: np.ndarray
>>> w3d = Wrap3D(np.arange(2), np.arange(3, 5), np.arange(6, 8))
>>> np.add(w3d, w3d)
Wrap3D(x=array([0, 2]), y=array([6, 8]), z=array([12, 14]))
>>> np.concatenate((w3d, w3d))
Wrap3D(x=array([0, 1, 0, 1]), y=array([3, 4, 3, 4]), z=array([6, 7, 6, 7]))
Assisting Groups of Overloads#
In the previous examples we wrote implementations for a single NumPy function. Overloading the full set of NumPy functions this way would take a long time.
Wouldn’t it be better if we could write many fewer, based on groups of NumPy functions?
>>> add_funcs = {np.add, np.subtract}
>>> @W_FUNCS.assists(add_funcs, types=Wrap1D, dispatch_on=Wrap1D)
... def add_assists(cls, func, w1, w2, *args, **kwargs):
... return cls(*(func(getattr(w1, f.name), getattr(w2, f.name), *args, **kwargs)
... for f in fields(cls)))
>>> stack_funcs = {np.vstack, np.hstack, np.dstack, np.column_stack, np.row_stack}
>>> @W_FUNCS.assists(stack_funcs, types=Wrap1D, dispatch_on=Wrap1D)
... def stack_assists(cls, func, ws, *args, **kwargs):
... return cls(*(func(tuple(getattr(v, f.name) for v in ws), *args, **kwargs)
... for f in fields(cls)))
Checking these work:
>>> np.subtract(w2d, w2d)
Wrap2D(x=array([0, 0, 0]), y=array([0, 0, 0]))
>>> np.vstack((w1d, w1d))
Wrap1D(x=array([[0, 1, 2],
[0, 1, 2]]))
>>> np.hstack((w1d, w1d))
Wrap1D(x=array([0, 1, 2, 0, 1, 2]))
We would also like to implement the accumulate
method for all the
add_funcs
overloads:
>>> @add_assists.register("accumulate")
... def add_accumulate_assists(cls, func, w1, *args, **kwargs):
... return cls(*(func(getattr(w1, f.name), *args, **kwargs)
... for f in fields(cls)))
>>> np.subtract.accumulate(w2d)
Wrap2D(x=array([ 0, -1, -3]), y=array([ 3, -1, -6]))
Where to go from here?#
This page is meant to demonstrate a few initial things you may want to do with
overload_numpy
. There is much more functionality that you can discover by
perusing the user guide. Some other
commonly-used functionality includes: