CODEWITHSUNDEEP
pythonnumpy
acdr
acdr

Reputation: 4706

Numpy functions return array class instance when called on subclass of ndarray

Some numpy functions (logically) return scalars:

>>> my_arr = np.ndarray(shape=(1,))
>>> type(np.max(my_arr))
<type 'numpy.float64'>

but only when called with an ndarray, rather than a subclass:

>>> class CustomArray(np.ndarray):
...     pass
>>> my_arr = CustomArray(shape=(1,))
>>> type(np.max(my_arr))
<class '__main__.CustomArray'>

Why is this? I'd expect either both to return a scalar (of type <type 'numpy.float64'>, or the former to return a np.ndarray instance and the latter a CustomArray instance. But instead, I get a combination of these two behaviours. Can I change this behaviour through changing my own class?

I don't see anything that would explain this on the doc page discussing subclassing ndarray (http://docs.scipy.org/doc/numpy-1.9.2/user/basics.subclassing.html).

(Running Python 2.7.10, numpy 1.9.2, in case it matters.)

Upvotes: 4

Views: 1142

Answers (2)

nfrasser
nfrasser

Reputation: 415

I also ran into this, solved it for all aggregation/reduction operations in NumPy by implementing a custom __array_wrap__:

import numpy as np

class CustomArray(np.ndarray):
    def __array_wrap__(self, obj, **kwargs):
        if obj.shape == ():
            return obj[()]
        else:
            return super().__array_wrap__(obj, **kwargs)

Example return types for various operations:

>>> a = CustomArray(shape=(3,))
>>> type(np.max(a))
<class 'numpy.float64'>
>>> type(np.median(a))
<class 'numpy.float64'>
>>> type(np.exp(a))
<class '__main__.CustomArray'>
>>>

Upvotes: 0

Sudeep Juvekar
Sudeep Juvekar

Reputation: 5068

This is because max() is not overloaded in CustomArray. If you try it, my_array.max() returns an object of CustomArray instead of scalar. 

my_array = CustomArray(shape=(1,))
print my_array.max()
>> CustomArray(9.223372036854776e+18)

np.max internally calls np.amax, which ends up calling np.maximum.reduce. This is the standard reduce of map-reduce and returns a base-object returned by max. Hence, the type returned by np.max is in fact the type returned by max() method called on your object. You can override it as:

class CustomArray(np.ndarray):
   def max(self, axis, out):
      return np.ndarray(self.shape, buffer=self).max(axis, out)

type(np.max(my_arr))
>> numpy.float64

The trick is to upcast self as an np.ndarray and find max using it.

Upvotes: 1

Related Questions