Commit f29e4910 authored by Chris Jewell's avatar Chris Jewell
Browse files

Initial commit

parents
# Chris's Doob Gillespie simulation code
import numpy as np
import scipy.spatial.distance as dist
import matplotlib.pyplot as plt
class LisaKernel:
def __init__(self, coords, X, c):
self.D = dist.squareform(dist.pdist(np.array(coords)))-c
self.X = np.array(X)
def __call__(self,params):
K = np.exp(params['beta'] + self.X) * np.exp(-params['kappa']*self.D)
return K
def shape(self):
return self.D.shape
class Gillespie:
def __init__(self, kernel):
"""Simulate an continuous time spatial SIR epidemic starting at individual I1
Parameters
==========
coords -- an nx2 np.array of coordinates
X -- a matrix of individual-level covariates
"""
# Set up distance matrix and kernel
self.kernel = kernel
def simulate(self, I1, params):
"""Runs a simulation starting at individual I1.
Parameters
==========
I1 -- 0-based index of initial infective
params -- a dict of parameters to pass to the kernel
Returns
=======
a nx2 np.array containing infection and removal times.
"""
# Calculate the full nxn infection kernel matrix (i.e. i to j)
# we only need to do this at the beginning.
# N.B. this won't scale to large populations!
kernel = self.kernel(params)
t = 0.
# nx2
eventTimes = np.full(self.kernel.shape(), np.inf)
eventTimes[I1,0] = t
idxS = np.where(eventTimes[:,0]>t)[0].tolist()
idxI = [I1]
gamma = params['gamma']
while(len(idxI) > 0):
# Infections
iRate = 0.
if len(idxS)>0:
# Calculate infec pressure
press = kernel[np.ix_(idxI,idxS)].sum(axis=0)
iRate = press.sum()
# Calculate removal rate
rRate = gamma * len(idxI)
# Time of next event
t += np.random.exponential(size=1, scale=1./(iRate+rRate))
# Infec or remove?
isInfec = (iRate / (iRate+rRate)) > np.random.uniform(size=1)
if isInfec[0]==True: # Infection
i = np.asscalar(np.nonzero(np.random.multinomial(n=1, pvals=press/iRate))[0][0])
idx = idxS[i]
del idxS[i]
idxI.append(idx)
eventTimes[idx,0] = t
else:
i = np.asscalar(np.random.randint(low=0,high=len(idxI),size=1))
idx = idxI[i]
del idxI[i]
eventTimes[idx,1] = t
return eventTimes
def getSimTimeSeries(eventTimes,res=200):
minT = np.min(eventTimes)
maxT = np.max(eventTimes[eventTimes[:,1]<np.inf,1])
T = np.linspace(minT,maxT+(maxT-minT)/res,res)
nIR = np.array([[np.sum((eventTimes[:,0]<=t) & (t<eventTimes[:,1])),
np.sum(eventTimes[:,1] < t)] for t in T])
nS = np.full(T.shape[0], eventTimes.shape[0]) - nIR.sum(axis=1)
return np.column_stack((T,nS,nIR))
def plotSim(coords, eventTimes):
coords = np.array(coords)
suscep = coords[eventTimes[:,0]==np.inf,:]
infec = coords[eventTimes[:,0]<np.inf,:]
f,ax = plt.subplots(1,2,figsize=(8,4))
ax[0].plot(suscep[:,0],suscep[:,1], 'bo')
ax[0].plot(infec[:,0],infec[:,1], 'ro')
ts = getSimTimeSeries(eventTimes)
ax[1].plot(ts[:,0],ts[:,1],'b-')
ax[1].plot(ts[:,0],ts[:,2],'r-')
ax[1].plot(ts[:,0],ts[:,3],'g-')
return ax
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment