Instantly share code, notes, and snippets.

# elofgren/gist:1390355 Created Nov 24, 2011

What would you like to do?
SIR with Random Parameters
 #Python implementation of continuous SIR model #Transmission parameter values are random #Import Necessary Modules import numpy as np import scipy.integrate as spi #Solving the differential equation. Solves over t for initial conditions PopIn def eq_system(startState,t): '''Defining SIR System of Equations''' beta = startState[3] gamma = startState[4] #Creating an array of equations Eqs= np.zeros((3)) Eqs[0]= -beta * startState[0]*startState[1] Eqs[1]= beta * startState[0]*startState[1] - gamma*startState[1] Eqs[2]= gamma*startState[1] return Eqs def model_solve(t): '''Stores all model parameters, runs ODE solver and tracks results''' #Setting up how long the simulation will run t_start = 1 t_end = t t_step = 0.02 t_interval = np.arange(t_start,t_end, t_step) n_steps = (t_end-t_start)/t_step #Setting up initial population state #n_params is the number of parameters (beta and gamma in this case) S0 = 0.99999 I0 = 0.00001 R0 = 0.0 beta = np.random.random_sample()+0.5 gamma = np.random.random_sample() n_params = 2 startPop = (S0, I0, R0, beta, gamma) #Create an array the size of the ODE solver output with the parameter values params = np.zeros((n_steps,n_params)) params[:,0] = beta params[:,1] = gamma timer = np.arange(n_steps).reshape(n_steps,1) SIR = spi.odeint(eq_system, startPop, t_interval) #Glue together ODE model output and parameter values in one big array output = np.hstack((timer,SIR,params)) return output testrun = model_solve(100) print testrun