<jupyter_start><jupyter_text>Bayesian inference of average causal effect with instrumental variables<jupyter_code>import stan
import numpy as np
from scipy.stats import beta
import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>We will implement inference of average causal effect using the finite-response treatment given by Pearl, Causality 2009.  Warm-up: using PyStan & beta-binomial model Below is a minimal working example for using PyStan for inference on a coin flip with a beta-distributed prior.<jupyter_code>x = np.random.choice([0,1], p=[0.2,0.8], size=1000) # data<jupyter_output><empty_output><jupyter_text>Declare the model using the Stan probabilistic programming language<jupyter_code>coin_flip = """
data {
  int<lower=0> N;
  int<lower=0, upper=1> x[N];   
}
parameters {
  real<lower=0, upper=1> theta;
}
model {
  theta ~ beta(0.5, 0.5);
  x ~ bernoulli(theta);
}
"""<jupyter_output><empty_output><jupyter_text>Compile the model<jupyter_code>coin_data = {"x": x, "N":len(x)}
posterior = stan.build(coin_flip, data=coin_data, random_seed=1)<jupyter_output><empty_output><jupyter_text>Perform inference<jupyter_code>fit = posterior.sample(num_chains=4, num_samples=1000)<jupyter_output><empty_output><jupyter_text>Compare to theory<jupyter_code>df = fit.to_frame()
a = 0.5 + np.sum(x)
b = 0.5 + len(x) - np.sum(x)
theta_sp = np.linspace(0.79, 0.86)
posterior_theory = beta.pdf(theta_sp, a, b)
fig, ax = plt.subplots()
df['theta'].hist(ax=ax, bins='auto', label='samples', density=True)
ax.plot(theta_sp, posterior_theory, '-k', label='theory')
ax.set_xlabel(r'$\theta$')
ax.set_ylabel(r'$P(\theta)$')
ax.legend()<jupyter_output><empty_output><jupyter_text>As expected  Finite-response treatment of instrumental variable model Imagine we are running a clinical trial, and the variable $Z \in \{0,1\}$ represents the treatment assigned to a patient. Let $X \in \{0,1\}$ denote whether the patient takes the treatment, and $Y \in \{0,1\}$ denote the observed response. There are factors (both observed and unobserved) $U$ that influence the way a subject responds to treatments, which may also affect the patient's choice to take the treatment. When $z \neq x$ we have imperfect compliance.   Pearl (Causality, 2009) shows that regardless of the domain of $U$, it can always be partitioned into four equivalence classes for the relationship between $X$ and $Y$, and similar for $Z$ and $X$, resulting in a 16 possible subpopulations to describe $U$ completely:    Let's generate the data $X$, $Y$, $Z$ according to these canonical partitions, so we have access to the ground truth for inference<jupyter_code>from scipy.stats import multinomial, dirichlet
p_subpop = dirichlet.rvs(0.5*np.ones(16), size=1, random_state=42)
p_subpop = p_subpop.squeeze(0)
multinomial.rvs(1, p_subpop, size=N, random_state=42)
N = 1000
R = multinomial.rvs(1, p_subpop, size=N, random_state=42).reshape((N,4,4))
_, rx, ry = np.where(R)
R[0]
rx[0], ry[0]
def make_response_mapping(u: int, r: int):
    if r==0:
        return 0
    elif r==1 and u==0:
        return 0
    elif r==1 and u==1:
        return 1
    elif r==2 and u==0:
        return 1
    elif r==2 and u==1:
        return 0
    elif r==3:
        return 1
    else:
        raise ValueError
_make_response_mapping = np.vectorize(make_response_mapping)
x = _make_response_mapping(z, rx)
y = _make_response_mapping(x, ry)
z.shape, x.shape, y.shape<jupyter_output><empty_output><jupyter_text>Make PyStan model<jupyter_code>canonical_partition_instrument = """
data {
  int<lower=0> N;
  int<lower=0, upper=1> x[N];
  int<lower=0, upper=1> y[N];
  int<lower=0, upper=1> z[N];
}
parameters {  
  vector<lower=0, upper=1>[16] vr;
  int<lower=1, upper=16> r;
}
model {
  r ~ categorical(vr);
  
  if((r == 1) || (r == 5) || (r==9) || (r==13))
    x=0;
  else if (((r == 2) || (r == 6) || (r==10) || (r==14)) && (z==0))
    x=0;
  else if (((r == 2) || (r == 6) || (r==10) || (r==14)) && (z==1))
    x=1;
  else if (((r == 3) || (r == 7) || (r==11) || (r==15)) && (z==0))
    x=1;
  else if (((r == 3) || (r == 7) || (r==11) || (r==15)) && (z==1))
    x=0;
  else if((r == 4) || (r == 8) || (r==12) || (r==16))
    x=1;

  if((r == 1) || (r == 2) || (r==3) || (r==4))
    y=0;
  else if (((r == 5) || (r == 6) || (r==7) || (r==8)) && (x==0))
    y=0;
  else if (((r == 5) || (r == 6) || (r==7) || (r==8)) && (x==1))
    y=1;
  else if (((r == 9) || (r == 10) || (r==11) || (r==12)) && (x==0))
    y=1;
  else if (((r == 9) || (r == 10) || (r==11) || (r==12)) && (x==1))
    y=0;
  else if (((r == 13) || (r == 14) || (r==15) || (r==16)) && (x==0))
    y=1;
      
}
"""<jupyter_output><empty_output><jupyter_text>Compile the model<jupyter_code>canonical_partition_instrument_data = {"x": x, "y":y, "z":z, "N":N}
posterior = stan.build(canonical_partition_instrument, data=canonical_partition_instrument_data, random_seed=1)<jupyter_output><empty_output>