如何加速 Numpy sum 和 Python for-loop？

Question

我有 2 个代码几乎做同样的事情。

代码 1：

from __future__ import division
import numpy as np

m = 1
gamma = 1
lam = 1
alpha = 1
step_num = 2 ** 16
dt = 0.02


def E_and_x(x0):
    xi = x0
    vi = 0
    f = 0
    xsum = 0
    Ei, xavg = 0, 0
    for i in range(step_num):
        vi += f / m * dt / 2
        xi += vi * dt
        f = - gamma * xi - lam * xi ** 2 - alpha * xi ** 3
        vi += f / m * dt / 2
        Ei = 1 / 2 * m * vi ** 2 + 1 / 2 * gamma * xi ** 2 + \
            1 / 3 * lam * xi ** 3 + 1 / 4 * alpha * xi ** 4
        xsum += xi
        xavg = xsum / (i + 1)
    return Ei, xavg

E, x = [], []
for x0 in np.linspace(0, 1, 40):
    mdresult = E_and_x(x0)
    E.append(mdresult[0])
    x.append(mdresult[1])

代码 2：

from __future__ import division
import numpy as np
from numba import jit

time = 50
niter = 2 ** 16  # number of iterations
t = np.linspace(0, time, num=niter, endpoint=True)


class MolecularDynamics(object):
    def __init__(self, time, niter, initial_pos):
        self.position = np.array([])
        self.velocity = np.array([])
        self.energy = np.array([])
        self.x_average = np.array([])
        self.vel = 0  # intermediate variable
        self.force = 0  # intermediate variable
        self.e = 0  # intermediate energy
        self.initial_pos = initial_pos  # initial position
        self.pos = self.initial_pos
        self.time = time
        self.niter = niter
        self.time_step = self.time / self.niter
        self.mass = 1
        self.k = 1  # stiffness coefficient
        self.lamb = 1  # lambda
        self.alpha = 1  # quartic coefficient

    @jit
    def iter(self):
        for i in np.arange(niter):
            # step 1 of leap frog
            self.vel += self.time_step / 2.0 * self.force / self.mass
            self.pos += self.time_step * self.vel
            # step 2 of leap frog
            self.force = - self.k * self.pos - self.lamb * self.pos ** 2 - self.alpha * self.pos ** 3
            self.vel += self.time_step / 2.0 * self.force / self.mass

            # calculate energy
            self.e = 1 / 2 * self.mass * self.vel ** 2 + \
                     1 / 2 * self.k * self.pos ** 2 + \
                     1 / 3 * self.lamb * self.pos ** 3 + \
                     1 / 4 * self.alpha * self.pos ** 4

            self.velocity = np.append(self.velocity, [self.vel])  # record vel after 1 time step
            self.position = np.append(self.position, self.pos)  # record pos after 1 time step
            self.energy = np.append(self.energy, [self.e])  # record e after 1 time step
            self.x_average = np.append(self.x_average, np.sum(self.position) / (i + 1))


mds = [MolecularDynamics(time, niter, xx) for xx in np.linspace(0, 1, num=40)]
[md.iter() for md in mds]  # loop to change value
mds_x_avg = [md.x_average[-1] for md in mds]
mds_e = [md.e for md in mds]

嗯，主要区别在于代码 2 使用 OO、Numpy 和 JIT。但是，代码 2 比代码 1 慢得多（计算需要很长时间）。

In [1]: %timeit code_1
10000000 loops, best of 3: 25.7 ns per loop

通过分析，我知道瓶颈是iter() 函数，更具体地说，在append 和sum 上。但尽我所能使用 Numpy， 我想知道为什么代码 2 慢得多，我该如何加快速度？

Answer 1

你做错了你的时间，只是测试你的第一个代码（稍微修改）：

from __future__ import division


def E_and_x(x0):
    m = 1
    gamma = 1
    lam = 1
    alpha = 1
    step_num = 2 ** 13   # much less iterations!
    dt = 0.02

    xi = x0
    vi = 0
    f = 0
    xsum = 0
    Ei, xavg = 0, 0
    for i in range(step_num):
        vi += f / m * dt / 2
        xi += vi * dt
        f = - gamma * xi - lam * xi ** 2 - alpha * xi ** 3
        vi += f / m * dt / 2
        Ei = 1 / 2 * m * vi ** 2 + 1 / 2 * gamma * xi ** 2 + \
            1 / 3 * lam * xi ** 3 + 1 / 4 * alpha * xi ** 4
        xsum += xi
        xavg = xsum / (i + 1)
    return Ei, xavg

时间不是在纳秒范围内：

%timeit [E_and_x(x0) for x0 in np.linspace(0, 1, 40)]  # 1 loop, best of 3: 3.46 s per loop

但是，如果numba 是一个选项，我强烈建议您使用E_and_x 函数：

import numba as nb

numbaE_and_x = nb.njit(E_and_x)

numbaE_and_x(1.2)  # warmup for the jit
%timeit [numbaE_and_x(x0) for x0 in np.linspace(0, 1, 40)]  # 100 loops, best of 3: 3.38 ms per loop

它已经快了 100 倍。如果您使用 PyPy（或 Cythonize）运行第一个代码，您应该会得到类似的结果。

除此之外：

• np.append 是一个糟糕的选择。因为np.append、np.concatenate、np.stack（所有变体）需要分配一个新数组并将所有其他数组复制到其中！你不需要对这些数组做任何事情，只需附加到它们。因此，您对 numpy 所做的唯一事情就是 numpy 非常不擅长的事情！
• 检查您是否可以矢量化任何使用 numpy-arrays 的计算（无需附加等）。
• 如果仍然要减慢所有这些 self.xxx 属性访问速度很慢，最好先读取它们一次，然后再设置它们。

Answer 2

除了 MSeifert 所说的，您可以将数组预分配到正确的大小，而不是附加到它们。所以不要像这样创建它们：

self.position = np.array([])  # No

你会写：

self.position = np.zeros(niter) # Yes

然后不要像这样追加：

self.velocity = np.append(self.velocity, [self.vel])

你可以这样填写：

self.velocity[i] = self.vel

这避免了在每次迭代时重新分配数组（您可以使用 array = [someValue]*size 对原始 python 列表进行完全相同的操作）。

矢量化

我一直想知道 OP 算法的矢量化能力。它似乎不可矢量化。引用Cornell Virtual Workshop：

写后读（“流”）依赖。 这种依赖是不可向量化的。当特定循环迭代（“读取”）中的变量值由先前的循环迭代（“写入”）确定时，就会发生这种情况。

在循环中可以看到“流”依赖关系，其中几个成员的值由同一成员的先前状态确定。例如：

self.vel += self.time_step / 2.0 * self.force / self.mass

这里，self.force 来自上一次迭代，是根据上一次 self.vel 计算得出的。