Page - 118 - in Programming for Computations – Python - A Gentle Introduction to Numerical Simulations with Python

118 4 SolvingOrdinaryDifferentialEquations # Ensure that any list/tuple returned from f_ is wrapped as array f_ = lambda u, t: asarray(f(u, t)) u = zeros((N_t+1, len(U_0))) t = 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 ...needsanexplanation. For auser,who justneeds to define thef in theODEsystem, it is convenient to insert thevariousmathematical expressions on the right-hand sides in a list and return that list. Obviously, we could demand the user to convert the list to anumpy array, but it is so easy to do a general such conversion in the ode_FE function aswell. Tomake sure that the result fromf is indeedanarray thatcanbeused forarraycomputingin the formula u[n] + dt*f(u[n], t[n]),we introduceanewfunctionf_ that calls theuser’s fandsends theresults throughthenumpy functionasarray,whichensures that its argument is converted toanumpyarray (if it is notalreadyanarray). Notealso theextraparenthesis requiredwhencallingzeroswith two indices. Let us showhow the previousSIRmodel can be solved using the newgeneral ode_FE thatcansolveanyvectorODE.Theuser’sf(u, t) functiontakesavector u,with threecomponentscorrespondingtoS,I, andRasargument,alongwith the currenttimepointt[n],andmustreturnthevaluesoftheformulasoftheright-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] Notethat theS,I, andRvaluescorrespondtoSn,In, andRn. Thesevaluesare then just inserted in thevarious formulas in thevectorODE.Herewecollect thevalues ina list since theode_FE functionwill anywaywrap this list inanarray.Wecould, ofcourse, returnedanarray instead: def f(u, t): S, I, R = u return array([-beta*S*I, beta*S*I - gamma*I, gamma*I]) Thelist version looksabitnicer, so that iswhyweprefera list and rather introduce f_ = lambda u, t: asarray(f(u,t)) in thegeneralode_FE function. We can now show a function that runs the previous SIR example, while using thegenericode_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]