Vectorization of the direct Euler method for a system of differential equations

Question

I numerically solve for x (t) for a system of differential equations of the first order. System:

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

I applied the Forward Euler method to solve this problem as follows:

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 

Now I want to further vectorize the code and store x and y in one array, I came up with the following solution:

 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 

Question: forward_euler_vector () works, but was this the best way to vectorize it? I ask because the vectorized version runs 20 ms slower on my laptop:

 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

Comment @Ophion very well explains what is happening. Calling array() inside f(x) introduces some overhead, which reduces the use of matrix multiplication in the expression h * f(x[step]) .

And, as he says, you might be interested in scipy.integrate look at scipy.integrate for a good set of numerical integrators.

To solve the problem by vectorizing your code, you want to avoid re-creating the array every time you call f . You want to initialize the array once and return it with every call. It looks like the static variable is in C / C ++.

You can achieve this with a mutable default argument, which is interpreted once during the definition of the function f(x) and has a local scope. Since it needs to be changed, you encapsulate it in a list of one element:

  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] 

With this modification of your code, the overhead of creating an array disappears, and on my machine I get a slight improvement:

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

This means that the gain in optimizing matrix multiplication using numpy is pretty small, at least on the problem.

You can immediately get rid of the function f by doing your operations in a for loop, getting rid of the overhead of the call. However, this default argument trick can also be applied to scipy more general time integrators, where you must provide the function f .

EDIT: as Jaime noted, another way is to treat static_tmp as an attribute of function f and create it after declaring the function, but before calling it:

  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)) 

Answer 2

There is always a trivial autojit solution:

 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 

Timings:

 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