SIR with Random Parameters

gistfile1.py

#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