9.2 Finite Difference Methods
299
9.2.5 Vectorization
Occasionally in this book, we show how to speed up code by replacing loops
over arrays by vectorized expressions. The present problem involves a loop for
computing the right-hand side:
for i in range(1, N):
rhs[i] = (beta/dx**2)*(u[i+1] - 2*u[i] + u[i-1]) + g(x[i], t)
This loop can be replaced by a vectorized expression with the following reasoning.
We want to set all the inner points at once: rhs[1:N-1] (this goes from index 1
up to, but not including, N). As the loop index i runs from 1 to N-1, the u[i+1]
term will cover all the inner u values displaced one index to the right (compared to
1:N-1), i.e., u[2:N]. Similarly, u[i-1] corresponds to all inner u values displaced
one index to the left: u[0:N-2]. Finally, u[i] has the same indices as rhs:
u[1:N-1]. The vectorized loop can therefore be written in terms of slices:
rhs[1:N-1] = (beta/dx**2)*(u[2:N+1] - 2*u[1:N] + u[0:N-1]) +
g(x[1:N], t)
This rewrite speeds up the code by about a factor of 10. A complete code is found
in the file rod_FE_vec.py.
9.2.6 Using Odespy to Solve the System of ODEs
Let us now show how to apply a general ODE package like Odespy (see Sect. 8.4.6)
to solve our diffusion problem. As long as we have defined a right-hand side function
rhs this is very straightforward:
import odespy
import numpy as np
solver = odespy.RKFehlberg(rhs)
solver.set_initial_condition(U_0)
T = 1.2
N_t = int(round(T/dt))
time_points = np.linspace(0, T, N_t+1)
u, t = solver.solve(time_points)
# Check how many time steps are required by adaptive vs
# fixed-step methods
if hasattr(solver, ’t_all’):
print(’# time steps:’, len(solver.t_all))
else:
print(’# time steps:’, len(t))
The very nice thing is that we can now easily experiment with many different
integration methods. Trying out some simple ones first, like RK2 and RK4, quickly
reveals that the time step limitation of the Forward Euler scheme also applies
to these more sophisticated Runge-Kutta methods, but their accuracy is better.
However, the Odespy package offers also adaptive methods. We can then specify a
much larger time step in time_points, and the solver will figure out the appropriate
299
9.2.5 Vectorization
Occasionally in this book, we show how to speed up code by replacing loops
over arrays by vectorized expressions. The present problem involves a loop for
computing the right-hand side:
for i in range(1, N):
rhs[i] = (beta/dx**2)*(u[i+1] - 2*u[i] + u[i-1]) + g(x[i], t)
This loop can be replaced by a vectorized expression with the following reasoning.
We want to set all the inner points at once: rhs[1:N-1] (this goes from index 1
up to, but not including, N). As the loop index i runs from 1 to N-1, the u[i+1]
term will cover all the inner u values displaced one index to the right (compared to
1:N-1), i.e., u[2:N]. Similarly, u[i-1] corresponds to all inner u values displaced
one index to the left: u[0:N-2]. Finally, u[i] has the same indices as rhs:
u[1:N-1]. The vectorized loop can therefore be written in terms of slices:
rhs[1:N-1] = (beta/dx**2)*(u[2:N+1] - 2*u[1:N] + u[0:N-1]) +
g(x[1:N], t)
This rewrite speeds up the code by about a factor of 10. A complete code is found
in the file rod_FE_vec.py.
9.2.6 Using Odespy to Solve the System of ODEs
Let us now show how to apply a general ODE package like Odespy (see Sect. 8.4.6)
to solve our diffusion problem. As long as we have defined a right-hand side function
rhs this is very straightforward:
import odespy
import numpy as np
solver = odespy.RKFehlberg(rhs)
solver.set_initial_condition(U_0)
T = 1.2
N_t = int(round(T/dt))
time_points = np.linspace(0, T, N_t+1)
u, t = solver.solve(time_points)
# Check how many time steps are required by adaptive vs
# fixed-step methods
if hasattr(solver, ’t_all’):
print(’# time steps:’, len(solver.t_all))
else:
print(’# time steps:’, len(t))
The very nice thing is that we can now easily experiment with many different
integration methods. Trying out some simple ones first, like RK2 and RK4, quickly
reveals that the time step limitation of the Forward Euler scheme also applies
to these more sophisticated Runge-Kutta methods, but their accuracy is better.
However, the Odespy package offers also adaptive methods. We can then specify a
much larger time step in time_points, and the solver will figure out the appropriate