Python:如何使类 `__call__` 方法接受 NumPy 数组和单个值？

Question

对于像这样的常规函数​​

def f(t):
    return t*t

我可以毫无问题地传递一个值或一个 NumPy 数组两者。例如。这有效:

T = 1
print(f(T))

times = np.mgrid[0 : T : 100j]
values = f(times)

现在我创建了一个带有__call__ 函数的类

class rnd_elemental_integrand:
    
    def __init__(self, n_sections, T):
        
        self.n_sections = n_sections
        self.T = T
        self.generate()
        
    def generate(self):
        self.values = norm.rvs(size = (self.n_sections + 1,), scale = 1)

    def __call__(self, t):
        ind = int(t * (self.n_sections/self.T))
        return self.values[ind]

但对于此类方法，我可以不传递一个 NumPy 数组。例如。这个

T = 5
elem_int_sections = 10

rnd_elem = rnd_elemental_integrand(elem_int_sections, T)

print(rnd_elem(T))
times = np.mgrid[0 : T : 100j]
values = rnd_elem(times)

产生输出

0.43978851468955377
Traceback (most recent call last):
  File "/Users/gnthr/Desktop/Programming/Python/StochAna/stochana.py", line 138, in <module>
    values = rnd_elem(times)
  File "/Users/gnthr/Desktop/Programming/Python/StochAna/stochana.py", line 117, in __call__
    ind = int(t * (self.n_sections/self.T))
TypeError: only size-1 arrays can be converted to Python scalars

从其他帖子我知道通过一些 np. 函数向量化一个 __call__ 方法是可行的，但是例如上面的函数 f 也没有向量化，并且可以很好地处理两种类型的输入。可以使此类 __call__ 方法接受两种参数类型( float 和 float 数组)吗？

Answer 1

修正:没有类型检查。

由于 np.array 可以接受 np.array 和标量的输入，我们可以创建一个新的 np.array 类型的 int

ind = np.array(t * (self.n_sections/self.T), dtype=int)

测试用例:

from scipy.stats import norm
T = 5
elem_int_sections = 10

rnd_elem = rnd_elemental_integrand(elem_int_sections, T)

print(rnd_elem(T))
times = np.mgrid[0 : T : 100j]
print (rnd_elem(times))

输出:

-0.7828585207846585
[-1.00037782 -1.00037782 -1.00037782 -1.00037782 -1.00037782 -1.00037782
 -1.00037782 -1.00037782 -1.00037782 -1.00037782  1.35744571  1.35744571
  1.35744571  1.35744571  1.35744571  1.35744571  1.35744571  1.35744571
  1.35744571  1.35744571  0.65442428  0.65442428  0.65442428  0.65442428
  0.65442428  0.65442428  0.65442428  0.65442428  0.65442428  0.65442428
  0.76685108  0.76685108  0.76685108  0.76685108  0.76685108  0.76685108
  0.76685108  0.76685108  0.76685108  0.76685108  0.48888641  0.48888641
  0.48888641  0.48888641  0.48888641  0.48888641  0.48888641  0.48888641
  0.48888641  0.48888641  0.62681856  0.62681856  0.62681856  0.62681856
  0.62681856  0.62681856  0.62681856  0.62681856  0.62681856  0.62681856
  1.05695641  1.05695641  1.05695641  1.05695641  1.05695641  1.05695641
  1.05695641  1.05695641  1.05695641  1.05695641 -0.0634099  -0.0634099
 -0.0634099  -0.0634099  -0.0634099  -0.0634099  -0.0634099  -0.0634099
 -0.0634099  -0.0634099  -0.00167191 -0.00167191 -0.00167191 -0.00167191
 -0.00167191 -0.00167191 -0.00167191 -0.00167191 -0.00167191 -0.00167191
  1.16756173  1.16756173  1.16756173  1.16756173  1.16756173  1.16756173
  1.16756173  1.16756173  1.16756173 -0.78285852]