インベンスルービン本の25章に記載されている，ベイズ的な操作変数法のPythonでの実装を簡単に解説します．使っているパッケージは基本的にはNumpyのみです．

問題設定やモデル，記法の詳細は本を参照してください．コードはGitHubに載せています．

データセットは長島先生のツイートを参考にhttps://mollyaredtana.wixsite.com/sta108 から取得しました．

※本に記載されている結果が一部再現できていないのですが，松浦先生のStanを使った再現とは結果がほぼあっているため，おそらく教科書に載っている結果が一部間違っているのではないかと考えています．

【目次】

実装の詳細
1. 遵守タイプのサンプリング
2. パラメータのサンプリング
因果効果の推定
サンプリング結果（更新予定）
雑感

1. 実装の詳細

遵守タイプとパラメータを（ブロック）ギブスサンプリング法によりサンプリングします．実装は大きく以下の2つのパートに分かれます．

遵守タイプ$G$のサンプリング：$p(G|Y, \theta)$からのサンプリング
パラメータ$\theta$のサンプリング：$p(\theta | Y, G)$からのサンプリング

1.1 遵守タイプのサンプリング

遵守タイプのサンプリングは以下の4つのパターンに分かれます．

$Z_i = 0$かつ$W_i = 1$のケース
$Z_i = 1$かつ$W_i = 0$のケース
$Z_i = 0$かつ$W_i = 0$のケース
$Z_i = 1$かつ$W_i = 1$のケース

$Z_i = 0$かつ$W_i = 1$のケース

割当反抗者は存在しないことを仮定しているので，$Z_i = 0$かつ$W_i = 1$のケースは常時処置受入者になります．

$Z_i = 1$かつ$W_i = 0$のケース

上記と同様に，割当反抗者は存在しないことを仮定しているので，$Z_i = 1$かつ$W_i = 0$のケースは常時処置拒否者になります．

$Z_i = 0$かつ$W_i = 0$のケース

$Z_i = 0$かつ$W_i = 0$のケースは割当遵守者(co)か常時処置拒否者(nt)の可能性があります．$Y_i$を観測したあとの患者$i$が割当遵守者である事後確率は以下です．

$$
p(G_i=co | Y_i, \theta) = \frac{p(G_i=co | \gamma_{at}, \gamma_{nt}) p(Y_i | G_i=co, \beta_{co,c})}{p(G_i=co | \gamma_{at}, \gamma_{nt}) p(Y_i | G_i=co, \beta_{co,c}) + p(G_i=nt | \gamma_{at}, \gamma_{nt}) p(Y_i | G_i=nt, \beta_{nt})}
$$

上記の確率で遵守タイプは割当遵守者となり，そうでない場合は常時処置拒否者になります．

$Z_i = 1$かつ$W_i = 1$のケース

$Z_i = 1$かつ$W_i = 1$のケースは割当遵守者(co)か常時処置受入者(at)の可能性があります．$Y_i$を観測したあとの患者$i$が割当遵守者である事後確率は以下です．

$$
p(G_i=co | Y_i, \theta) = \frac{p(G_i=co | \gamma_{at}, \gamma_{nt}) p(Y_i | G_i=co, \beta_{co, t})}{p(G_i=co | \gamma_{at}, \gamma_{nt}) p(Y_i | G_i=co, \beta_{co, t}) + p(G_i=at | \gamma_{at}, \gamma_{nt}) p(Y_i | G_i=nt, \beta_{at})}
$$

上記の確率で遵守タイプは割当遵守者となり，そうでない場合は常時処置受入者になります．

コード

遵守タイプのサンプリングのコードは以下です．$G$が遵守タイプの配列で，常時処置受入者は$0$，常時処置拒否者は$1$，割当遵守者は$2$を割り当てています．

速度のために，配列で処理しています．また，計算精度のために，事後確率は対数をとっており，$\log(1+x)$はnumpyのlog1pという関数を使って計算しています．

def G_sampler(self, gamma_at, gamma_nt, beta_at, beta_nt, beta_co_c, beta_co_t):
    """
    G: compliance type
    0: at
    1: nt
    2: co
    """
    G = np.zeros(self.N)
    
    # pattern for (Z_i = 0) & (W_i = 1)
    G[(self.Z == 0) & (self.W == 1)] = 0

    # pattern for (Z_i = 1) & (W_i = 0)
    G[(self.Z == 1) & (self.W == 0)] = 1

    # pattern for (Z_i = 0) & (W_i = 0)
    X_z0_w0 = self.X[(self.Z == 0) & (self.W == 0)]
    Y_z0_w0 = self.Y[(self.Z == 0) & (self.W == 0)]
    N_z0_w0 = len(Y_z0_w0)
    # log p(G_i = co | Y, \theta)
    log_prob_co = - np.log1p(np.exp(X_z0_w0.dot(gamma_at)) + np.exp(X_z0_w0.dot(gamma_nt)))
    # log p(G_i = nt | Y, \theta)
    log_prob_nt = X_z0_w0.dot(gamma_nt) - np.log1p(np.exp(X_z0_w0.dot(gamma_at)) + np.exp(X_z0_w0.dot(gamma_nt)))
    # log p(Y = 1, | G_i = co, \theta)
    log_prob_y1_given_co_c = X_z0_w0.dot(beta_co_c) - np.log1p(np.exp(X_z0_w0.dot(beta_co_c)))
    # log p(Y = 1, | G_i = nt, \theta)
    log_prob_y1_given_nt = X_z0_w0.dot(beta_nt) - np.log1p(np.exp(X_z0_w0.dot(beta_nt)))
    # log p(Y = 0, | G_i = co, \theta)
    log_prob_y0_given_co_c = - np.log1p(np.exp(X_z0_w0.dot(beta_co_c)))
    # log p(Y = 0, | G_i = nt, \theta)
    log_prob_y0_given_nt = - np.log1p(np.exp(X_z0_w0.dot(beta_nt)))
    # log p(G_i = co | Y_i, \theta)
    logp = Y_z0_w0 * (log_prob_co + log_prob_y1_given_co_c - np.log(np.exp(log_prob_co + log_prob_y1_given_co_c) + np.exp(log_prob_nt + log_prob_y1_given_nt)))
    logp += (1-Y_z0_w0) * (log_prob_co + log_prob_y0_given_co_c - np.log(np.exp(log_prob_co + log_prob_y0_given_co_c) + np.exp(log_prob_nt + log_prob_y0_given_nt)))
    # deciding co or nt
    threshold = np.log(np.random.rand(N_z0_w0))
    G_z0_w0 = np.ones(N_z0_w0)
    G_z0_w0[threshold < logp] = 2
    G[(self.Z == 0) & (self.W == 0)] = G_z0_w0

    # pattern for (Z_i = 1) & (W_i = 1)
    X_z1_w1 = self.X[(self.Z == 1) & (self.W == 1)]
    Y_z1_w1 = self.Y[(self.Z == 1) & (self.W == 1)]
    N_z1_w1 = len(Y_z1_w1)
    # log p(G_i = co | Y, \theta)
    log_prob_co = - np.log1p(np.exp(X_z1_w1.dot(gamma_at)) + np.exp(X_z1_w1.dot(gamma_nt)))
    # log p(G_i = at | Y, \theta)
    log_prob_at = X_z1_w1.dot(gamma_at) - np.log1p(np.exp(X_z1_w1.dot(gamma_at)) + np.exp(X_z1_w1.dot(gamma_nt)))
    # log p(Y = 1, | G_i = co, \theta)
    log_prob_y1_given_co_t = X_z1_w1.dot(beta_co_t) - np.log1p(np.exp(X_z1_w1.dot(beta_co_t)))
    # log p(Y = 1, | G_i = at, \theta)
    log_prob_y1_given_at = X_z1_w1.dot(beta_at) - np.log1p(np.exp(X_z1_w1.dot(beta_at)))
    # log p(Y = 0, | G_i = co, \theta)
    log_prob_y0_given_co_t = - np.log1p(np.exp(X_z1_w1.dot(beta_co_t)))
    # log p(Y = 0, | G_i = at, \theta)
    log_prob_y0_given_at = - np.log1p(np.exp(X_z1_w1.dot(beta_at)))
    # log p(G_i = co | Y_i, \theta)
    logp = Y_z1_w1 * (log_prob_co + log_prob_y1_given_co_t - np.log(np.exp(log_prob_co + log_prob_y1_given_co_t) + np.exp(log_prob_at + log_prob_y1_given_at)))
    logp += (1-Y_z1_w1) * (log_prob_co + log_prob_y0_given_co_t - np.log(np.exp(log_prob_co + log_prob_y0_given_co_t) + np.exp(log_prob_at + log_prob_y0_given_at)))
    # deciding co or at
    threshold = np.log(np.random.rand(N_z1_w1))
    G_z1_w1 = np.zeros(N_z1_w1)
    G_z1_w1[threshold < logp] = 2
    G[(self.Z == 1) & (self.W == 1)] = G_z1_w1

    return G

1.2 パラメータのサンプリング

次にパラメータは$\gamma_{at}, \gamma_{nt}, \beta_{at}, \beta_{nt}, \beta_{co,c}, \beta_{co, t}$という6つに分けてメトロポリスヘイスティング法でサンプリングします．

例えば，$p(\gamma_{at} | G, Y, \gamma_{nt}, \beta_{at}, \beta_{nt}, \beta_{co,c}, \beta_{co, t})$からのサンプリングを考えます．

メトロポリスヘイスティング法では提案分布$q(\gamma_{at}^{\star} | \gamma_{at}^{(i)})$から候補$\gamma_{at}^{\star}$をサンプリングし，その候補を確率$\rm{min}(1, p(\gamma_{at}^{\star} | G, Y, \gamma_{nt})q(\gamma_{at}^{\star} | \gamma_{at}^{(i)}) / p(\gamma_{at}^{(i)} | G, Y, \gamma_{nt})q(\gamma_{at}^{(i)} | \gamma_{at}^{\star}))$で採択します．$q(\gamma_{at}^{\star} | \gamma_{at}^{(i)}) = q(\gamma_{at}^{(i)} | \gamma_{at}^{\star})$となるような提案分布（例えば正規分布）を選ぶと，提案分布の項は無視できます．また，正規化項は分母分子で打ち消し合うので無視できます．正規化項を無視した事後分布は以下です．

\begin{align*}
&p(\gamma_{at} | G, Y, \gamma_{at}) \\
&\propto p(\gamma_{at}) \prod_{i:G_i=at} p(G_i=at|\gamma_{at}, \gamma_{nt})\prod_{i:G_i=nt} p(G_i=nt|\gamma_{at}, \gamma_{nt}) \prod_{i:G_i=co} p(G_i=co|\gamma_{at}, \gamma_{nt}) \\
&\propto \left( \prod_i \frac{1}{1+\exp(X_i\gamma_{at})+\exp(X_i\gamma_{nt})} \frac{\exp(X_i\gamma_{at})}{1+\exp(X_i\gamma_{at})+\exp(X_i\gamma_{nt})} \frac{\exp(X_i\gamma_{nt})}{1+\exp(X_i\gamma_{at})+\exp(X_i\gamma_{nt})}\right)^{(N_a / 12 N)} \\
& \prod_{i:G_i=co} \frac{1}{1+\exp(X_i\gamma_{at})+\exp(X_i\gamma_{nt})} \prod_{i:G_i=at} \frac{\exp(X_i\gamma_{at})}{1+\exp(X_i\gamma_{at})+\exp(X_i\gamma_{nt})} \prod_{i:G_i=nt} \frac{\exp(X_i\gamma_{nt})}{1+\exp(X_i\gamma_{at})+\exp(X_i\gamma_{nt})}
\end{align*}

コードの該当箇所は以下です．計算精度のために事後分布は対数をとって評価しています．また，提案分布は正規分布を利用しています．

def gamma_at_sampler(self, G, gamma_at, gamma_nt):
    # proposing gamma_at_new
    gamma_at_new = gamma_at + self.prop_scale['gamma_at'] * np.random.randn(self.dim)

    # calculating log likelihood
    if sum(G==0) > 0:
        X_at = self.X[G==0]
        log_likelihood_old = X_at.dot(gamma_at).sum() - np.log1p(np.exp(self.X.dot(gamma_at)) + np.exp(self.X.dot(gamma_nt))).sum()
        log_likelihood_new = X_at.dot(gamma_at_new).sum() - np.log1p(np.exp(self.X.dot(gamma_at_new)) + np.exp(self.X.dot(gamma_nt))).sum()
    else:
        log_likelihood_old = - np.log1p(np.exp(self.X.dot(gamma_at)) + np.exp(self.X.dot(gamma_nt))).sum()
        log_likelihood_new = - np.log1p(np.exp(self.X.dot(gamma_at_new)) + np.exp(self.X.dot(gamma_nt))).sum()

    # calculating log prior
    log_prior_old = (self.N_a / 12 / self.N) * (self.X.dot(gamma_at).sum() - 3 * np.log1p(np.exp(self.X.dot(gamma_at)) + np.exp(self.X.dot(gamma_nt))).sum())
    log_prior_new = (self.N_a / 12 / self.N) * (self.X.dot(gamma_at_new).sum() - 3 * np.log1p(np.exp(self.X.dot(gamma_at_new)) + np.exp(self.X.dot(gamma_nt))).sum())

    # deciding accept or not     
    log_acceptance_ratio = log_likelihood_new + log_prior_new - log_likelihood_old - log_prior_old
    if log_acceptance_ratio > np.log(np.random.rand()):
        return gamma_at_new
    else:
        return gamma_at

その他のパラメータも同様の方法でサンプリングできます．

※ちなみに，ハミルトニアンモンテカルロ法も実装しており，問題設定によってはそちらを使ったほうがミキシングが良いと思います．

2. 因果効果の推定

上記の方法で得たパラメータのサンプルから局所平均因果効果$\tau_{late}$を推定するために，観測できていない$Y$を補完する必要があります．

まず，$G_i = co$かつ$Z_i = 0$のケースでは，$Y_i(1)$が観測できていませんが，確率$\frac{\exp(X_i\beta_{co,c})}{1+\exp(X_i\beta_{co,c})}$で$1$をとります．また，$G_i = co$かつ$Z_i = 1$のケースでは，$Y_i(0)$が観測できていませんが，確率$\frac{\exp(X_i\beta_{co,t})}{1+\exp(X_i\beta_{co,t})}$で$1$をとります．

観測できていない$Y$を補完すれば，
$$
\tau_{late} = \frac{1}{N_{co}} \sum_{i:G_i=co} Y_i(1) - Y_i(0)
$$
として，$\tau_{late}$をサンプリングできます．

コードの該当箇所は以下です．

def Y_obs_and_mis_sampler(self, G, beta_co_c, beta_co_t):
    Y_obs_and_mis = np.full((self.N, 2), np.nan)
    
    # pattern for at
    Y_obs_and_mis[G==0, 1] = self.Y[G==0]
    # pattern for nt
    Y_obs_and_mis[G==1, 0] = self.Y[G==1]
    
    N_co_c = int(sum((G==2) & (self.Z==0)))
    N_co_t = int(sum((G==2) & (self.Z==1)))

    # pattern for co, c
    if N_co_c > 0:
        X_co_c = self.X[(G==2) & (self.Z==0)]
        log_prob_y1_given_co_c = X_co_c.dot(beta_co_c) - np.log1p(np.exp(X_co_c.dot(beta_co_c)))
        Y_obs_and_mis[(G==2) & (self.Z==0), 0] = self.Y[(G==2) & (self.Z==0)]
        Y_obs_and_mis[(G==2) & (self.Z==0), 1] = (np.log(np.random.rand(N_co_c)) < log_prob_y1_given_co_c).astype(float)

    # pattern for co, t
    if N_co_t > 0:
        X_co_t = self.X[(G==2) & (self.Z==1)]
        log_prob_y1_given_co_t = X_co_t.dot(beta_co_t) - np.log1p(np.exp(X_co_t.dot(beta_co_t)))
        Y_obs_and_mis[(G==2) & (self.Z==1), 0] = (np.log(np.random.rand(N_co_t)) < log_prob_y1_given_co_t).astype(float)
        Y_obs_and_mis[(G==2) & (self.Z==1), 1] = self.Y[(G==2) & (self.Z==1)]

    return Y_obs_and_mis

def tau_late_sampler(self, Y_mis_and_obs, G):
    N_co = int(sum(G==2))

    if N_co > 0:
        tau_late = (Y_mis_and_obs[(G==2), 1] - Y_mis_and_obs[(G==2), 0]).sum() / N_co
        return tau_late
    else:
        return np.nan

3. サンプリング結果

本コードにおけるサンプリングの実行方法はGitHubのexamples.ipynbを参照してください．

サンプリング結果と本に記載されている推定結果を比較します．

以下のヒストグラムが今回得られたサンプルで，黒の実線が本に記載されている中央値，黒の破線が本に記載されている95%区間です．

$\beta_{\star,age}$以外は本の結果と概ね一致しているように見えますが，$\beta_{\star,age}$のみ大きく外れています．しかし，以下の理由から教科書に記載されている結果が誤っている可能性が高いと考えています．

松浦先生のStanを使った再現とは結果がほぼあっていること．
確実に常時処置拒否者（つまり，$Z_i=1$かつ$W_i=0$）であるデータのみとりだしてロジスティック回帰を行ったところ，$\beta_{nt, age}$は$-0.18$となり，教科書に記載されている中央値よりも今回得られたサンプルの中央値のほうが近いこと
同様のデータにおける漸近分散の平方根は約$0.11$であったが，教科書に記載されている区間幅が$1/10$程度に小さくなっていること．観測数の増加が要因として考えられるが，$\beta_{nt,intercept}, \beta_{nt,copd}, \beta_{nt,heart}$では同様の傾向がないこと．
常時処置拒否者か割当遵守者（つまり，$Z_i=0$かつ$W_i=0$）のデータを合わせて同様の解析を行っても，上記２つの性質が確認されること
とはいえ，遵守タイプの偏りにより本に記載されているような結果が得られる可能性もゼロではありません．いずれにせよ，本のコードが公開されていないため，要因の特定は困難です．

また，以下が各遵守タイプごとの$Y$の期待値と$\tau_{late}$のヒストグラムです．黒の実線が本に記載されている中央値，黒の破線が本に記載されている95%区間です．
中央値は全てのケースで概ねあっていますが，$\mathbb{E}[Y_i(1)|G_i=co]$の信頼区間が狭くなっており，結果として$\tau_{late}$の信用区間が小さく推定されています．

4. 雑感

以下雑感を列挙します．

教科書が間違っている可能性を排除できないのでコードを公開して欲しいです．
今回は学習も兼ねてNumpyのみで実装しましたが，実用上はStanやJuliaを使ったほうがよいと思います．
ロジットモデルはPolya-Gamma分布などを使うとより実装が速くなるかと思います．
$N$に比例する事前分布は，尤度原理に反しますし，逐次推定が困難になるため，個人的にはあまり好みではありません．