New pymc/theano user. The following is a very simplified version of the actual model I’m trying to implement. Its an SIRD model where the spread of infections is moderated by the amount of deaths observed over some previous span of time. We want to solve for the disease parameters (spread, recovery rate and fatality rate).
I know with a regular SIR model I could use pymc3.ode, but as far as I know I can’t use that with this model because I can’t access the previous states of the system (i.e. the lags of D for the memory) while its running.
A for loop slowed my model to a grinding halt so I tried coding up a loop as a theano.scan. I haven’t actually figured out how to do the “memory” bit as a theano.scan yet but my bigger concern is it seems that the model parameters cannot be an input to the scan. Is this the case or is there something I’m misunderstanding?
There doesn’t seem to be a way to avoid a loop to compute the next state of the epidemic dynamics, is there any way to use pymc3 for a model that loops like this?
import numpy as np
import pymc3 as pm
import theano
import theano.tensor as T
import arviz as az
#simple SIRD model with a "memory" feature
def sirdMemory(N ,beta, gamma, pi, kappa, D):
#beta is a contact rate, gamma is a recovery rate and pi is case fatality rate
S_old=N-1
I_old=1
#we don't actually care about R, we focus on D as a flow not a stock
for k in range(len(D)):
D[k]=pi*gamma*I_old
S=-beta/N*S_old*I_old*(kappa/(kappa+np.sum(D[max(1,k-14):k])))+S_old
I=beta/N*S_old*I_old*(kappa/(kappa+np.sum(D[max(1,k-14):k])))-gamma*I_old+I_old
S_old=S
I_old=I
return D
N=10000
beta=2
gamma=1.5
pi=0.02
kappa=100
D=np.zeros(50)
deaths=sirdMemory(N,beta,gamma,pi,kappa,D)
#running an sir(d) function from theano.scan
#running an sir(d) function from theano.scan
#running an sir(d) function from theano.scan
#running a simpler sird model with theano.scan
#-----------------------------------------------------------------------------
t=T.iscalar('t')
N_in=T.iscalar('N')
S=T.scalar('S')
I=T.scalar('I')
beta_in=T.scalar('beta')
gamma_in=T.scalar('gamma')
pi_in=T.scalar('pi')
def innersir2(s,i,beta,gamma, pi, n):
sflow=-beta/n*s*i
iflow=beta*s*i/n-gamma*i
d=pi*gamma*i
return [s+sflow, i+iflow,d]
results, updates=theano.scan(fn=innersir2, outputs_info=[S,I,None],
non_sequences=[beta_in, gamma_in, pi_in,N_in], n_steps=t)
final_result = results[-1]
sirloop2 = theano.function(inputs=[S, I,beta_in,gamma_in,pi_in,N_in,t], outputs=final_result, updates=updates)
print(sirloop2(N-1,1,beta,gamma,pi,N,len(D)))
#Trying theano.scan in the model section
#-----------------------------------------------------------------------------
#define theano variables
t_sym=theano.tensor.iscalar('t_syn')
S_sym=theano.tensor.scalar('S_sym')
I_sym=theano.tensor.scalar('I_sym')
N_sym=theano.tensor.iscalar('N_sym')
beta_sym=theano.tensor.scalar('beta_sym')
gamma_sym=theano.tensor.scalar('gamma_sym')
pi_sym=theano.tensor.scalar('pi_sym')
with pm.Model() as model6:
theano.config.compute_test_value = "ignore"
beta = pm.Normal('lam', 2, 0.2)
gamma = pm.Normal('gamma', 0.8, 0.1)
pi = pm.Normal('pi', 0.02, 0.005)
results, updates=theano.scan(fn=innersir2, outputs_info=[S_sym,I_sym,None],
non_sequences=[beta_sym, gamma_sym, pi_sym,N_sym], n_steps=t_sym)
final_result = results[-1]
sirloop2 = theano.function(inputs=[S_sym, I_sym,beta_sym,gamma_sym,pi_sym,N_sym,t_sym], outputs=final_result, updates=updates)
dhat=sirloop2(N-1.0,1.0,beta,gamma,pi,N,len(deaths[1:]))
Y=pm.Poisson('Y', mu=dhat, observed=deaths[1:])
trace = pm.sample(10000, tune=1000, cores=1, chains=2, progressbar=True, start={'lam':1,'pi':0.03, 'gamma': 0.7}, step=pm.Metropolis())
data = az.from_pymc3(trace=trace)