python - 如何高效使用numpy进行迭代求和

Question

问题的定义

我可以这样表达我试图解决的问题:

给定 A [Nx1] 和 n[Nx1] 和 x[Mx1]，我想执行此操作

S = sum([A[i]*x**n[i] for i in range(len(n))])

使用numpy。我想我可以使用 numpy 中的广播之类的东西来做到这一点，但我无法理解 numpy 文档。有人可以帮我弄清楚如何在 numpy 中有效地做到这一点吗？

我有一个针对下面这个问题的有效 cython 解决方案，速度非常快，我想知道我是否可以使用 numpy 更轻松地完成它并完全避免 cython。

Cython 解决方案

这是使用 cython 来演示该问题的此问题的有效实现:

cimport cython
import numpy as np

@cython.boundscheck(False)
cpdef  sum_function( double [:] A, double [:] x, double [:] n, double [:] out):
    cdef int i,j
    cdef int Nx = len(x)
    cdef int Nn = len(n)

    out[:] = 0

    for i in xrange(Nx):
        for j in range(Nn):
            out[i] += A[j]*x[i]**n[j]

Answer 1

您可以使用np.sum(A * x.reshape(-1, 1)**n, axis=1):

In [40]: A
Out[40]: array([ 1.,  2., -1.,  3.])

In [41]: n
Out[41]: array([ 2.,  1.,  1.,  3.])

In [42]: x
Out[42]: array([  5.,  10.,   1.])

In [43]: S = sum([A[i]*x**n[i] for i in range(len(n))])

In [44]: S
Out[44]: array([  405.,  3110.,     5.])

In [45]: np.sum(A * x.reshape(-1, 1)**n, axis=1)
Out[45]: array([  405.,  3110.,     5.])

x.reshape(-1, 1) 具有形状 (3,1)，n 具有形状 (4,)，因此 x.reshape(-1, 1)**n 的(广播)结果具有形状 (3,4)； k 列包含 x**n[k]。 A 具有形状 (4,)，因此 A * x.reshape(-1, 1)**n 具有形状 (3,4)；该乘积的元素 (i,j) 包含 A[j]*x[i]**n[j]。所需的结果是该数组沿 axis=1 的总和。

这是具有相同数据的 cython 版本:

In [46]: out = np.zeros_like(x)

In [47]: sum_function(A, x, n, out)

In [48]: out
Out[48]: array([  405.,  3110.,     5.])