新型肺炎COVID-19の感染推移を時間依存パラメータを含むSEIRSモデルで推定しようと思って断念していたが出来るかもしれない

R Rpackage Rを使いこなす rstan 数学 数理モデル 統計 感染症学

諦めていた。
mikuhatsune.hatenadiary.com
というのも、SEIRSモデルでS→Eになるパラメータ\beta が、時間依存的なパラメータとして
\beta(t)=\gamma R_0 \left[\frac{f}{2}\cos(\frac{2\pi}{52}(t-\phi))-(1-\frac{f}{2})\right]
で定義されるが、これはrstanではできなさそうである。というのが、integrate_ode関数が時間依存のパラメータをどう頑張っても取れないみたいだからである。
ODEs with time dependent parameters - Modeling - The Stan Forums

というわけでstanの神が助言してくれた。


S/E/I/R の各コンパートメントは、ある時刻t においてすべて足せば和は1になるので、simplexを必要な時刻数分だけリストか何かで持っておいて、transformed parametersブロック内でぶん回せばよさそうである。
まずは普通に\frac{dy}{dt}を地味に計算してSEIRモデルをシミュレーションした。
f:id:MikuHatsune:20200421183612p:plain

さて、rstan でtransformed parametersブロックにsimplexを各時刻に定義してリストとして持ち、上記式がf\sim U(0.18, 0.25), \phi\sim U(1.5, 2), R_0\sim U(2, 3) でサンプリングされて\beta(t)が決まる、としてコードを実行した。
推定区間がついたSEIRの推移が推定できた。
f:id:MikuHatsune:20200421183631p:plain

times <- 0:200
f <- 0.21
phi <- 2
R0 <- 2.2
Rt <- R0*(f/2*cos(2*pi/52*(times-phi)) + (1-f/2))
plot(times, Rt)

g <- 1.2    # gamma
d <- 0.3    # delta
b <- g * Rt # beta
e <- 0.4   # eta

y0 <- c(0.999, 0.001, 0, 0)
names(y0) <- c("S", "E", "I", "R")
res <- rbind.data.frame(y0, matrix(0, nr=length(times)-1, nc=4))
colnames(res) <- c("S", "E", "I", "R")
for(i in 2:length(times)){
  dS <- -b[i-1]*res$S[i-1]*res$I[i-1] + e*res$R[i-1]
  dE <-  b[i-1]*res$S[i-1]*res$I[i-1] - d*res$E[i-1]
  dI <-  d*res$E[i-1] - g*res$I[i-1]
  dR <-  g*res$I[i-1] - e*res$R[i-1]
  res[i, ] <- res[i-1,] + c(dS, dE, dI, dR)
}
cols <- c("black", "orange", "red", "lightgreen")
matplot(res, type="l", lty=1, lwd=3, col=cols, xlab="Time", ylab="Proportion")
legend("topright", legend=colnames(res), pch=15, col=cols)

# simplex でやってみる
code <- "
data {
   int<lower=1> T; // 時間の数
   real TS[T]; // 各時刻
   real Y0[4]; // Y の初期値 SEIR
   real b[T];
   real d;
   real g;
   real e;
}
parameters {
  real a; // dummy
}
transformed parameters {
   simplex[4] y_[T];
   y_[1] = to_vector(Y0);
   for(i in 2:T){
     y_[i][1] = y_[i-1][1] + -b[i-1]*y_[i-1][1]*y_[i-1][3] + e*y_[i-1][4];
     y_[i][2] = y_[i-1][2] +  b[i-1]*y_[i-1][1]*y_[i-1][3] - d*y_[i-1][2];
     y_[i][3] = y_[i-1][3] +  d*y_[i-1][2] - g*y_[i-1][3];
     y_[i][4] = y_[i-1][4] +  g*y_[i-1][3] - e*y_[i-1][4];
   }
}
model {
  a ~ uniform(0, 1); // 何もしない
}
"
m0 <- stan_model(model_code=code)
standata <- list(T=length(times), TS=times, Y0=y0, b=b, d=d, g=g, e=e)
fit <- sampling(m0, standata, iter=200, warmup=100, chain=2)
ex <- extract(fit, pars=head(fit@model_pars, -1))

# beta をb(t) にする
code <- "
data {
   int<lower=1> T; // 時間の数
   real TS[T]; // 各時刻
   real Y0[4]; // Y の初期値 SEIR
   // real b[T];
   real d;
   real g;
   real e;
   //real b01[2]; // beta の範囲
   real p01[2]; // phi  の範囲
   real f01[2]; // f の範囲
   real R01[2]; // R0 の範囲
}
parameters {
   real<lower=p01[1], upper=p01[2]> phi;
   real<lower=f01[1], upper=f01[2]> f;
   real<lower=R01[1], upper=R01[2]> R0;
}
transformed parameters {
   simplex[4] y_[T];
   real<lower=0> b[T];
   y_[1] = to_vector(Y0);
   for(i in 1:T){
     b[i] = g * R0 * (f/2*cos(2*pi()/52*(TS[i]-phi)) + (1-f/2));
   }
   // b = g * Rt;
   for(i in 2:T){
     y_[i][1] = y_[i-1][1] + -b[i-1]*y_[i-1][1]*y_[i-1][3] + e*y_[i-1][4];
     y_[i][2] = y_[i-1][2] +  b[i-1]*y_[i-1][1]*y_[i-1][3] - d*y_[i-1][2];
     y_[i][3] = y_[i-1][3] +  d*y_[i-1][2] - g*y_[i-1][3];
     y_[i][4] = y_[i-1][4] +  g*y_[i-1][3] - e*y_[i-1][4];
   }
}
model {
   f   ~ uniform(f01[1], f01[2]);
   phi ~ uniform(p01[1], p01[2]);
   R0  ~ uniform(R01[1], R01[2]);
}
"

m0 <- stan_model(model_code=code)
standata <- list(T=length(times), TS=times, Y0=c(0.999, 0.001, 0, 0), d=d, g=g, e=e,
                 f01=c(0.18, 0.25), p01=c(1.5, 2.4), R01=c(2, 3))
fit <- sampling(m0, standata, iter=1000, warmup=500, chain=4)
ex <- extract(fit, pars=head(fit@model_pars, -1))

alpha <- 0.05
cia <- c(alpha/2, 0.5, 1-alpha/2)
m <- abind::abind(mapply(function(z) cbind(c(0, y0[z], 0), apply(ex$y_[,,z], 2, quantile, cia)), seq(y0), SIMPLIFY=FALSE), along=3)

parms <- apply(ex$b, 2, quantile, cia)

lab <- c("S", "E", "I", "R")
# 透明にしながら実線も描けるようにする
cols256 <- list(S=c(0, 255, 0), E=c(255, 202, 5), I=c(255, 0, 0), R=c(0, 255, 255))
cols <- mapply(function(z) rgb(z[1], z[2], z[3], maxColorValue=256), cols256)
colsa <- mapply(function(z) rgb(z[1], z[2], z[3], 20, maxColorValue=256), cols256)

par(mar=c(5, 5, 3, 5), cex.lab=1.5, las=1)
matplot(m[2,,], type="n", col=colsa, lty=1, lwd=3, xlab="Time (days)", ylab="Proportion of S/E/I/R patients")
#legend(0, N, names(y0), ncol=length(y0), col=cols, pch=15, yjust=-0.1, xpd=TRUE, bty="n", cex=2)
legend(par()$usr[2], 1, names(y0), col=cols, pch=15, xpd=TRUE, bty="n", cex=2)
for(i in seq(dim(m)[3])){
  lines(c(0, times), m[2,,i], col=cols[i], lwd=4)
  polygon(c(c(0, times), rev(c(0, times))), c(m[1,,i], rev(m[3,,i])), col=colsa[i], border=NA)

}