NPArrayFuncOverloadMixin

class overload_numpy.NPArrayFuncOverloadMixin

Bases: object

Mixin for adding __array_function__ to a class.

This mixin adds the method __array_function__. Subclasses must define a class variable NP_OVERLOADS and optionally NP_FUNC_TYPES.

Notes

When compiled this class is a mypyc trait() and permits interpreted subclasses (see https://mypyc.readthedocs.io/en/latest/native_classes.html#inheritance).

Examples

First, some imports:

>>> from dataclasses import dataclass, fields
>>> from typing import ClassVar
>>> import numpy as np
>>> from overload_numpy import NumPyOverloader, NPArrayFuncOverloadMixin

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(NPArrayFuncOverloadMixin):
...     '''A simple array wrapper.'''
...     x: np.ndarray
...     NP_OVERLOADS: ClassVar[NumPyOverloader] = W_FUNCS

Now numpy functions can be overloaded and registered for Wrap1D.

>>> @W_FUNCS.implements(np.concatenate, Wrap1D)
... def concatenate(w1ds):
...     return Wrap1D(np.concatenate(tuple(w.x for w in w1ds)))

Time to check this works:

>>> w1d = Wrap1D(np.arange(3))
>>> np.concatenate((w1d, w1d))
Wrap1D(x=array([0, 1, 2, 0, 1, 2]))

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.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 this works:

>>> w2d = Wrap2D(np.arange(3), np.arange(3, 6))
>>> 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.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 this works:

>>> 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.concatenate((w3d, w3d))
Wrap3D(x=array([0, 1, 0, 1]), y=array([3, 4, 3, 4]), z=array([6, 7, 6, 7]))

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?

>>> 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 this works:

>>> 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]))

Attributes:

NP_OVERLOADSNumPyOverloader

How the overrides are registered. A class-attribute of an instance of NumPyOverloader.

NP_FUNC_TYPESfrozenset[type | TypeConstraint] | python:None, optional

Default type constraints. A class-attribute of None or a frozenset of type or TypeConstraint. If None, then the types argument for overloading functions becomes mandatory. If a frozenset (including blank) then the self-type + the contents are used in the types check. See __array_function__ for details of the types argument.

Attributes Summary

NP_FUNC_TYPES

Attributes Documentation

NP_FUNC_TYPES: ClassVar[frozenset[type | TypeConstraint] | None] = frozenset({})