Create a gist now

Instantly share code, notes, and snippets.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment