使用两个不同维度的参数来矢量化/加速numpy函数

时间:2015-11-07 20:47:56

标签: python numpy parallel-processing vectorization

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

my_func(t, x, p, a)

t是一个标量浮点数 x是1D numpy浮点阵列 pa是字典。

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

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?

我看到的另一种可能性是预先切割数组并使用正确的参数形成元组列表,并使用某种形式的“map”功能。

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

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元素的向量。 在这种情况下,该函数仅返回一个值。但是可能会出现返回列表的情况。在这种情况下,返回的值在输出中形成一列。希望能澄清一点。

1 个答案:

答案 0 :(得分:0)

您的for循环不是untythonic。 Python喜欢循环。 :)你的意思是,它可能没有充分利用numpy。这是不是numpy-onic?

我不确定将它投入到诸如numexpr,vectorize甚至apply_along_axis之类的黑盒子中会更加难以理解。理想的是理解问题并在该结构内/周围工作以使其处理更大的结构。

让我们尝试一个示例函数(就像你应该给我们的那样?):

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.没有帮助。

这个TX广播就好了,例如T*X就可以了。

vectorize的问题在于它会向您的函数传递从您的TX获取的标量值元组。它的意思是向量化一个带有标量的函数,而不是标量和向量。

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

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循环就像它得到的一样好。大多数替代方案难以做到正确,更加模糊,而且可能不会更快。