docs/07-do-operator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Do-Operator and Bayesian Models\n",
    "\n",
    "The do-operator is a foundational idea in causal methods,\n",
    "used to express, in precise mathematical language,\n",
    "what counterfactual interventions would look like.\n",
    "The do-operator from causal inference has a tight connection to probabilistic modelling.\n",
    "But what exactly is that connection?\n",
    "\n",
    "Having myself been previously confused about the link between graphical models,\n",
    "causal structure, and more,\n",
    "here's my current understanding of the connection.\n",
    "What I hope you'll see here is that the do-operator\n",
    "gives us the ability to simulate what an experiment might look like...\n",
    "without actually doing the experiment,\n",
    "and that at its implementation core,\n",
    "it ends up being nothing more than \"set this thing to a particular value\".\n",
    "If you're curious now, it's time to read on :)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Causal models and Bayesian models, in brief\n",
    "\n",
    "Let’s start first with a brief overview of the connection between \n",
    "causal models and Bayesian models. \n",
    "In doing [my own study on causality](https://github.com/ericmjl/causality), \n",
    "it soon became clear to me that causal models \n",
    "can be read off from mathematical equations quite easily. \n",
    "For a moment, let's assume we have the following set of equations \n",
    "that we presume to describe some data that we observe:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "$$b = \\delta d + \\epsilon e + \\sigma_b$$\n",
    "\n",
    "and \n",
    "\n",
    "$$a = \\beta b + \\gamma c + \\sigma_a$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "My convention here is that the English letters represent data, \n",
    "while the Greek letters represent parameters of the model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make this abstract example a bit more concrete,\n",
    "let's connect these equations to something in real life,\n",
    "such as scholastic achievement.\n",
    "In this highly contrived and fictitious example,\n",
    "let's say that $a$ is the scholastic achievement of a student \n",
    "as measured by the sum total of subject scores\n",
    "across 11 subjects in the O-level examinations \n",
    "administered by the University of Cambridge,\n",
    "and it is thought to be a function of:\n",
    "\n",
    "- $b$, the intelligence quotient (IQ) of the student,\n",
    "- $c$, the cost of a student's tuition fees in thousands of dollars (private vs. public school), and\n",
    "- $\\sigma_a$, the intrinsic variation in student performance as a whole\n",
    "\n",
    "For the term $b$, we think it is a function of:\n",
    "\n",
    "- $d$, the crime rate of their neighborhood, \n",
    "- $e$, the household income of that student, in thousands of dollars, and\n",
    "- $\\sigma_b$, the intrinsic variation in IQ.\n",
    "\n",
    "The astute reader will immediately see the controversy in this example.\n",
    "The causal path we assume here isn't going to be the only plausible model,\n",
    "and there are many, many variables we're leaving out here.\n",
    "However, I'd ask that you let your mind suspend the controversy for a moment;\n",
    "it's always a big brain teaser to try to come up with plausible but fictitious examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**To read off the causal model from this set of equations, \n",
    "everything on the right is causal for the stuff on the left.**\n",
    "By those equations, \n",
    "we state that the values of $a$ that we observe \n",
    "are caused by values of $b$ and $c$ and their Greek coefficients $\\beta$ and $\\gamma$\n",
    "plus some noise $\\sigma_a$,\n",
    "while the values of $b$ that we observe \n",
    "are caused by the values of $d$ and $e$ and their Greek coefficients $\\delta$ and $\\epsilon$\n",
    "plus some noise $\\sigma_b$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**With Bayesian inference on this model, \n",
    "our task here is to estimate the coefficients**, \n",
    "$\\beta$, $\\gamma$, $\\delta$, and $\\epsilon$. \n",
    "More precisely, we are interested in estimating their expected value and uncertainty. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to express the model in such a way that reflects its probabilistic in nature, \n",
    "we need some syntax to express the idea that \n",
    "because $\\delta$, $\\epsilon$ and $\\sigma_b$ are uncertain \n",
    "and are modelled by a probability distribution, \n",
    "therefore $b$ is uncertain and has a corresponding probability distribution too. \n",
    "\n",
    "At the same time, \n",
    "we need an analogous syntax to express that\n",
    "because $\\beta$, $\\gamma$, and $\\sigma_a$ are uncertain\n",
    "and are modelled by a probability distribution,\n",
    "therefore $a$ is also uncertain and has a corresponding probability distribution too. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Moreover, if we assume that $c$, $d$, $e$ are stochastic \n",
    "because they are drawn from a distribution, \n",
    "then we have ~~a classic case where everything is unknown and we can’t do anything~~ \n",
    "an awesome modelling problem at hand! 🙂 \n",
    "In any case, for the first equation, \n",
    "our expression for the distribution of $b$ \n",
    "conditioned on everything on the right would look something like:\n",
    "\n",
    "$$P(b | \\delta, \\epsilon, \\sigma_b, d, e)$$\n",
    "\n",
    "And likewise, the distribution of $a$ conditioned on everything on the right would look something like this:\n",
    "\n",
    "$$P(a | \\beta, \\gamma, \\sigma_a, b, c)$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, in Bayesian inference, \n",
    "we usually regard $(b, c, d, e)$ as being fixed (and hence sacred), \n",
    "because they are the *data that we have observed.* \n",
    "By convention, in many probabilistic modelling problems, \n",
    "we’re not really concerned about the data generating processes for $(b, c, d, e)$ \n",
    "because they aren’t usually the sample-invariant, \n",
    "intrinsic parameter of a system we’re interested in,\n",
    "unlike the Greek-letter symbols which take on their values\n",
    "independent of any given measured sample."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Causal inference vs Bayesian inference, in brief\n",
    "\n",
    "I’d now like to address “causal inference” vs. “bayesian inference”; I think the use of a contrast is going to be helpful here.\n",
    "\n",
    "Bayesian inference asks the question, “given the observed data and a presumed model with parameters, what is the expectation and uncertainty in the parameters that could have generated the observed data?”\n",
    "\n",
    "Causal inference asks the question, “given the observed data, what are plausible structures of the model that could have generated the observed data?”"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structures: in graphics and in equations\n",
    "\n",
    "By structures, we’re really asking about the relationships between variables. \n",
    "Using the above example again, \n",
    "if $a$ is caused by $b$ and $c$, \n",
    "then in abstract, we'd write that $a = f(b, c)$. \n",
    "We would then draw the following diagram to express the relationship:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_format='retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import daft\n",
    "\n",
    "scale = 2.0\n",
    "pgm = daft.PGM()\n",
    "pgm.add_node(\"a\", r\"$a$\", 1.5, 1)\n",
    "pgm.add_node(\"b\", r\"$b$\", 1, 2)\n",
    "pgm.add_node(\"c\", r\"$c$\", 2, 2)\n",
    "\n",
    "a_color = {\"ec\": \"blue\"}\n",
    "\n",
    "pgm.add_edge(\"b\", \"a\", plot_params=a_color)\n",
    "pgm.add_edge(\"c\", \"a\", plot_params=a_color)\n",
    "pgm.render();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Likewise, since we presume that $b$ is caused by $d$ and $e$, \n",
    "then the functional form of the causal relationship will be $b = g(d, e)$. \n",
    "We would then draw the following diagram to express the relationship:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pgm = daft.PGM()\n",
    "pgm.add_node(\"b\", r\"$b$\", 1.5, 1)\n",
    "pgm.add_node(\"d\", r\"$d$\", 1, 2)\n",
    "pgm.add_node(\"e\", r\"$e$\", 2, 2)\n",
    "\n",
    "b_color = {\"ec\": \"red\"}\n",
    "\n",
    "pgm.add_edge(\"d\", \"b\", plot_params=b_color)\n",
    "pgm.add_edge(\"e\", \"b\", plot_params=b_color)\n",
    "pgm.render();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And taken together, the full model would look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pgm = daft.PGM()\n",
    "pgm.add_node(\"a\", r\"$a$\", 1.5, 1)\n",
    "pgm.add_node(\"b\", r\"$b$\", 2, 2)\n",
    "pgm.add_node(\"c\", r\"$c$\", 1, 2)\n",
    "pgm.add_node(\"d\", r\"$d$\", 1.5, 3)\n",
    "pgm.add_node(\"e\", r\"$e$\", 2.5, 3)\n",
    "\n",
    "pgm.add_edge(\"c\", \"a\", plot_params=a_color)\n",
    "pgm.add_edge(\"b\", \"a\", plot_params=a_color)\n",
    "pgm.add_edge(\"d\", \"b\", plot_params=b_color)\n",
    "pgm.add_edge(\"e\", \"b\", plot_params=b_color)\n",
    "\n",
    "pgm.add_text(3, 1.25, r\"$a = f(b, c) = \\beta b + \\gamma c + \\sigma_a$\")\n",
    "pgm.add_text(3, 2.4, r\"$b = f(d, e) = \\delta d + \\epsilon e + \\sigma_b$\")\n",
    "\n",
    "pgm.render();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In economics, the term “structural equation models” \n",
    "refers to the system of equations \n",
    "that form the econometric model that economists build. \n",
    "Usually, those models are linear in nature. \n",
    "On occasion, though, additional functional forms might be used \n",
    "(sigmoidal, piecewise linear, neural network, etc.), \n",
    "if they help model the phenomena at hand. \n",
    "Whatever the form of the equation gets encapsulated into $f(v_1, v_2, ..., v_n)$,\n",
    "where $v_1... v_n$ refer to the variables like $b$ and $c$ above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Structure, then, can refer to both the graphical structure of the model \n",
    "and the particular form of equations. \n",
    "As far as I have seen, \n",
    "most causal models assume some kind of linear equation between variables, \n",
    "though there are exceptions; \n",
    "in addition, causal inference, as taught, \n",
    "is usually concerned with inferring the *graphical* relationship between variables, \n",
    "presuming some linear form underneath. \n",
    "Inferring the model structure is what we call \"model inference\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*A note for deep learners: this is what model inference actually is: \n",
    "inferring the structure of a model. \n",
    "More generally, if you think about linguistic convention, \n",
    "“X inference” usually refers to the tools and processes used in inferring X. \n",
    "Calling model inference the forward pass through the model breaks linguistic convention, \n",
    "and hence introduces viscosity in communication with others \n",
    "who adopt said linguistic convention!*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Restating the distinction\n",
    "\n",
    "At this point, I think it's an appropriate moment to try to re-state clearly\n",
    "what the relationship between a \"causal inference\" and a \"Bayesian inference\" is.\n",
    "They are **both** concerned with the system of equations that make up our model.\n",
    "However, in causal inference, \n",
    "we are primarily concerned with the relationship between observed variables, \n",
    "expressed as math equations.\n",
    "In Bayesian inference, we are primarily concerned with the parameters of those equations\n",
    "and their uncertainty."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The do-operator\n",
    "\n",
    "Finally, we reach the point where we can touch the do-operator!\n",
    "This is such a key and crucial idea to causal methods,\n",
    "as it allows us to do **counterfactual** arguments conditioned on a presumed model.\n",
    "To illustrate what we mean by the do-operator,\n",
    "I'm going to rely on code, prose, and equations together.\n",
    "To start, let's implement the full probabilistic model above in Python code.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "from scipy.stats import norm\n",
    "import pandas as pd \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Equations\n",
    "beta = 10\n",
    "gamma = 0.3\n",
    "delta = -30\n",
    "epsilon = 0.1\n",
    "variance = 10\n",
    "sigma_a_dist = norm(0, variance)\n",
    "sigma_b_dist = norm(0, variance)\n",
    "\n",
    "\n",
    "def a_equation(b, c):\n",
    "    return beta * b + gamma * c + sigma_a_dist.rvs(len(b))\n",
    "\n",
    "def b_equation(d, e):\n",
    "    return delta * d + epsilon * e + sigma_b_dist.rvs(len(d))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate the data, we always begin with the nodes\n",
    "that have no parental nodes.\n",
    "In more complicated networks, \n",
    "we would leverage tools from network science, in this case, the _topological sort_,\n",
    "to identify the exact order in which we need to simulate observations.\n",
    "The observed data that we would end up collecting \n",
    "for this system looks like the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1_000\n",
    "cc = norm(50, 5).rvs(N)\n",
    "dd = norm(-4, 1).rvs(N)\n",
    "ee = norm(30, 4).rvs(N)\n",
    "\n",
    "bb = b_equation(dd, ee)\n",
    "aa = a_equation(bb, cc)\n",
    "\n",
    "\n",
    "data = pd.DataFrame({\"a\": aa, \"b\": bb, \"c\": cc, \"d\": dd, \"e\": ee})\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And if we were to visualize the marginal distribution of each of the variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axes = plt.subplots(figsize=(15, 3), ncols=5)\n",
    "\n",
    "for var, ax in zip([aa, bb, cc, dd, ee], axes.flatten()):\n",
    "    ax.hist(var)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, plotting the joint distributions, we get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns \n",
    "\n",
    "sns.pairplot(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a feel for what kind of data we'll collect,\n",
    "let's explore the three ways that the do-operator can be used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The do-operator: local counterfactuals\n",
    "\n",
    "At an elementary level, the do-operator expresses mathematically the following quesiton,\n",
    "“For a given sample, what if one of its variables took on a different value?” \n",
    "You can think of this as a \"personalized\" counterfactual for a given sample,\n",
    "whether that sample is a patient, student, employee, citizen, or some other thing.\n",
    "\n",
    "Some concrete examples of this in action are:\n",
    "\n",
    "1. “For this given patient, what if the treatment given was the alternative treatment?”\n",
    "2. “For this given student, what if they came from a wealthier household?”\n",
    "3. “For this given tennis player, what if the court they played on were grass instead of clay?”\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "As you can see, one of the questions we’re asking when asking counterfactuals \n",
    "are in fact *personalized* for a given sample.\n",
    "\n",
    "If we go back to the structural equations above, \n",
    "we could isolate a given observational data point $(a, b, c)$ and then ask the question, \n",
    "“What if in row $i$, $b$ took on a certain value $B$ instead of $b_i$?” \n",
    "This question, when asked, results in our probabilistic structure changing a bit:\n",
    "\n",
    "1. We are now asking for $P(a_i | do(b_i=B), \\beta, \\gamma, c_i)$, where $i$ refers to the particular sample index.\n",
    "2. Our counterfactual question presumes a known value of $b$, and hence no longer requires us to generate it from $(d, e)$. We can effectively cut $(d, e)$ out of the picture.\n",
    "\n",
    "As long as we preserve uncertainty in the parameter values, \n",
    "we can obtain counterfactual uncertainty as well.\n",
    "To illustrate how we do personalized counterfactuals in a Bayesian setting,\n",
    "let's see how to do it with PyMC."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation with PyMC\n",
    "\n",
    "We're going to start first by implementing the models.\n",
    "There's no do-operations happening just yet,\n",
    "we're just writing down the equations in PyMC first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc as pm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first implement the set of equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as model_a:\n",
    "    a = pm.Data(\"a\", data[\"a\"])\n",
    "    b = pm.Data(\"b\", data[\"b\"])\n",
    "    c = pm.Data(\"c\", data[\"c\"])\n",
    "    # Priors for model parameters\n",
    "    beta = pm.Normal(\"beta\", mu=10, sigma=3)\n",
    "    gamma = pm.Normal(\"gamma\", mu=0, sigma=10)\n",
    "    sigma_a = pm.Exponential(\"sigma_a\", lam=1)\n",
    "    pm.Normal(\"like\", mu=beta * b + gamma * c, sigma=sigma_a, observed=a)\n",
    "\n",
    "\n",
    "with pm.Model() as model_b:\n",
    "    b = pm.Data(\"b\", data[\"b\"])\n",
    "    d = pm.Data(\"d\", data[\"d\"])\n",
    "    e = pm.Data(\"e\", data[\"e\"])\n",
    "    # Priors for model parameters\n",
    "    delta = pm.Normal(\"delta\", mu=-15, sigma=15)\n",
    "\n",
    "    epsilon = pm.Normal(\"epsilon\", mu=0, sigma=1)\n",
    "    sigma_b = pm.Exponential(\"sigma_b\", lam=1)\n",
    "    pm.Normal(\"like\", mu=delta * d + epsilon * e, sigma=sigma_b, observed=b)\n",
    "\n",
    "\n",
    "with pm.Model() as model_joint:\n",
    "    a = pm.Data(\"a\", data[\"a\"])\n",
    "    b = pm.Data(\"b\", data[\"b\"])\n",
    "    c = pm.Data(\"c\", data[\"c\"])\n",
    "    d = pm.Data(\"d\", data[\"d\"])\n",
    "    e = pm.Data(\"e\", data[\"e\"])\n",
    "    # Priors for model parameters\n",
    "    delta = pm.Normal(\"delta\", mu=-15, sigma=15)\n",
    "\n",
    "    epsilon = pm.Normal(\"epsilon\", mu=0, sigma=1)\n",
    "    sigma_b = pm.Exponential(\"sigma_b\", lam=1)\n",
    "    pm.Normal(\"like_b\", mu=delta * d + epsilon * e, sigma=sigma_b, observed=b)\n",
    "    # Priors for model parameters\n",
    "    beta = pm.Normal(\"beta\", mu=10, sigma=3)\n",
    "    gamma = pm.Normal(\"gamma\", mu=0, sigma=10)\n",
    "    sigma_a = pm.Exponential(\"sigma_a\", lam=1)\n",
    "    pm.Normal(\"like_a\", mu=beta * b + gamma * c, sigma=sigma_a, observed=a)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now plot the graphical model provided by PyMC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pm.model_to_graphviz(model_a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this model, we see that $a$, $b$, and $c$ are all observed data,\n",
    "nested within the plate representing 1,000 data points.\n",
    "$\\beta$, $\\gamma$, and $\\sigma_a$ are the parameters of the model\n",
    "that are invariant to any particular data point\n",
    "and hence are located outside of the plate.\n",
    "Those are our system-level parameters.\n",
    "An analogous diagram exists for $b$'s model as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pm.model_to_graphviz(model_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pm.model_to_graphviz(model_joint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameter inference\n",
    "\n",
    "We can now use PyMC's inference machinery, the Inference Buttom (tm),\n",
    "to infer the values of the parameters above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model_a:\n",
    "    idata_a = pm.sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model_b:\n",
    "    idata_b = pm.sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model_joint:\n",
    "    idata_joint = pm.sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I will be skipping over a rigorous Bayesian workflow here\n",
    "because it is not the point of the notebook.\n",
    "Nonetheless, here is a set of plots for our posteriors;\n",
    "you can check-them against the original Gaussian distributions above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz as az\n",
    "\n",
    "az.plot_posterior(idata_a, ref_val=[10, 0.3, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(idata_b, ref_val=[-30, 0.1, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(idata_joint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In all of the cases, we're pretty close, though a bit off. \n",
    "As the goal here is to show a personalized counterfactual,\n",
    "we're not going to worry too much about the parameter recovery accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Personalized counterfactuals in PyMC\n",
    "\n",
    "Personalized counterfactuals ask the question, \n",
    "\"What would have happened for a particular sample,\n",
    "i.e. one row in the dataset,\n",
    "had one of its observed variables been a different value?\"\n",
    "In this case, we're going to take one of the observations \n",
    "and ask what would happen if we counterfactually set $c$ to a different value.\n",
    "\n",
    "We'll start by isolating a sample of interest:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = data.sample(random_state=491)  # random state for reproducibility\n",
    "sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then simulate what the posterior for that sample would look like\n",
    "under the original case vs. the counterfactual case.\n",
    "In the first code block below, we set our data variables to be a single sample\n",
    "and then leverage the posterior predictive samples to simulate\n",
    "what `a` would look like under the original data.\n",
    "In the second code block, we set our data variables to be a single sample as well,\n",
    "except that we have changed `c` to be a totally different value from what it was before.\n",
    "This gives us our counterfactual scenario!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "\n",
    "with model_a:\n",
    "    # Simulate posterior predictive under original observed data.\n",
    "    pm.set_data({\"a\": sample[\"a\"], \"b\": sample[\"b\"], \"c\": sample[\"c\"]})\n",
    "    ppc_or = pm.sample_posterior_predictive(idata_a)\n",
    "\n",
    "    # Simulate posterior predictive under a counterfactual scenario.\n",
    "    # ****This is the do-operator in action!****\n",
    "    pm.set_data({\"a\": sample[\"a\"], \"b\": sample[\"b\"], \"c\": sample[\"c\"] / 10})\n",
    "    ppc_cf = pm.sample_posterior_predictive(idata_a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ppc_or.posterior_predictive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(10, 5), ncols=2)\n",
    "az.plot_posterior(ppc_or.posterior_predictive[\"like\"], ax=axes[0])\n",
    "axes[0].set_title(\"Original\")\n",
    "\n",
    "az.plot_posterior(ppc_cf.posterior_predictive[\"like\"], ax=axes[1])\n",
    "axes[1].set_title(\"Counterfactual\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's really cool here is that we've done a personalized counterfactual\n",
    "for that particular student.\n",
    "In this case, we set the cost of the student's tuition fees to 1/10 the original, \n",
    "akin to sending them from a private school to a public school.\n",
    "Under this particularly flawed model, we would expect their sum of grades to go down,\n",
    "though not by a large magnitude.\n",
    "\n",
    "The key trick to enabling this\n",
    "was using `pm.Data()` containers for our data,\n",
    "thus registering them with PyMC (and Aesara underneath)\n",
    "as being hot-swappable entities with `pm.set_data()`.\n",
    "\n",
    "Finally, because our _posterior_ distributions \n",
    "contain our probabilistic beliefs having seen the data,\n",
    "or more colloquially, \"fits of our parameters\",\n",
    "we sample from the _posterior predictive_ distribution\n",
    "to identify what we would have gotten in a counterfactual situation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The do-operator: global counterfactuals\n",
    "\n",
    "Let’s also think about another possible counterfactual question: \n",
    "What if the system parameters were different?\n",
    "\n",
    "Recall here that “system parameters” refer to the linear coefficients. \n",
    "They aren’t properties of any observation (or sample), \n",
    "but are properties of the entire system as a whole, \n",
    "hence the moniker “global” counterfactuals.\n",
    "\n",
    "To do global counterfactuals, \n",
    "we actually need to condition the coefficients’ values on a particular value, \n",
    "just like we did for data on the personalized counterfactual. \n",
    "For example, if we conditioned $\\beta$  to be equal to 0,\n",
    "in other words severing the relationship between $a$ and $b$, then:\n",
    "\n",
    "1. We are now asking for $P(a|do(\\beta=0), \\gamma, b, c)$. (Notice the omission of subscript $i$, it’s intentional!)\n",
    "2. Our counterfactual question presumes a known value of $\\beta$, but not a presumed known value of any of $(b, c)$. (We aren't asking about any particular sample, after all!)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One way we can implement this do-operation is to reconstruct the model from our posteriors\n",
    "while hard-coding the value of $\\beta$ to 0.\n",
    "\n",
    "Let's see how to make this happen.\n",
    "\n",
    "Firstly, in order to take advantage of our fitted posteriors,\n",
    "we will use a `from_posterior` function written implemented in the PyMC how-to guides."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import stats \n",
    "\n",
    "# Credit: https://docs.pymc.io/projects/examples/en/latest/pymc3_howto/updating_priors.html\n",
    "def from_posterior(param, samples):\n",
    "    smin, smax = np.min(samples), np.max(samples)\n",
    "    width = smax - smin\n",
    "    x = np.linspace(smin, smax, 100)\n",
    "    y = stats.gaussian_kde(samples)(x)\n",
    "\n",
    "    # what was never sampled should have a small probability but not 0,\n",
    "    # so we'll extend the domain and use linear approximation of density on it\n",
    "    x = np.concatenate([[x[0] - 1 * width], x, [x[-1] + 1 * width]])\n",
    "    y = np.concatenate([[0], y, [0]])\n",
    "    return pm.Interpolated(param, x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we'll need to rewrite the model a little bit,\n",
    "this time hard-coding one of the variables to a particular value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as model_a_modified:\n",
    "    a = pm.Data(\"a\", data[\"a\"])\n",
    "    b = pm.Data(\"b\", data[\"b\"])\n",
    "    c = pm.Data(\"c\", data[\"c\"])\n",
    "    # Model parameters\n",
    "    beta = 0  # ****This is the do-operator in action!****\n",
    "    gamma = from_posterior(\"gamma\", idata_a.posterior[\"gamma\"].values.flatten())\n",
    "    sigma_a = from_posterior(\"sigma_a\", idata_a.posterior[\"sigma_a\"].values.flatten())\n",
    "    pm.Normal(\"like\", mu=beta * b + gamma * c, sigma=sigma_a, observed=a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, since our posteriors have become priors in this new model,\n",
    "we sample from the prior predictive distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model_a_modified:\n",
    "    trace = pm.sample_prior_predictive()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we're able to examine how our observations have changed\n",
    "based on the intervention at the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_global_counterfactual(idx: int):\n",
    "    az.plot_posterior(trace.prior_predictive[\"like\"][0, :, idx])\n",
    "    plt.gca().set_title(f\"Original value: {data.loc[idx, 'a']:.2f}\")\n",
    "\n",
    "\n",
    "plot_global_counterfactual(2)\n",
    "plot_global_counterfactual(3)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see here that by severing the connection between the IQ of a student and their grades,\n",
    "each student's total sum of grades goes down to 1/20th of their original.\n",
    "Those of us smarter than this author would probably have intuited this point without needing to code it up (by examining the magnitude of the slope coefficients),\n",
    "but in the case of exotic functional forms with interaction terms (or a neural net structure),\n",
    "the value of an explicit, global perturbation of an interpretable parameter is quite evident!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The do-operator: combining global and personalized counterfactuals\n",
    "\n",
    "If we’re thinking logically here, \n",
    "we’ll soon realize that it’s also possible \n",
    "to combine the two aforementioned counterfactuals together. \n",
    "We can ask the question, \n",
    "what is $P(a_i | do(\\beta=0, b_i=B), \\gamma, c_i)$? \n",
    "(Note again the presence of the index $i$!)\n",
    "\n",
    "Algorithmically, this question essentially translates to:\n",
    "\n",
    "1. picking out sample $i$, \n",
    "2. setting $b=B$, \n",
    "3. setting $beta=3.14$,\n",
    "4. and evaluating what $a_i$ would look like under those two conditions.\n",
    "\n",
    "For the sake of illustration, here it is in PyMC code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with model_a_modified:\n",
    "    pm.set_data({\"a\": sample[\"a\"], \"b\": sample[\"b\"], \"c\": sample[\"c\"]})\n",
    "    ppc_or = pm.sample_prior_predictive()\n",
    "\n",
    "    # cf = \"counterfactual\"\n",
    "    pm.set_data({\"a\": sample[\"a\"], \"b\": sample[\"b\"], \"c\": sample[\"c\"] * 0.1})\n",
    "    ppc_cf = pm.sample_prior_predictive()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(15, 5), ncols=2)\n",
    "az.plot_posterior(ppc_or.prior_predictive[\"like\"], ax=axes[0])\n",
    "axes[0].set_title(r\"Distribution of $a$ with $\\beta=0$ and $c = c_i$\")\n",
    "\n",
    "az.plot_posterior(ppc_cf.prior_predictive[\"like\"], ax=axes[1])\n",
    "axes[1].set_title(r\"Distribution of $a$ with $\\beta=0$ and $c = \\frac{c_i}{10}$\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Framework\n",
    "\n",
    "Having worked through this example, we've finally arrived at a framework for thinking through the connection between Bayesian models and causal models. \n",
    "\n",
    "To know which kind of counterfactual we need to employ, we have to be extremely clear on the exact question we’re trying to ask. Are we trying to ask:\n",
    "\n",
    "1. A personalized question? (”What would have happened to this particular sample had its dependent variable been set to a particular value?”)\n",
    "2. A systems-level question? (”What would happen to all of our observations if a system parameter was set to a particular value?”)\n",
    "3. A hybrid question? (”What would have happened to this particular sample had its dependent variable been set to a particular value *and* the system parameter set to a particular value?”)\n",
    "\n",
    "If we have our variables’ dependencies clearly and explicitly stated, then it becomes easy to ask the do-operator family of questions, which basically are asking, “What happens if we set something in the model to a particular value?”"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we have seen how to do causal inference on Bayesian models written in PyMC. \n",
    "We started with the contrast and comparison between causal and Bayesian inference.\n",
    "Then, we went through the logical framework of asking counterfactuals on (1) a per-sample basis, (2) a system-wide basis, and (3) both together.\n",
    "\n",
    "The value of the do-operator, and by extension causal methods,\n",
    "is that they provide us the ability to ask these \"what if\" questions\n",
    "in cases where conducting an experiment would be unethical, cost-prohibitive, or logistically challenging.\n",
    "Combining it with Bayesian models gives us the ability to assess\n",
    "what would have happened in a counterfactual world (causal)\n",
    "while also calculating the full range of possibilities (Bayesian).\n",
    "Blending the two together is nothing more than a matter of logic!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Credits\n",
    "\n",
    "I would like to credit Thomas Wiecki and Ben Vincent for reviewing an early draft of this notebook."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
  },
  "kernelspec": {
   "display_name": "Python 3.8.2 64-bit ('base': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}