新型肺炎COVID-19の日本の実効再生産数を推定したrstanのコードを解説してみる

西浦先生が日本の実効再生産数を推定した。
コードはrstanで下記から取れる。
https://nbviewer.jupyter.org/github/contactmodel/COVID19-Japan-Reff/tree/master/

解説動画を見逃したのでコードと関連論文からのお勉強になるが、肝としては、
・知りたいのは「感染した日」である。
・診断日もしくは報告日は、データを収集して統計を取っているのでわかる。
・診断されるには検査される必要があるから、だいたい症状か接触歴があって、発症日はそこそこデータがある。
・感染した瞬間、はもちろん発病（はほとんど）していないのでわからない。
という前提がある。PDFの「患者」の観測データについて、の項。
（誰からから感染させられる）ー感染日ー発症日ー診断日/報告日という一連の流れについて、まったく情報がないわけではなく、いままでの数理モデルの知見からは
・誰かから感染させられるー感染日：serial interval もしくはgeneration time と書いてある。厳密には違うようだが、使用しているパラメータはワイブル分布から同一のものとしている。PDFでは「世代時間<=明示的にまだ得られていない間、発病間隔と同じと想定」の項。
・感染日ー発症日：incubation period. 対数正規分布から平均5.6日くらいと適当に設定している（適当にというのはランダム、というわけではなく、過去の推定で当てはまりがよさそうな、という意味合いで、対数正規分布ではくワイブル分布やガンマ分布でもやれないことはない）。
・発症日ー報告日：Right-truncated reporting delay. 打ち切りを考慮したワイブル分布を考慮しているようで、尤度関数 $L$ は
$L(\theta|S_k,O_k,T)=\displaystyle\prod_{k=1}^K\frac{h(S_k-O_k|\theta)}{H(T-O_k|\theta)}$
$S_k$ ： $k$ さんの報告日
$O_k$ ： $k$ さんの発症日
$T$ ：カレンダー上の直近の観測日
$h$ ：ワイブル分布の確率密度（PDF）
$H$ ：ワイブル分布の累積確率密度（CDF）
となる。

さて、たぶんこれが一番大事な話だと思うが、感染日ー発症日、そして発症日ー診断日/報告日というのをなんとか推定するのに、backprojection というものを利用している。報告日はそれなりにデータがあるので、発症日←診断日/報告日のbackprojection をして発症日を推定し、感染日←発症日を再度して感染日を推定している。rstan のコードではこの部分がsurveillance というパッケージからbackprojNPという関数を使って最尤推定法で求めた点推定値を流用して実効再生産数を求めている。backprojection についての説明は皆無なので、むしろこれを理解するのに時間がかかった。

A method of non-parametric back-projection and its application to AIDS data. - PubMed - NCBI
今回のCOVID-19のように、診断日ははっきりしているが感染日がわからない疾患として、HIVを題材にしている。ある日付 $t, t=\{1,2.\dots T\}$ のHIVの診断数が $Y_t$ とすると、 $d$ の期間のincubation period は $f_d$ の確率（PDF）に従うとする。incubation period を考慮した真の推定値 $\lambda_t$ はポアソン分布に従うとして、尤度関数は
$\displaystyle\prod_{t=1}^T(\displaystyle\sum_{i=1}^t \lambda_i f_{t-i})^{y_t} \exp(-\displaystyle\sum_{i=1}^t \lambda_i f_{t-i})$
を最大化することで得られる、となっているらしい。 $\displaystyle\sum_{i=1}^t \lambda_i f_{t-i}$ の部分が一般的なポアソン分布のパラメータ $\lambda$ 部分だが $y_t!$ で割らなくていいのだろうか。
これはGLMで解くのが難しかったらしいので、EMアルゴリズムで解けるようにしたのが論文である。
$\lambda_t$ というのが日付 $t=\{1,2,\dots, T\}$ における真の（平均）感染数とすると、EステップとMステップを合体させた式が
$\lambda_t^{new}=\frac{\lambda_t^{old}}{F_{T-t}}\displaystyle\sum_{d=0}^{T-t}\frac{Y_{t+d} f_d}{\displaystyle\sum_{i=1}^{t+d} \lambda_t^{old} f_{t+d-i}}$
となる。ここで、 $F_{T-t}=\displaystyle\sum_{d=0}^{T-t}f_d$ である。 $\lambda_t$ が更新されたら差分でもなんでもとって適当な閾値 $\epsilon$ を下回れば終了する。
ただし、このまま $\lambda_t$ を使うと結構ギザギザなので、smoothing を行なっている。これがExpectation-Maximization-Smoothing (EMS) と書いてある所以のようで、上で求めたのを $\phi_t^{new}\gets \lambda_t^{new}$ にしておく。やり方は $\phi_t^{new}$ の前後での重み付け平均をしているだけで、重み $w_i={k \choose i}/2^k$ とする。 $k=2$ だと[0.25, 0.5, 0.25] となる。これを使って
$\lambda_t^{new}=\displaystyle\sum_{i=0}^k w_i \phi_{t+i-\frac{k}{2}}^{new}$ とする。 $k$ は偶数とし、ある $i$ を中心に前後の重み付け平均をとる。

surveillanceではsurveillance:::backprojNP.fitおよびsurveillance:::em.step.beckerでEMSを行なっている。普通にやると遅いので、RccpでC++を呼び出して行なっている。eq3aというmethod 名が論文での式3a（eq3a）である。
こんな感じになる。感染日と報告日は10-14日くらいずれていそうな感じがある。
f:id:MikuHatsune:20200513203350p:plain

入力するデータがそれなりに整ったので、後はrstanで実効再生産数 $R_t$ を推定する。平均的な時間 $t$ での感染者数は（domestic とimported でデータは分かれているが、簡便のため略）
$E(C)=R_t\displaystyle\sum_{\tau =0}^{t-1}C_{total}(t-\tau)g(\tau)\frac{F(T-t)}{F(T-t+\tau)}$
ここで
$C_{total}(t)$ ：感染者数の総計
$g$ ：generation interval の確率密度（PDF）
$F$ ：感染日から報告日までの遅れの累積確率（CDF）。これはincubation period のPDFと発症日ー報告日の遅れのPDFとの畳み込みになる（感染日→発症日→報告日）。
である。後はポアソン分布を仮定する。

ここまでくるとrstanのコードがだいぶわかるようになってくる。
冒頭に使用するCPUコア数をオプションで指定すれば、利用可能な最大CPUで並列してくれる。コードでは2chain しか使ってないが、収束が非常に速いので困りはしない。

library(rstan)
rstan_options(auto_write=TRUE)
options(mc.cores=parallel::detectCores())

functionsブロックでは、好きな関数を定義してコード内で使えるようにする。rstanは離散確率分布は別に定義しないと柔軟に扱えないことが多く、今回のようにある日（離散）での確率密度はPMFになるので、累積確率分布の差分を取って柔軟にrstan内で扱えるようにする。
generation interval にはワイブル分布、incubation period には対数正規分布なので、それぞれ定義する。

functions {
    /* discretesized version of lognormal distribution */
    vector plognormal(real mu, real sigma, int K) {
        vector[K] res; 
        for (k in 1:K)
            res[k] = exp(lognormal_lcdf(k | mu, sigma));

        return append_row(res[1], res[2:K]-res[1:(K-1)]);        
    }

    /* discretesized version of Weibull distribution */
    vector pweibull(real kappa, real theta, int K) {
        vector[K] res; 
        for (k in 1:K)
            res[k] = -expm1(-pow(1.0*k/theta, kappa));

        return append_row(res[1], res[2:K]-res[1:(K-1)]);  
    }

確率分布の畳み込みを行うのも関数を定義する。rstanにはreverse関数がないようなのでYrevというのを後のtransformed dataブロックで作っているが、X[1:k]とY[k:1]でやれば同じである。ただし、rstanは普通のRみたいにY[k:1]とやるのは受け付けてくれなかった。

/* calculating the convolutions */
// X: first function, Yrev: reversed version of the second function
// K: length of X and Yrev
// the result is the vector of length K-1, because the first element equal zero is omitted
vector convolution(vector X, vector Yrev, int K) {
    vector[K-1] res;
    res[1] = X[1]*Yrev[K];
    for (k in 2:K-1) 
        res[k] = dot_product(head(X, k), tail(Yrev, k));

    return res;        
}

// Special form of the convolution with adjusting to the delay
// F: cumulative distribution function of the reporting delay
vector convolution_with_delay(vector X, vector Yrev, vector F, int K) {
    vector[K-1] res;
    vector[K] Z = X ./ F;

    res[1] = F[2]*Z[1]*Yrev[K];
    for (k in 3:K) 
        res[k-1] = F[k]*dot_product(head(Z, k-1), tail(Yrev, k-1));

    return res;

modelブロック、 $R_t$ が2.4 くらいにいるだろうという事前分布はいいとして、scaleパラメータが1のガンマ分布からサンプリングしたのは単にRt .* conv が期待値になるようにしたかったからなんでしょうか。

model {
    Rt ~ normal(2.4, 2.0);
    Rt_adj ~ normal(2.4, 2.0);
    
    target += gamma_lpdf(tail(domestic_backproj, K-1) | Rt .* conv + 1e-13, 1.0)
            + gamma_lpdf(tail(domestic_backproj, K-1) | Rt_adj .* conv_delay_adj + 1e-13, 1.0);
}

backprojection はEMSでなくてもベイズ的に推定できるはずなので、これもrstan に組み込めばよさそうだが、まだこれができていないので保留。