Use NumPy function on arrays containing both numbers and objects

Question

I have a number-like class where I have implemented the sqrt, exp, etc. methods so that NumPy functions will broadcast on them when they are in ndarrays.

class A:
    def sqrt(self):
        return 1.414

This work perfectly in an array of these:

import numpy as np
print(np.sqrt([A(), A()]))  # [1.414 1.414]

Obviously, sqrt also works with pure numbers:

print(np.sqrt([4, 9]))  # [2. 3.]

However, this does not work when numbers and objects are mixed:

print(np.sqrt([4, A()]))

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-38-0c4201337685> in <module>()
----> 1 print(np.sqrt([4, A()]))

AttributeError: 'int' object has no attribute 'sqrt'

This happens because the dtype of the heterogenous array is object and numpy functions broadcast by calling a method of the same name on each object, but numbers do not have methods with these names.

How do I work around this?

Answer 1

Not sure about the efficiency, but as a work around, you can use boolean indexing created with map and isinstance and then apply to both slices the same operation, changing the type of the element that are not of class A to be able to use the numpy method.

ar = np.array([4, A(), A(), 9.])
ar_t = np.array(list(map(lambda x: isinstance(x, A), ar)))

ar[~ar_t] = np.sqrt(ar[~ar_t].astype(float)) 
ar[ar_t] = np.sqrt(ar[ar_t])

print(ar)
# array([2.0, 1.414, 1.414, 3.0], dtype=object)

Note: in the astype, I used float, but not sure it will be good for your requirement