bym_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on PyMC v4.1.3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/conor/miniconda3/envs/pymc_non_dev/lib/python3.10/site-packages/pymc/sampling_jax.py:36: UserWarning: This module is experimental.\n",
      "  warnings.warn(\"This module is experimental.\")\n"
     ]
    }
   ],
   "source": [
    "import aesara \n",
    "import aesara.tensor as at\n",
    "import arviz as az\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np \n",
    "import pymc as pm \n",
    "import xarray as xr \n",
    "import pandas as pd \n",
    "import pymc.sampling_jax\n",
    "import numpyro \n",
    "\n",
    "print(f\"Running on PyMC v{pm.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data: loading in the Scottish lip cancer data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_scot_cancer = pd.read_csv(pm.get_data(\"scotland_lips_cancer.csv\"))\n",
    "\n",
    "# observed cancer counts\n",
    "y = df_scot_cancer[\"CANCER\"].values\n",
    "\n",
    "# number of observations\n",
    "N = len(y)\n",
    "\n",
    "# expected cancer counts E for each county: this is calculated using age-standardized rates of the local population\n",
    "E = df_scot_cancer[\"CEXP\"].values\n",
    "logE = np.log(E)\n",
    "\n",
    "# proportion of the population engaged in agriculture, forestry, or fishing\n",
    "x = df_scot_cancer[\"AFF\"].values / 10.0\n",
    "\n",
    "# spatial adjacency information: column `ADJ` contains list entries which are preprocessed to obtain adj as a list of lists\n",
    "adj = (\n",
    "    df_scot_cancer[\"ADJ\"].apply(lambda x: [int(val) for val in x.strip(\"][\").split(\",\")]).to_list()\n",
    ")\n",
    "\n",
    "# change to Python indexing (i.e. -1)\n",
    "for i in range(len(adj)):\n",
    "    for j in range(len(adj[i])):\n",
    "        adj[i][j] = adj[i][j] - 1\n",
    "\n",
    "# storing the adjacency matrix as a two-dimensional np.array\n",
    "adj_matrix = np.zeros((N, N), dtype=\"int32\")\n",
    "\n",
    "for area in range(N):\n",
    "    adj_matrix[area, adj[area]] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# getting the necessary information to use for a spatial model specification \n",
    "N_edges = (adj_matrix == 1).sum()\n",
    "node1 = np.where(adj_matrix == 1)[0] \n",
    "node2 = np.where(adj_matrix == 1)[1] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords = {\"num_areas\": N, \n",
    "          \"num_neighbours\": N_edges}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining relevant functions for an ICAR prior\n",
    "\n",
    "Writing a term proportional to the $\\log p$ of the ICAR component, in terms of pairwise differences you get\n",
    "\\begin{align}\n",
    "\\log p(\\boldsymbol{\\phi})\\propto \\frac{-1}{2}\\sum_{j\\sim i}(\\phi_i - \\phi_j)^2.\n",
    "\\end{align}\n",
    "Note that this is for the instance where $\\tau=1$, and so below we experiment with including the scale parameter within the `pm.Potential` call. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# creating the pairwise specification for the ICAR prior \n",
    "def icar_pairwise_diff(phi, node1, node2):\n",
    "    return -0.5*((phi[node1]-phi[node2])**2).sum()\n",
    "# pairwise specification of the ICAR prior, with the inclusion of the scale parameter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BYM Model\n",
    "\n",
    "The BYM model that we will now fit to the cancer data is written as \n",
    "\n",
    "\\begin{align}\n",
    "y_i &\\sim \\text{Poisson}(\\mu_i)\\\\\n",
    "\\log \\mu_i &= \\beta_0 + \\beta_1 x_i + \\log E_i + \\theta_i + \\phi_i \\\\\n",
    "\\boldsymbol{\\theta} &\\sim \\text{Normal}\\big (0, \\frac{1}{\\tau_{\\theta}^2}\\big ) \\\\\n",
    "\\boldsymbol{\\phi} &\\sim \\text{Normal}\\big (0, \\big (\\tau_\\phi^2\\big[\\mathbf{D}-\\mathbf{W} \\big]\\big )^{-1}\\big )\\\\\n",
    "\\tau_\\theta^2 &\\sim \\text{Gamma}(3.2761, 1.81)\\\\\n",
    "\\tau_\\phi^2 &\\sim \\text{Gamma}(1, 1)\\\\\n",
    "\\beta_0, \\beta_1 &\\sim \\text{Normal}(0, 5).\n",
    "\\end{align}\n",
    "\n",
    "The random effect for each of the area's $i$ can be considered as \n",
    "\\begin{align}\n",
    "\\psi_i = \\theta_i + \\phi_i,  \n",
    "\\end{align}\n",
    "where $\\phi_i$ is the spatial random effect modelled using an ICAR prior, and $\\theta_i$ is an independent random effect, that accounts for the additional variation in the data set, that can neither be explained by spatial dependencies or variatibility from the Poisson distribution. \n",
    "\n",
    "To ensure identifiability, we need the vectors $\\boldsymbol{\\phi}$ and $\\boldsymbol{\\theta}$ to independently sum to $0$. We will try and write the BYM model below through two different parameterizations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model(coords={\"num_areas\": np.arange(N)}) as bym_model:\n",
    "    # precision priors \n",
    "    tau_theta = pm.Gamma(\"tau_theta\", alpha=3.2761, beta=1.81)\n",
    "    tau_phi = pm.Gamma(\"tau_phi\", alpha=1, beta=1)\n",
    "\n",
    "    # transform to standard deviation \n",
    "    sigma_phi = pm.Deterministic(\"sigma_phi\", 1/at.sqrt(tau_phi))\n",
    "\n",
    "    # independent random effect prior\n",
    "    theta = pm.Normal(\"theta\", mu=0, sigma=1/at.sqrt(tau_theta), dims=\"num_areas\")\n",
    "    # spatial ICAR random effect prior \n",
    "    phi = pm.Flat(\"phi\", dims=\"num_areas\")\n",
    "    pm.Potential(\"spatial_diff\", icar_pairwise_diff(phi, node1, node2)) \n",
    "\n",
    "    # constraints for each of the random effect vectors \n",
    "    zero_constraint = pm.Normal.dist(mu=0.0, sigma=np.sqrt(0.001))\n",
    "    pm.Potential(\"zero_sum_theta\", pm.logp(zero_constraint, pm.math.sum(theta))) \n",
    "    pm.Potential(\"zero_sum_phi\", pm.logp(zero_constraint, pm.math.sum(phi)))\n",
    "\n",
    "    # regression coefficient priors \n",
    "    beta0 = pm.Normal(\"beta0\", mu=0, sigma=5)\n",
    "    beta1 = pm.Normal(\"beta1\", mu=0, sigma=5)\n",
    "    \n",
    "    # linear predictor \n",
    "    eta = pm.Deterministic(\"eta\", logE + beta0 + beta1*x + sigma_phi*phi + theta, dims=\"num_areas\") \n",
    "\n",
    "    # likelihood\n",
    "    obs = pm.Poisson(\"obs\", at.exp(eta), observed=y, dims=\"num_areas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Auto-assigning NUTS sampler...\n",
      "INFO:pymc:Auto-assigning NUTS sampler...\n",
      "Initializing NUTS using jitter+adapt_diag...\n",
      "INFO:pymc:Initializing NUTS using jitter+adapt_diag...\n",
      "Multiprocess sampling (4 chains in 4 jobs)\n",
      "INFO:pymc:Multiprocess sampling (4 chains in 4 jobs)\n",
      "NUTS: [tau_theta, tau_phi, theta, phi, beta0, beta1]\n",
      "INFO:pymc:NUTS: [tau_theta, tau_phi, theta, phi, beta0, beta1]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='8000' class='' max='8000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      100.00% [8000/8000 03:04<00:00 Sampling 4 chains, 0 divergences]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/conor/miniconda3/envs/pymc_non_dev/lib/python3.10/site-packages/scipy/stats/_continuous_distns.py:624: RuntimeWarning: overflow encountered in _beta_ppf\n",
      "  return _boost._beta_ppf(q, a, b)\n",
      "/Users/conor/miniconda3/envs/pymc_non_dev/lib/python3.10/site-packages/scipy/stats/_continuous_distns.py:624: RuntimeWarning: overflow encountered in _beta_ppf\n",
      "  return _boost._beta_ppf(q, a, b)\n",
      "/Users/conor/miniconda3/envs/pymc_non_dev/lib/python3.10/site-packages/scipy/stats/_continuous_distns.py:624: RuntimeWarning: overflow encountered in _beta_ppf\n",
      "  return _boost._beta_ppf(q, a, b)\n",
      "/Users/conor/miniconda3/envs/pymc_non_dev/lib/python3.10/site-packages/scipy/stats/_continuous_distns.py:624: RuntimeWarning: overflow encountered in _beta_ppf\n",
      "  return _boost._beta_ppf(q, a, b)\n",
      "Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 197 seconds.\n",
      "INFO:pymc:Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 197 seconds.\n"
     ]
    }
   ],
   "source": [
    "with bym_model:\n",
    "    idata = pm.sample()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f977084116a5d06351f0969c43dbc528856b1012f23a32e53dbf92034586142c"
  },
  "kernelspec": {
   "display_name": "Python 3.10.5 ('pymc_non_dev')",
   "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.10.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}