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是一个选项，我肯定会建议jit 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列表BTW执行完全相同的操作。）

<强> Vectorizability

我继续想知道OP算法的可引导性。似乎它不可矢量化。引用Cornell Virtual Workshop：

写入后读取（＆＃34;流程＆＃34;）依赖。 这种依赖不可矢量化。它发生在特定循环迭代中的变量值（＆＃34;读取＆＃34;）由前一个循环迭代（＆＃34;写＆＃34;）确定。

＆＃34;流程＆＃34;在循环中看到依赖性，其中有几个成员＆＃39;值由同一成员的先前状态确定。例如：

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

此处，self.force来自上一次迭代，是根据之前的self.vel计算的。

如何加速Numpy sum和Python for循环？

2 个答案: