具有两个不同维度的参数的矢量化/加速 numpy 函数

Question

我不确定以前是否有人问过这个问题。我在 SO 或 Google 上找不到很多相关结果。无论如何，这就是我想要做的。我有一个在运行时创建的函数，它接受 4 个参数。

my_func(t, x, p, a)

t 是一个标量浮点数。
x 是一个一维 numpy 浮点数数组。
p 和 a 是字典。

我有一个 numpy 数组 T 和一个 2D numpy 数组 X。我想用 T 的每个元素和 X 的每一列调用该函数，并将结果存储在另一个 numpy 数组中。该函数返回一个构成解决方案列的 numpy 数组。 我当前的、天真的、非常不合 Python 的实现如下：

for i in range(len(T)):
    _u = my_func(T[i],X[:,i],p,a)
    sol.u[:,i] = _u   # The length of u is unrelated to the other variables

如果它允许自定义函数，我可以使用 numexpr。我还查看了 numpy.vectorize 并尝试像这样使用它：

f = np.vectorize(my_func)
f.excluded.add(2)
f.excluded.add(3)
f(T,X,p,a)

这给了我来自函数内部某处的错误“IndexError：标量变量的无效索引”。我显然不明白 numpy 广播是如何工作的。

是否可以使用 np.vectorize 做到这一点？
或者有什么方法可以使用 np.apply_along_axis？

我看到的另一种可能性是预先对数组进行切片，并形成一个具有正确参数的元组列表，并将它们与某种形式的“映射”功能一起使用。

更新：我找到的另一个解决方案：

f = lambda _t, _X: my_func(_t,_X,p,a)
u = np.array(list(map(f, T, list(X.T)))).T

它似乎工作。但这是最有效的方法吗？

更新 2： 这是一个示例函数：

def my_func(_t,_X,_p,_aux):
    [x,y,v,lamX,lamY,lamV,tf,] = _X[:7]

    g = _aux['const']['g']

    theta = -2*atan((lamX*v + sqrt(g**2*lamV**2 - 2*g*lamV*lamY*v + lamX**2*v**2 + lamY**2*v**2))/(g*lamV - lamY*v))
    return theta

_aux 的一个例子是：

_aux = {'const': {'g':-9.81} }

_t 和 _p 在这种情况下不使用。它期望 _X 是一个 7 元素向量。 在这种情况下，该函数只返回一个值。但在某些情况下，它可能会返回一个列表。在这种情况下，返回的值在输出中形成一列。希望能澄清一点。

Answer 1

您的for 循环并非不符合 Python 标准。 Python 喜欢 for 循环。 :) 您的意思是，它可能没有充分利用numpy。这不是 numpy-onic 吗？

我不确定将其放入像 numexpr、vectorize 甚至 apply_along_axis 这样的黑盒中是否更适合 numpy。理想的做法是了解问题并围绕/在该结构内工作以使其能够处理更大的结构。

让我们尝试一个示例函数（就像你应该给我们的那样？）：

In [77]: def myfunction(t,x,p,a):
     print(t.shape)
     print(x.shape)
     print(p,a)
     return t*x

In [78]: f=np.vectorize(myfunction)
In [79]: f.excluded.add(2)   # so you can pass p,a 
In [80]: f.excluded.add(3)
In [81]: T=np.arange(5)
In [83]: X=np.ones((4,5))
In [85]: for i in range(T.shape[0]):
    print(myfunction(T[i],X[:,i],{},{}))
   ....:     
()
(4,)
{} {}
[ 0.  0.  0.  0.]
()
(4,)
{} {}
[ 1.  1.  1.  1.]
...

In [87]: f(T,X,{},{})
()
()
{} {}
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-87-b585bb8fb6bc> in <module>()
----> 1 f(T,X,{},{})
...
/usr/lib/python3/dist-packages/numpy/lib/function_base.py in func(*vargs)
   1566                     the_args[_i] = vargs[_n]
   1567                 kwargs.update(zip(names, vargs[len(inds):]))
-> 1568                 return self.pyfunc(*the_args, **kwargs)
   1569 
   1570             vargs = [args[_i] for _i in inds]

<ipython-input-77-a72fc1b2ad5e> in myfunction(t, x, p, a)
      1 def myfunction(t,x,p,a):
----> 2      print(t.shape)
      3      print(x.shape)
      4      print(p,a)
      5      return t*x

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

所以t 是一个标量。我得到了一个不同的错误，但我认为它与你得到的错误一致。您的函数是否在某个地方使用了t[0]？

---

更正 - 您的 t 可以是标量，x 应该是向量。你的批评者This gives me the error "IndexError: invalid index to scalar variable" from somewhere inside the function. 没有帮助。

---

这个T 和X 一起广播就好了，例如T*X 有效。

vectorize 的问题在于它向您的函数传递了一个从您的T 和X 获取的标量值元组。它旨在向量化一个接受标量的函数，而不是标量和向量。

让我们重新定义函数，使其不关心输入的形状或类型：

In [101]: def myfunction(t,x):
     print(t)
     print(x)
     return t*x
   .....: 
In [102]: f=np.vectorize(myfunction)
In [103]: f(T,X)
0
1.0
0
1.0
1
...
1.0
Out[103]: 
array([[ 0.,  1.,  2.,  3.,  4.],
       [ 0.,  1.,  2.,  3.,  4.],
       [ 0.,  1.,  2.,  3.,  4.],
       [ 0.,  1.,  2.,  3.,  4.]])

尽管 T 是 1d 且 X 是 2d，但它在每次迭代时将 2 个标量传递给函数。

---

但首先让我退后一步。

vectorize 可以使“广播”值更容易，但它并不能加快速度。它仍然遍历值，为每个集合调用你的函数。它根本不会改变你的功能。它只是一个包装器，就像你的 for 循环一样。

appply_along/over_axis 也是一个迭代包装器。再次没有加速。

地图表达式的 Dito。

...

对于采用标量和向量且无法在内部更改的函数，您的 for 循环几乎可以做到。大多数替代方案将更难正确，更晦涩难懂，而且可能不会更快。