Page - 299 - in Programming for Computations – Python - A Gentle Introduction to Numerical Simulations with Python 3.6, Volume Second Edition

9.2 FiniteDifferenceMethods 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 computingthe right-handside: 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 loopcan be replacedby a vectorizedexpressionwith 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 coverall the inneruvaluesdisplaced one index to the right (compared to 1:N-1), i.e.,u[2:N].Similarly,u[i-1]correspondstoall inneruvaluesdisplaced one index to the left: u[0:N-2]. Finally, u[i] has the same indices as rhs: u[1:N-1].Thevectorized loopcan thereforebewritten in termsof 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 thefilerod_FE_vec.py. 9.2.6 UsingOdespytoSolvetheSystemofODEs LetusnowshowhowtoapplyageneralODEpackagelikeOdespy(seeSect.8.4.6) tosolveourdiffusionproblem.Aslongaswehavedefinedaright-handsidefunction rhs this isverystraightforward: 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, likeRK2 andRK4, 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 adaptivemethods. We can then specify a muchlargertimestepintime_points,andthesolverwillfigureouttheappropriate