对系统微分方程向前欧拉法进行向量化

Question

我正在通过数值方法求解一个一阶微分方程组中的x(t)。这个系统的方程是：

dx/dt = y
dy/dt = -x - a*y(x^2 + y^2 -1)

我已经用前向欧拉法来解决这个问题，代码如下：

def forward_euler():
    h = 0.01
    num_steps = 10000

    x = np.zeros([num_steps + 1, 2]) # steps, number of solutions
    y = np.zeros([num_steps + 1, 2])
    a = 1.

    x[0, 0] = 10. # initial condition 1st solution
    y[0, 0] = 5.

    x[0, 1] = 0.  # initial condition 2nd solution
    y[0, 1] = 0.0000000001

    for step in xrange(num_steps):
        x[step + 1] = x[step] + h * y[step]
        y[step + 1] = y[step] + h * (-x[step] - a * y[step] * (x[step] ** 2 + y[step] ** 2 - 1))

    return x, y

现在我想进一步优化代码，把x和y放在同一个数组里。我想出了以下的解决方案：

def forward_euler_vector():
    num_steps = 10000
    h = 0.01

    x = np.zeros([num_steps + 1, 2, 2]) # steps, variables, number of solutions
    a = 1.

    x[0, 0, 0] = 10. # initial conditions 1st solution
    x[0, 1, 0] = 5.  

    x[0, 0, 1] = 0.  # initial conditions 2nd solution
    x[0, 1, 1] = 0.0000000001

    def f(x): 
        return np.array([x[1],
                         -x[0] - a * x[1] * (x[0] ** 2 + x[1] ** 2 - 1)])

    for step in xrange(num_steps):
        x[step + 1] = x[step] + h * f(x[step])

    return x

问题是：forward_euler_vector()这个函数可以运行，但这样向量化的方式是否是最好的？我问这个问题是因为在我的笔记本上，向量化版本的运行速度慢了大约20毫秒：

In [27]: %timeit forward_euler()
1 loops, best of 3: 301 ms per loop

In [65]: %timeit forward_euler_vector()
1 loops, best of 3: 320 ms per loop

Answer 1

这里有一个简单的 autojit 解决方案：

def forward_euler(initial_x, initial_y, num_steps, h):

     x = np.zeros([num_steps + 1, 2]) # steps, number of solutions
     y = np.zeros([num_steps + 1, 2])
     a = 1.

     x[0, 0] = initial_x[0] # initial condition 1st solution
     y[0, 0] = initial_y[0]

     x[0, 1] = initial_x[1]  # initial condition 2nd solution
     y[0, 1] = initial_y[1]

     for step in xrange(int(num_steps)):
         x[step + 1] = x[step] + h * y[step]
         y[step + 1] = y[step] + h * (-x[step] - a * y[step] * (x[step] ** 2 + y[step] ** 2 - 1))

     return x, y

时间记录：

from numba import autojit
jit_forward_euler = autojit(forward_euler)

%timeit forward_euler([10,0], [5,0.0000000001], 1E4, 0.01)
1 loops, best of 3: 385 ms per loop

%timeit jit_forward_euler([10,0], [5,0.0000000001], 1E4, 0.01)
100 loops, best of 3: 3.51 ms per loop

Answer 2

@Ophion 的评论很好地解释了发生了什么。在 f(x) 中调用 array() 会增加一些额外的开销，这样就失去了在表达式 h * f(x[step]) 中使用矩阵乘法的好处。

正如他所说的，你可能会想看看 scipy.integrate，它提供了一些很不错的数值积分工具。

要解决你代码向量化的问题，你需要避免每次调用 f 时都重新创建数组。你希望只初始化一次数组，并在每次调用时返回修改后的数组。这有点像 C/C++ 中的 static 变量。

你可以通过使用一个可变的默认参数来实现这一点，这个参数在定义函数 f(x) 时只会被解释一次，并且具有局部作用域。因为它需要是可变的，所以你可以把它放在一个只有一个元素的列表中：

 def f(x,static_tmp=[empty((2,2))]): 
    static_tmp[0][0]=x[1]
    static_tmp[0][1]=-x[0] - a * x[1] * (x[0] ** 2 + x[1] ** 2 - 1)
    return static_tmp[0]

通过对你的代码进行这个修改，数组创建的开销就消失了，在我的机器上我得到了小幅的性能提升：

%timeit forward_euler()        #258ms
%timeit forward_euler_vector() #248ms

这意味着，使用 numpy 优化矩阵乘法的收益相对较小，至少在当前的问题上是这样的。

你也可以考虑直接去掉函数 f，把它的操作放在 for 循环中，这样可以消除调用的开销。不过，这个默认参数的技巧也可以应用于 scipy 更通用的时间积分器，在那里你必须提供一个函数 f。

编辑：正如 Jaime 指出的，另一种方法是把 static_tmp 当作函数 f 的一个属性，在声明函数后但在调用之前创建它：

 def f(x): 
    f.static_tmp[0]=x[1]
    f.static_tmp[1]=-x[0] - a * x[1] * (x[0] ** 2 + x[1] ** 2 - 1)
    return f.static_tmp
 f.static_tmp=empty((2,2))