Subclassing Numpy Array - Propagate Attributes
import numpy as np
class ExampleTensor(np.ndarray):
def __new__(cls, input_array):
return np.asarray(input_array).view(cls)
def __array_finalize__(self, obj) -> None:
if obj is None: return
# This attribute should be maintained!
default_attributes = {"attr": 1}
self.__dict__.update(default_attributes) # another way to set attributes
Implement the array_ufunc method like this
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): # this method is called whenever you use a ufunc
f = {
"reduce": ufunc.reduce,
"accumulate": ufunc.accumulate,
"reduceat": ufunc.reduceat,
"outer": ufunc.outer,
"at": ufunc.at,
"__call__": ufunc,
}
output = ExampleTensor(f[method](*(i.view(np.ndarray) for i in inputs), **kwargs)) # convert the inputs to np.ndarray to prevent recursion, call the function, then cast it back as ExampleTensor
output.__dict__ = self.__dict__ # carry forward attributes
return output
Test
x = ExampleTensor(np.array([1,2,3]))
x.attr = 2
y0 = np.add(x, x)
print(y0, y0.attr)
y1 = np.add.outer(x, x)
print(y1, y1.attr) # works even if called with method
[2 4 6] 2
[[2 3 4]
[3 4 5]
[4 5 6]] 2
Explanation in the comments.
Here is an attempt that works for operators that are not arrays and even when our subclass is specified as output of a numpy ufunc (explanations in the comments):
import numpy as np
class ArraySubclass(np.ndarray):
'''Subclass of ndarray MUST be initialized with a numpy array as first argument.
'''
def __new__(cls, input_array, a=None, b=1):
obj = np.asarray(input_array).view(cls)
obj.a = a
obj.b = b
return obj
def __array_finalize__(self, obj):
if obj is None: # __new__ handles instantiation
return
'''we essentially need to set all our attributes that are set in __new__ here again (including their default values).
Otherwise numpy's view-casting and new-from-template mechanisms would break our class.
'''
self.a = getattr(obj, 'a', None)
self.b = getattr(obj, 'b', 1)
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): # this method is called whenever you use a ufunc
'''this implementation of __array_ufunc__ makes sure that all custom attributes are maintained when a ufunc operation is performed on our class.'''
# convert inputs and outputs of class ArraySubclass to np.ndarray to prevent infinite recursion
args = ((i.view(np.ndarray) if isinstance(i, ArraySubclass) else i) for i in inputs)
outputs = kwargs.pop('out', None)
if outputs:
kwargs['out'] = tuple((o.view(np.ndarray) if isinstance(o, ArraySubclass) else o) for o in outputs)
else:
outputs = (None,) * ufunc.nout
# call numpys implementation of __array_ufunc__
results = super().__array_ufunc__(ufunc, method, *args, **kwargs) # pylint: disable=no-member
if results is NotImplemented:
return NotImplemented
if method == 'at':
# method == 'at' means that the operation is performed in-place. Therefore, we are done.
return
# now we need to make sure that outputs that where specified with the 'out' argument are handled corectly:
if ufunc.nout == 1:
results = (results,)
results = tuple((self._copy_attrs_to(result) if output is None else output)
for result, output in zip(results, outputs))
return results[0] if len(results) == 1 else results
def _copy_attrs_to(self, target):
'''copies all attributes of self to the target object. target must be a (subclass of) ndarray'''
target = target.view(ArraySubclass)
try:
target.__dict__.update(self.__dict__)
except AttributeError:
pass
return target
and here are the corresponding unittests:
import unittest
class TestArraySubclass(unittest.TestCase):
def setUp(self):
self.shape = (10, 2, 5)
self.subclass = ArraySubclass(np.zeros(self.shape))
def test_instantiation(self):
self.assertIsInstance(self.subclass, np.ndarray)
self.assertIs(self.subclass.a, None)
self.assertEqual(self.subclass.b, 1)
self.assertEqual(self.subclass.shape, self.shape)
self.assertTrue(np.array_equal(self.subclass, np.zeros(self.shape)))
sub2 = micdata.arrayasubclass.ArraySubclass(np.zeros(self.shape), a=2)
self.assertEqual(sub2.a, 2)
def test_view_casting(self):
self.assertIsInstance(np.zeros(self.shape).view(ArraySubclass),ArraySubclass)
def test_new_from_template(self):
self.subclass.a = 5
bla = self.subclass[3, :]
self.assertIsInstance(bla, ArraySubclass)
self.assertIs(bla.a, 5)
self.assertEqual(bla.b, 1)
def test_np_min(self):
self.assertEqual(np.min(self.subclass), 0)
def test_ufuncs(self):
self.subclass.b = 2
self.subclass += 2
self.assertTrue(np.all(self.subclass == 2))
self.subclass = self.subclass + np.ones(self.shape)
self.assertTrue(np.all(self.subclass == 3))
np.multiply.at(self.subclass, slice(0, 2), 2)
self.assertTrue(np.all(self.subclass[:2] == 6))
self.assertTrue(np.all(self.subclass[2:] == 3))
self.assertEqual(self.subclass.b, 2)
def test_output(self):
self.subclass.a = 3
bla = np.ones(self.shape)
bla *= 2
np.multiply(bla, bla, out=self.subclass)
self.assertTrue(np.all(self.subclass == 5))
self.assertEqual(self.subclass.a, 3)
P.s. tempname123 got it almost right. However, his answer fails for operators that are not arrays and when his class is specified as output of a ufunc:
>>> ExampleTensor += 1
AttributeError: 'int' object has no attribute 'view'
>>> np.multiply(np.ones((5)), np.ones((5)), out=ExampleTensor)
RecursionError: maximum recursion depth exceeded in comparison