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

8.3 SpreadingofDisease:ASystemofFirstOrderODEs 233 For all this to work, the complete numerical solution must be represented by a two-dimensional array, created byu = zeros((N_t+1, m+1)). The first index counts the time points and the second the componentsof the solution vector at one timepoint.That is,u[n,i]correspondsto themathematicalquantityuni . Whenwe use only one index, as in u[n], this is the same as u[n,:] and picks out all the components in the solution at the time point with index n. Then the assignment u[n+1] = ... becomes correct because it is actually an in-place assignment u[n+1, :] = ....Thenice featureof these facts is that thesamepieceofPython codeworks forbothascalar ODEandasystem ofODEs! Theode_FEfunctionforthevectorODEisplacedinthefileode_system_FE.py andwaswrittenas follows: import numpy as np import matplotlib.pyplot as plt def ode_FE(f, U_0, dt, T): N_t = int(round(T/dt)) # Ensure that any list/tuple returned from f_ is wrapped as array f_ = lambda u, t: np.asarray(f(u, t)) u = np.zeros((N_t+1, len(U_0))) t = np.linspace(0, N_t*dt, len(u)) u[0] = U_0 for n in range(N_t): u[n+1] = u[n] + dt*f_(u[n], t[n]) return u, t The linef_ = lambda ...needsan explanation.For a user,who just needs to define thef in the ODE system, it is convenient to insert the variousmathematical expressions on the right-hand sides in a list and return that list. Obviously, we could demand the user to convert the list to a numpy array, but it is so easy to do a general such conversion in the ode_FE function as well. To make sure that the result fromf is indeedanarray thatcanbeused forarraycomputing in the formula u[n] + dt*f(u[n], t[n]),we introducea new functionf_ that calls the user’s fandsends the results throughthenumpy functionasarray,whichensures that its argument is convertedtoanumpyarray(if it isnotalreadyanarray). Notealso theextraparenthesis requiredwhencallingzeroswith two indices. Let us show how the previous SIR model can be solved using the new general ode_FE thatcansolveanyvectorODE.Theuser’sf(u, t) functiontakesavector u,with threecomponentscorrespondingtoS, I, andR asargument,alongwith the currenttimepointt[n],andmustreturnthevaluesoftheformulasof theright-hand sides in thevectorODE.Anappropriate implementation is def f(u, t): S, I, R = u return [-beta*S*I, beta*S*I - gamma*I, gamma*I] Note that theS,I, andRvaluescorrespondtoSn,In, andRn.Thesevaluesare then just inserted in the various formulas in the vector ODE. Here we collect the values ina list since theode_FE functionwill anywaywrap this list inan array.We could, ofcourse, returnedanarray instead: def f(u, t): S, I, R = u return array([-beta*S*I, beta*S*I - gamma*I, gamma*I])