在numpy中使用数组优化操作

时间:2014-01-16 22:50:46

标签: python optimization numpy numerical-computing

我必须运用一些我写过的数学公式 在python中:

    for s in range(tdim):
        sum1 = 0.0
        for i in range(dim):
            for j in range(dim):
                sum1+=0.5*np.cos(theta[s]*(i-j))*
                eig1[i]*eig1[j]+eig2[i]+eig2[j])-0.5*np.sin(theta[s]*(i-j))*eig1[j]*eig2[i]-eig1[i]*eig2[j])

        PHi2.append(sum1)

现在,这是正确的,但显然效率低下,另一种方法是:

for i in range(dim):
            for j in range(dim):
                PHi2 = 0.5*np.cos(theta*(i-j))*(eig1[i]*eig1[j]+eig2[i]+eig2[j])-0.5*np.sin(theta*(i-j))*(eig1[j]*eig2[i]-eig1[i]*eig2[j])

然而,第二个例子在PHi2的所有元素中给出了相同的数字,所以这个 更快,但答案是错误的。你怎么能正确而有效地做到这一点?

注意:eig1和eig2的尺寸d相同,theta和PHi2的尺寸D相同, 但是d!= D。

3 个答案:

答案 0 :(得分:5)

您可以使用暴力广播方法,但是您正在创建一个形状为(D, d, d)的中间数组,如果您的数组甚至是中等大小,它可能会失控。此外,在使用没有改进的广播时,您将从最里面的循环重新计算大量计算,您只需要执行一次。如果您首先为i - j的所有可能值计算必要参数并将它们加在一起,则可以在外循环上重用这些值,例如:

def fast_ops(eig1, eig2, theta):
    d = len(eig1)
    d_arr = np.arange(d)
    i_j = d_arr[:, None] - d_arr[None, :]
    reidx = i_j + d - 1
    mult1 = eig1[:, None] * eig1[ None, :] + eig2[:, None] + eig2[None, :]
    mult2 = eig1[None, :] * eig2[:, None] - eig1[:, None] * eig2[None, :]
    mult1_reidx = np.bincount(reidx.ravel(), weights=mult1.ravel())
    mult2_reidx = np.bincount(reidx.ravel(), weights=mult2.ravel())

    angles = theta[:, None] * np.arange(1 - d, d)

    return 0.5 * (np.einsum('ij,j->i', np.cos(angles), mult1_reidx) -
                  np.einsum('ij,j->i', np.sin(angles), mult2_reidx))

如果我们将M4rtini的代码重写为比较函数:

def fast_ops1(eig1, eig2, theta):
    d = len(eig1)
    D = len(theta)
    s = np.array(range(D))[:, None, None]
    i = np.array(range(d))[:, None]
    j = np.array(range(d))
    ret = 0.5 * (np.cos(theta[s]*(i-j))*(eig1[i]*eig1[j]+eig2[i]+eig2[j]) -
                 np.sin(theta[s]*(i-j))*(eig1[j]*eig2[i]-eig1[i]*eig2[j]))
    return ret.sum(axis=(-1, -2))

我们编制了一些数据:

d, D = 100, 200
eig1 = np.random.rand(d)
eig2 = np.random.rand(d)
theta = np.random.rand(D)

速度提升非常明显,超过原始代码的115倍,导致速度超过9000倍:

In [22]: np.allclose(fast_ops1(eig1, eig2, theta), fast_ops(eig1, eig2, theta))
Out[22]: True

In [23]: %timeit fast_ops1(eig1, eig2, theta)
10 loops, best of 3: 145 ms per loop

In [24]: %timeit fast_ops(eig1, eig2, theta)
1000 loops, best of 3: 1.85 ms per loop

答案 1 :(得分:1)

这是通过广播来实现的 对于tdim = 200 and dim = 100
14秒与原件 120毫秒的版本。

s = np.array(range(tdim))[:, None, None]
i = np.array(range(dim))[:, None]
j = np.array(range(dim))
PHi2 =(0.5*np.cos(theta[s]*(i-j))*(eig1[i]*eig1[j]+eig2[i]+eig2[j])-0.5*np.sin(theta[s]*(i-j))*(eig1[j]*eig2[i]-eig1[i]*eig2[j])).sum(axis=2).sum(axis=1)

答案 2 :(得分:0)

在代码的第一位,您有0.5*np.cos(theta[s]*(i-j))...但在第二位是0.5*np.cos(theta*(i-j))...。除非你为第二位代码定义了不同的,否则这很可能是造成麻烦的原因。