在同时包含数字和对象的数组上使用NumPy函数

Question

我有一个类似number的类，我在其中实现了sqrt、exp等方法，这样当NumPy函数在ndarray中时就会对它们进行广播。

class A:
    def sqrt(self):
        return 1.414

这在以下数组中可以完美地工作：

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

显然，sqrt也适用于纯数字：

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

但是，当数字和对象混合在一起时，这不起作用：

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'

这是因为异构数组的dtype是object，并且numpy函数通过在每个对象上调用相同名称的方法来广播，但是numbers没有使用这些名称的方法。

我该如何解决这个问题？

Answer 1

不确定效率，但作为一种解决办法，您可以使用使用map和isinstance创建的布尔索引，然后对两个切片应用相同的操作，更改不属于A类的元素的类型，以便能够使用numpy方法。

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)

注意:在astype中，我使用了float，但不确定它是否适合您的要求