Initial commit

python/gillespieSim.py

+# 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
+