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

234 8 SolvingOrdinaryDifferentialEquations Thelist version looksabitnicer, so that iswhyweprefera list andrather introduce f_ = lambda u, t: asarray(f(u,t)) in thegeneralode_FE function. Wecannowshowafunctionthat runsthepreviousSIRexample,whileusing the genericode_FE function: def demo_SIR(): """Test case using a SIR model.""" def f(u, t): S, I, R = u return [-beta*S*I, beta*S*I - gamma*I, gamma*I] beta = 10./(40*8*24) gamma = 3./(15*24) dt = 0.1 # 6 min D = 30 # Simulate for D days N_t = int(D*24/dt) # Corresponding no of time steps T = dt*N_t # End time U_0 = [50, 1, 0] u, t = ode_FE(f, U_0, dt, T) S = u[:,0] I = u[:,1] R = u[:,2] fig = plt.figure() l1, l2, l3 = plt.plot(t, S, t, I, t, R) fig.legend((l1, l2, l3), (’S’, ’I’, ’R’), ’center right’) plt.xlabel(’hours’) plt.show() # Consistency check: N = S[0] + I[0] + R[0] eps = 1E-12 # Tolerance for comparing real numbers for n in range(len(S)): SIR_sum = S[n] + I[n] + R[n] if abs(SIR_sum - N) > eps: print(’*** consistency check failed: S+I+R={:g} != {:g}’\ .format(SIR_sum, N)) if __name__ == ’__main__’: demo_SIR() Recall that the u returned from ode_FE contains all components (S, I, R) in the solutionvectoratall timepoints.Wethereforeneedtoextract theS,I, andRvalues in separatearraysfor furtheranalysis andeasy plotting. Another key feature of this higher-quality code is the consistency check. By adding the three differential equations in the SIR model, we realize that S′ + I′ +R′ = 0, which means that S + I +R = const. We can check that this relation holds by comparing Sn + In +Rn to the sum of the initial conditions. Exercise 8.6 suggests another method for controlling the quality of the numerical solution.