矢量化前向欧拉方法的微分方程组

时间:2014-05-15 13:14:47

标签: python numpy vectorization

我在数值上求解x(t)的一阶微分方程组。该系统是:

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

我已经实现了Forward Euler方法来解决这个问题,如下所示:

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

2 个答案:

答案 0 :(得分:3)

@Ophion评论很好地解释了发生了什么。在array()内对f(x)的调用会引入一些开销,这会在表达式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的一个属性,并在声明函数之后但在调用函数之前创建它: / p>
 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))

答案 1 :(得分:3)

总有一些简单的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