radon-v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a180963a",
   "metadata": {
    "tags": [],
    "toc": true
   },
   "source": [
    "<h1>Table of Contents<span class=\"tocSkip\"></span></h1>\n",
    "<div class=\"toc\"><ul class=\"toc-item\"><li><ul class=\"toc-item\"><li><span><a href=\"#Notebook-setup\" data-toc-modified-id=\"Notebook-setup-0.1\">Notebook setup</a></span></li><li><span><a href=\"#Radon-dataset:--home-radon-measurements,-per-county-soil-uranium-levels\" data-toc-modified-id=\"Radon-dataset:--home-radon-measurements,-per-county-soil-uranium-levels-0.2\">Radon dataset:  home radon measurements, per-county soil uranium levels</a></span></li><li><span><a href=\"#Plotting-basics-with-plotnine\" data-toc-modified-id=\"Plotting-basics-with-plotnine-0.3\">Plotting basics with plotnine</a></span></li></ul></li><li><span><a href=\"#Best-Practice:-preliminary-data-analysis\" data-toc-modified-id=\"Best-Practice:-preliminary-data-analysis-1\">Best Practice: preliminary data analysis</a></span></li><li><span><a href=\"#Best-Practice:-start-with-a-simple-model\" data-toc-modified-id=\"Best-Practice:-start-with-a-simple-model-2\">Best Practice: start with a simple model</a></span></li><li><span><a href=\"#Linear-regression-in-Stan\" data-toc-modified-id=\"Linear-regression-in-Stan-3\">Linear regression in Stan</a></span><ul class=\"toc-item\"><li><span><a href=\"#Fitting-Models-in-CmdStanPy\" data-toc-modified-id=\"Fitting-Models-in-CmdStanPy-3.1\">Fitting Models in CmdStanPy</a></span></li><li><span><a href=\"#Extracting-fit-information\" data-toc-modified-id=\"Extracting-fit-information-3.2\">Extracting fit information</a></span></li><li><span><a href=\"#Excercise\" data-toc-modified-id=\"Excercise-3.3\">Excercise</a></span></li><li><span><a href=\"#Visualizing-model-estimates-with-plotnine\" data-toc-modified-id=\"Visualizing-model-estimates-with-plotnine-3.4\">Visualizing model estimates with plotnine</a></span></li></ul></li><li><span><a href=\"#Best-Practice:-posterior-predictive-checks\" data-toc-modified-id=\"Best-Practice:-posterior-predictive-checks-4\">Best Practice: posterior predictive checks</a></span><ul class=\"toc-item\"><li><span><a href=\"#Adding-county-uranium-as-a-predictor\" data-toc-modified-id=\"Adding-county-uranium-as-a-predictor-4.1\">Adding county uranium as a predictor</a></span></li><li><span><a href=\"#Discussion\" data-toc-modified-id=\"Discussion-4.2\">Discussion</a></span></li></ul></li><li><span><a href=\"#Multilevel-Regression\" data-toc-modified-id=\"Multilevel-Regression-5\">Multilevel Regression</a></span><ul class=\"toc-item\"><li><span><a href=\"#Modeling-the-regression-intercept-term\" data-toc-modified-id=\"Modeling-the-regression-intercept-term-5.1\">Modeling the regression intercept term</a></span></li><li><span><a href=\"#Writing-the-Stan-model\" data-toc-modified-id=\"Writing-the-Stan-model-5.2\">Writing the Stan model</a></span></li><li><span><a href=\"#Visualizations\" data-toc-modified-id=\"Visualizations-5.3\">Visualizations</a></span></li><li><span><a href=\"#Posterior-predictive-checks\" data-toc-modified-id=\"Posterior-predictive-checks-5.4\">Posterior predictive checks</a></span></li><li><span><a href=\"#Next-Steps\" data-toc-modified-id=\"Next-Steps-5.5\">Next Steps</a></span></li></ul></li><li><span><a href=\"#Discussion\" data-toc-modified-id=\"Discussion-6\">Discussion</a></span></li><li><span><a href=\"#References-and-Resources\" data-toc-modified-id=\"References-and-Resources-7\">References and Resources</a></span></li><li><span><a href=\"#Appendix-A:-Linear-regression-review-(chapters-2-and-3-of-Gelman-and-Hill)\" data-toc-modified-id=\"Appendix-A:-Linear-regression-review-(chapters-2-and-3-of-Gelman-and-Hill)-8\">Appendix A: Linear regression review (chapters 2 and 3 of Gelman and Hill)</a></span></li></ul></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25194d5e",
   "metadata": {},
   "source": [
    "### Notebook setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fe0ea6e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import all libraries used in this notebook\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from cmdstanpy import CmdStanModel\n",
    "\n",
    "# gis data\n",
    "import geopandas as gpd\n",
    "import libpysal as sa\n",
    "\n",
    "# plotting libs\n",
    "import matplotlib.pyplot as plt\n",
    "import splot as splt\n",
    "import plotnine as p9\n",
    "%matplotlib inline\n",
    "\n",
    "# suppress plotnine warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# setup plotnine look and feel\n",
    "p9.theme_set(\n",
    "  p9.theme_grey() + \n",
    "  p9.theme(text=p9.element_text(size=10),\n",
    "        plot_title=p9.element_text(size=14),\n",
    "        axis_title_x=p9.element_text(size=12),\n",
    "        axis_title_y=p9.element_text(size=12),\n",
    "        axis_text_x=p9.element_text(size=8),\n",
    "        axis_text_y=p9.element_text(size=8)\n",
    "       )\n",
    ")\n",
    "xlabels_90 = p9.theme(axis_text_x = p9.element_text(angle=90, hjust=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2534095e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep notebook outputs clean - demos only\n",
    "import logging\n",
    "logging.getLogger('cmdstanpy').setLevel(logging.CRITICAL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44fef876-55c5-4dc5-b383-da2b6a658806",
   "metadata": {},
   "source": [
    "**Best Practice:  avoid meaningless precision**\n",
    "\n",
    "The precision of an estimate is inversely proportional to square root of the amount of data.\n",
    "For the Minnesota data, only 2 decimal places is necessary.\n",
    "To change the default print behavior for pandas DataFrames, we use the pandas global option [display.precision](https://pandas.pydata.org/docs/user_guide/options.html#frequently-used-options)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18d5ef6-988e-4686-b9b9-7ce921e9084f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pd.set_option('display.precision', 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5980f3b",
   "metadata": {},
   "source": [
    "###  Radon dataset:  home radon measurements, per-county soil uranium levels\n",
    "\n",
    "The data comes from US Environmental Protection Agency (EPA) surveys at the state and national level\n",
    "carried out in the 1990s.\n",
    "It is available from the Gelman and Hill [ARM website](http://www.stat.columbia.edu/~gelman/arm/examples/radon/), together with the R scripts used to produce the examples in the book.\n",
    "\n",
    "The goal of this study is to estimate home radon levels for all counties in the US.\n",
    "Localizing data to the county provides guidelines for individual homeowners as well\n",
    "as public health and planning agencies.\n",
    "\n",
    "**Raw data**\n",
    "\n",
    "The study data is in two separate files: one for the radon home survey, one for county soil uranium levels.\n",
    "These are distributed as\n",
    "[srrs2.dat](http://www.stat.columbia.edu/~gelman/arm/examples/radon/srrs2.dat)\n",
    "and \n",
    "[cty.dat](http://www.stat.columbia.edu/~gelman/arm/examples/radon/cty.dat)\n",
    "respectively.\n",
    "\n",
    "More data is available in folder http://www.stat.columbia.edu/~gelman/arm/examples/radon-complete.\n",
    "There are a total of 63K home radon measurements from 3000 US counties and tribal lands.\n",
    "There are roughly 4K records from tribal lands records. These are not coded by US county but\n",
    "instead coded with by EPA region, thus they cannot be easily integrated with the rest of the data.\n",
    "\n",
    "The per-county measurements follow the population density; there are few or no measurements\n",
    "for sparsely populated counties, i.e. rural counties and correspondingly more for metropolitan counties.\n",
    "\n",
    "\n",
    "**Processed data**\n",
    "\n",
    "For this analysis, we combine the radon data, soil-uranium data, and GIS data from the US Census bureau, which provides clean, standardized county names.\n",
    "The [data prep notebook](data_prep_radon_uranium_geodata.ipynb) shows the processing details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "139dd19e",
   "metadata": {},
   "source": [
    "**USA data**\n",
    "\n",
    "Radon, uranium levels for all US counties, cross-indexed by county name, FIPS\n",
    "\n",
    "+ [data/us_radon.csv](data/us_radon.csv) contains all home radon measurements.\n",
    "+ [data/us_counties.csv](data/us_counties.csv) contains county-level counts of homes, soil uranium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e225acf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "us_radon = pd.read_csv(os.path.join('data','us_radon.csv'))\n",
    "print(f'number of houses: {len(us_radon)}')\n",
    "us_radon.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c62fba7",
   "metadata": {},
   "source": [
    "The object of this study is home radon levels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caed4a34",
   "metadata": {},
   "outputs": [],
   "source": [
    "us_radon.radon.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ef6caac",
   "metadata": {},
   "source": [
    "On the log scale, this distribution looks like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ced8275",
   "metadata": {},
   "outputs": [],
   "source": [
    "us_radon.log_radon.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98b40d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "us_counties = pd.read_csv(os.path.join('data','us_counties.csv'))\n",
    "print(f'number of counties: {len(us_counties)}')\n",
    "us_counties.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8908792",
   "metadata": {},
   "source": [
    "On average, there are few observations per county."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eae0306",
   "metadata": {},
   "outputs": [],
   "source": [
    "us_counties.homes.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e97232cd",
   "metadata": {},
   "source": [
    "**Minnesota**\n",
    "\n",
    "The Minnesota data files are used to replicate Gelman and Hill analysis chapter 12.\n",
    "These are cross-indexed by FIPS code, county name, and a synthetic county_id.\n",
    "\n",
    "+ [data/mn_radon.csv](data/mn_radon.csv) contains the individual home radon measurements.\n",
    "+ [data/mn_counties.csv](data/mn_counties.csv) contains county-level data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3146fc97-3ea9-4216-93e2-abc6e433b515",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_radon = pd.read_csv(os.path.join('data','mn_radon.csv'))\n",
    "print(f'number of houses: {len(mn_radon)}')\n",
    "mn_radon.head(7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703b40ce",
   "metadata": {},
   "source": [
    "There are 87 counties in Minnesota, but the EPA radon dataset doesn't contain measurements for 2 of these counties.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2def98f7-b871-4a21-a73d-ae2414efa1ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_counties = pd.read_csv(os.path.join('data','mn_counties.csv'))\n",
    "print(f'number of counties: {len(mn_counties)}')\n",
    "mn_counties.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6000158e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_radon.radon.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16255967",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_counties.homes.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c4cfd0f-80dc-49af-8c3f-0550fa912dde",
   "metadata": {},
   "source": [
    "### Plotting basics with plotnine\n",
    "\n",
    "For plotting we use the plotnine package, which\n",
    "is a Python implementation of a grammar of graphics based on the R ggplot2 package.\n",
    "The grammar allows users to compose plots from one or more layers.\n",
    "A good resource is [Data Visualization with Plotnine](https://f0nzie.github.io/rmarkdown-python-plotnine/),\n",
    "a Python translation of [R for Data Science](https://r4ds.had.co.nz/).\n",
    "\n",
    "A grammar of graphics defines a plot in terms of:\n",
    "\n",
    "- A dataset in the form of a pandas.DataFrame.\n",
    "- A set of mappings from dataset variables to graph elements called \"aesthetics\".\n",
    "- A coordinate system, default Cartesian, x,y axes, where x is on the horizontal and y is on the vertical axis.\n",
    "- A facet specification based on a categorical variable which results in per-category subplots, default None.\n",
    "- One or more layers, each layer takes as arguments:\n",
    "   + a dataset and aesthetic mapping - by default, the plot dataset and mappings are used\n",
    "   + one geometric object (\"geom\") - the geometric building blocks of the plot, e.g., point, line, polygon.\n",
    "   + one statistical transformation (\"stat\"), default \"identity\"\n",
    "   + one position adjustment, default \"identity\"\n",
    "\n",
    "A plot layer is based on a geom, or a stat paired with a geom.\n",
    "As a first example, we create a plot which contains a single layer which is a \n",
    "[plotnine geom_histogram](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_histogram.html).\n",
    "We provide the minimum plot specification:  the geom, data, and mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0bde7b-2065-41db-b353-ac2f8f1ab7be",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "p9.ggplot(data=mn_radon, mapping=p9.aes(x='log_radon')) + p9.geom_histogram()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2ef81de-d728-471b-83c8-d70e86273d25",
   "metadata": {},
   "source": [
    "Each layer works off a dataset and mapping which can be specified either\n",
    "as arguments to the [plotnine.ggplot](https://plotnine.readthedocs.io/en/stable/generated/plotnine.ggplot.html#plotnine.ggplot) object,\n",
    "as in the above example, or as arguments to the geom object.\n",
    "An alternative way to create the same plot is:\n",
    "\n",
    "```python\n",
    "p9.ggplot() + p9.geom_histogram(data=mn_radon, mapping=p9.aes(x='log_radon'))\n",
    "```\n",
    "\n",
    "In this example, the only aesthetic is \"x\" - the position of data along the x-axis;\n",
    "the y axis are the values of the histogram bins.\n",
    "Different geoms require different mappings. \n",
    "Were this a plot such as a scatterplot, mapping for value \"y\" would be required as well.\n",
    "Other common mappings are:\n",
    "\n",
    "- \"shape\" - of points\n",
    "- \"linetype\"\n",
    "- \"color\" - color of lines and outlines\n",
    "- \"fill\" - color inside shapes\n",
    "- \"size\"\n",
    "\n",
    "All of these can be mapped to features of the data, so that a single plot layer\n",
    "can express many aspects of the dataset.\n",
    "To see how this works, we add information to the histogram by using aesthetics\n",
    "\"color\" and \"fill\" to factor radon measurements by floor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a36907c2-bfea-4270-ba23-9fa81d89c3d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# overlay histograms\n",
    "(p9.ggplot(data=mn_radon, mapping=p9.aes(x='log_radon', color='factor(floor)', fill='factor(floor)'))\n",
    "    + p9.geom_histogram(alpha=0.7, binwidth=0.2)\n",
    "    + p9.scale_color_manual(['darkgreen','darkblue'])\n",
    "    + p9.scale_fill_manual(['orange','violet'])\n",
    "    + p9.theme(figure_size=(6,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d453315-fc0f-4ffe-b155-91c049dbd34c",
   "metadata": {},
   "source": [
    "In analyzing and plotting the data, the county name is a categorical value.\n",
    "We update the `mn_radon` dataframe accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4c1b447-4ea0-4f0a-8870-f8fe81d72faa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use home radon data county name as a categorical variable\n",
    "mn_radon['county'] = (mn_radon['county'].astype('category', copy=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c74236-5102-4fea-b89c-6bd5068a879e",
   "metadata": {},
   "source": [
    "## Best Practice: preliminary data analysis\n",
    "\n",
    "A sometimes overlooked issue when doing model criticism and model comparison is the fact that we are evaluating the fit of the model _to the data_.  The size and shape of the data informs our choice of model.  Finally, the data collected is not always the data expected.  Therefore we start with plots and summaries of the raw data.\n",
    "\n",
    "**First questions: amount of data, variable of interest**\n",
    "\n",
    "How much data is there for Minnesota?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "601ad455",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'number of houses: {len(mn_radon)}')\n",
    "print(f'number of counties: {len(mn_counties)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c2c9095-215d-4598-aa05-bdac057cd727",
   "metadata": {},
   "source": [
    "The goal of our analysis is to estimate home radon levels; therefore the outcome variable of interest is `log_radon`.\n",
    "We use the [pandas.DataFrame.describe](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.describe.html) function to get summary statistics over the observed outcome `log_radon`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfcfd10d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'log_radon summary statistics\\n{mn_radon[\"log_radon\"].describe()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07cf3160",
   "metadata": {},
   "source": [
    "**Relationship between radon and floor**\n",
    "\n",
    "As the radon pathways diagram shows, radon comes from the soil, therefore the floor level on which the measurement was taken should be a good predictor of the observed radon level. This is coded as \"0\" for basement and \"1\" for ground floor level.  Most of the observations in the survey were taken on the basement level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d08f36",
   "metadata": {},
   "outputs": [],
   "source": [
    "pct_1 = round((mn_radon.floor.sum() / len(mn_radon) * 100))\n",
    "pct_0= round(100 - pct_1)\n",
    "print(f'floor 0: {pct_0}%\\nfloor 1: {pct_1}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0c1b556-ddfd-4949-8c0a-c453b3455a9d",
   "metadata": {},
   "source": [
    "Plotting the histogram of raw counts of number of observations by floor clearly shows the differing amount of per-floor data, but the trend in log radon levels is unclear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b46709d-05d7-45e9-8569-989687dfcd4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# overlay histograms\n",
    "(p9.ggplot(data=mn_radon, mapping=p9.aes(x='log_radon', color='factor(floor)', fill='factor(floor)'))\n",
    "    + p9.geom_histogram(alpha=0.7, binwidth=0.2)\n",
    "    + p9.scale_color_manual(['darkgreen','blue'], name='floor')\n",
    "    + p9.scale_fill_manual(['orange','violet'], name='floor')\n",
    "    + p9.xlab(\"log radon levels\")\n",
    "    + p9.theme(figure_size=(6,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b54bea7c",
   "metadata": {},
   "source": [
    "The [plotnine stat_density](https://plotnine.readthedocs.io/en/stable/generated/plotnine.stats.stat_density.html)\n",
    "shows per-floor trends, but obscures the difference in amounts of observations.\n",
    "This is an example of a layer constructed from a plotnine stat paired with a geom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10294c2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p9.ggplot(data=mn_radon, mapping=p9.aes(x='log_radon', color='factor(floor)'))\n",
    "    + p9.stat_density(geom='line')\n",
    "    + p9.scale_color_manual(['darkorange','purple'], name='floor')\n",
    "    + p9.xlab(\"log radon levels\")\n",
    "    + p9.theme(figure_size=(6,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7aba069",
   "metadata": {},
   "source": [
    "The [plotnine geom_point](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_point.html) plots two variables as (x, y) points.\n",
    "A scatterplot of points (`floor`, `log_radon`) will only have 2 distinct x-axis values:  0 and 1.\n",
    "The [plotnine geom_jitter](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_jitter.html) adds jitter to the (x, y) points,\n",
    "which reduces the amount of overplotting and therefore allows for a better visualization of the amount of data being plotted.\n",
    "Therefore we use the latter to visualize the differences between the the radon measurements by floor.\n",
    "\n",
    "Because `log_radon` is the outcome variable of interest for this example, whenever possible, we plot it on the y axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4df24560",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_radon_floor = (p9.ggplot(data=mn_radon, mapping=p9.aes(x='floor', y='log_radon')) \n",
    "    + p9.geom_jitter(width=0.1, alpha=0.5, fill='orange', color='darkred')\n",
    "    + p9.scale_x_continuous(breaks=range(0,2), minor_breaks=[])\n",
    "    + p9.ggtitle(\"Radon measurements by floor\")\n",
    "    + p9.theme(figure_size=(4,4))\n",
    ")\n",
    "plot_radon_floor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00b985d5",
   "metadata": {},
   "source": [
    "**County-level information**\n",
    "\n",
    "Because the most home radon measurements were taken at the basement floor level and because most of the counties have relatively few home measurements taken, there are many counties where all measurements are from the basement floor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b00ab2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f'Number of counties: {mn_counties.shape[0]}')\n",
    "print(f'Counties with measurements from floor 0: {mn_radon[mn_radon[\"floor\"]==0][\"county\"].nunique()}')\n",
    "print(f'Counties with measurements from floor 1: {mn_radon[mn_radon[\"floor\"]==1][\"county\"].nunique()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9540c15",
   "metadata": {},
   "source": [
    "At the county level we have many home radon measurements from the relatively few counties with metropolitan areas, and very few home radon measurements from the rest.\n",
    "A basic way to see this distribution of homes per county is a histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d5ae9ca-84d8-4c47-8307-5551e53ae542",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(p9.ggplot()\n",
    "    + p9.geom_histogram(data=mn_counties, mapping=p9.aes(x='homes'), bins=40)\n",
    "    + p9.xlab(\"homes per county\")\n",
    "    + p9.ylab(\"counties per bin\")\n",
    "    + p9.theme(figure_size=(12,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05b1b869-6bde-4604-ab68-639b3ec213ec",
   "metadata": {},
   "source": [
    "**Sort order: homes per county, ascending**\n",
    "\n",
    "The amount of data per county directly affects the precision of our estimates.\n",
    "\n",
    "First we create a list of county names in order of number of homes per county, ascending."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45d870a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_asc = mn_counties.sort_values(by='homes').reset_index(drop=True).county.to_frame().reset_index()\n",
    "obs_asc.rename(columns={'index':'sort_homes_asc'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "313f6960-d80c-4a4d-a630-df021382be0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_radon = mn_radon.merge(obs_asc, on='county')\n",
    "mn_radon.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd440dc0-a54e-48ff-b0aa-a254f6c111a5",
   "metadata": {},
   "source": [
    "**Boxplot visualizations**\n",
    "\n",
    "Another way to visualize the amount and spread of data per county is by using a\n",
    "[plotnine.geom_boxplot](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_boxplot.html#plotnine.geoms.geom_boxplot)\n",
    "to generate a [box_and_whiskers plot](https://en.wikipedia.org/wiki/Box_plot) for each set of per-county radon measurements.\n",
    "The box encloses the central 25% - 75% quantiles; this is the \n",
    "[interquartile range (IQR)](https://en.wikipedia.org/wiki/Interquartile_range).\n",
    "The the whiskers extend to the values which are a distance of 1.5 the IQR, and values beyond that are plotted as points - these are the outliers.\n",
    "\n",
    "Setting the width of the box to be proportional to the square root of the number of observations\n",
    "shows amount of data per county, as well as its spread."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "422fe4d3",
   "metadata": {},
   "source": [
    "First we set up a dataframe whose index provides the sort order for the boxplot,\n",
    "which is then used as an argument to the plotnine [scales.scale_x_discrete](https://plotnine.readthedocs.io/en/stable/generated/plotnine.scales.scale_x_discrete.html#plotnine.scales.scale_x_discrete) object, then we set up the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd9bcf5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sort_asc = mn_radon[['county', 'sort_homes_asc']].drop_duplicates().sort_values(by='sort_homes_asc').reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f4d8b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_cnty_by_pop = (p9.ggplot(data=mn_radon, mapping=p9.aes(x='county', y='log_radon'))\n",
    "    + p9.geom_boxplot(width=2, varwidth=True, outlier_alpha=0.4)\n",
    "    + p9.scale_x_discrete(limits=obs_asc['county'], expand=(0,3))\n",
    "    + p9.ggtitle(\"Counties ordered by number of observations per county\")\n",
    "    + p9.ylab(\"range of home radon measurements\")\n",
    "    + xlabels_90\n",
    "    + p9.theme(figure_size=(20,6))\n",
    ")\n",
    "p_cnty_by_pop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37ea52af",
   "metadata": {},
   "source": [
    "**Relationship between home radon and county-level soil uranium**\n",
    "\n",
    "At the county-level, we have information on the soil uranium level.  We plot the number of observations by soil uranium.  The points on the x-axis line up earlier histogram, but instead of bars, we have a series of points showing the different log_uranium levels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2310936d",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p9.ggplot(data=mn_counties, mapping=p9.aes(x='homes', y='log_uranium'))\n",
    "    + p9.geom_point(fill='orange', color='darkred')\n",
    "    + p9.geom_text(data=mn_counties[mn_counties['homes']>25],\n",
    "                   mapping=p9.aes(label='county'),\n",
    "                   size=8, nudge_x=4, nudge_y=0.1)\n",
    "    + p9.xlab(\"observations per county\") + p9.ylab(\"county soil log_uranium\")\n",
    "    + p9.theme(figure_size=(12,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "638f2e05",
   "metadata": {},
   "source": [
    "We plot the relationship between soil uranium level and the home radon measurement.\n",
    "We use plotnine's [facet_grid](https://plotnine.readthedocs.io/en/stable/generated/plotnine.facets.facet_grid.html#plotnine.facets.facet_gridplot) to get side-by-side per-floor plots.\n",
    "Because the soil uranium level measurement is the same for all homes in a county, for counties with many houses, i.e., metropolitan areas, the plot shows distinct vertical bands.\n",
    "\n",
    "Comparing the information in two side-by-side plots is difficult.  We have established that there are fewer measurements taken on the ground floor than on the basement level.   But we can't see whether or not the home radon measurements are consistently lower when taken on the ground floor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f108f7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plotnine wants all data in a single table.\n",
    "mn_radon = mn_radon.merge(mn_counties[['FIPS', 'log_uranium']], on='FIPS')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d22197",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p9.ggplot(data=mn_radon, mapping=p9.aes(x='log_uranium', y='log_radon'))\n",
    "     + p9.geom_point(alpha=0.9, fill='orange', color='darkred')\n",
    "     + p9.facet_grid(facets='~ floor', labeller='label_both')\n",
    "     + p9.xlab(\"county-level soil log_uranium\") + p9.ylab(\"home log_radon\") \n",
    "     + p9.theme(figure_size=(12,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7b3f1e8",
   "metadata": {},
   "source": [
    "Alternatively, we can use color to indicate floor, basement orange, ground floor purple, and add jitter to overplot the data.  Taken together, it's not clear whether or not the blue dots are generally lower than the orange ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf05dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p9.ggplot()\n",
    "    + p9.geom_jitter(data=mn_radon[mn_radon['floor']==0],\n",
    "                    mapping=p9.aes(x='log_uranium', y='log_radon'), \n",
    "                    width=0.01, alpha=0.7, fill='orange', color='darkred')\n",
    "    + p9.geom_jitter(data=mn_radon[mn_radon['floor']==1],\n",
    "                    mapping=p9.aes(x='log_uranium', y='log_radon'), \n",
    "                    width=0.01, alpha=0.7, fill='purple', color='darkblue')\n",
    "    + p9.xlab(\"county-level soil log_uranium\")\n",
    "    + p9.ylab(\"home log_radon\")\n",
    "    + p9.theme(figure_size=(8,4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1ec30a9",
   "metadata": {},
   "source": [
    "To see whether or not the soil uranium level might be a good predictor of the home radon measurements,\n",
    "we can repeat the above boxplot, this time ordering the counties on the x-axis by\n",
    "the per-county activity levels, ordered by uranium level per county, descending.\n",
    "\n",
    "\n",
    "Given the sparse data, the resulting plot is inconclusive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf568fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "uranium_desc = mn_counties.sort_values(by='log_uranium', ascending=False).reset_index()\n",
    "\n",
    "p_cnty_by_u = (p9.ggplot(data=mn_radon, mapping=p9.aes(x='county',y='log_radon'))\n",
    "    + p9.geom_boxplot(width=2, varwidth=True, outlier_alpha=0.4)\n",
    "    + p9.scale_x_discrete(limits=uranium_desc['county'], expand=(0,1))\n",
    "    + p9.ggtitle(\"Counties ordered by soil uranium high (left) to low (right)\")\n",
    "    + p9.ylab(\"range of home radon measurements\")\n",
    "    + xlabels_90\n",
    "    + p9.theme(figure_size=(20,6))\n",
    ")\n",
    "p_cnty_by_u"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57e541cd",
   "metadata": {},
   "source": [
    "**Spatial Data**\n",
    "\n",
    "Geographic information systems (GIS) data is any item which has a geographic location,\n",
    "either a single point or a set of bounding polygons.\n",
    "In order to manage, analyze, and visualize GIS data,\n",
    "we use specialized packages which can do the geographic math.\n",
    "In this notebook we use the following packages:\n",
    "\n",
    "* GeoPandas - manages a set of GIS records in tabular format\n",
    "* libpysal - spatial analysis package which can analyze distance between locations\n",
    "* splot - plots for libpysal objects\n",
    "* plotnine - object geom_map\n",
    "\n",
    "Cartographic data (maps) are encoded as a set of records, one per map region.\n",
    "The shapefile format is an open specification used to insure interoperatility among GIS software packages.\n",
    "When items in a dataset contain location labels, it is necessary to obtain a set of shapefiles\n",
    "for the corresponding map.\n",
    "\n",
    "The shapefiles for US counties are available from the US Census Bureau. For this analysis, we are using shapefiles where the boundary information is specified with the lowest possible resolution; this greatly speeds up analysis and plotting.\n",
    "These can be downloaded via URL: https://www2.census.gov/geo/tiger/GENZ2018/shp/cb_2018_us_county_20m.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61edc826",
   "metadata": {},
   "outputs": [],
   "source": [
    "shpfile = os.path.join('geo_data','cb_2018_us_county_20m', 'cb_2018_us_county_20m.shp')\n",
    "us_geodata = gpd.read_file(shpfile)\n",
    "mn_geodata = us_geodata[us_geodata.STATEFP == '27']\n",
    "# GEOID should be numeric\n",
    "mn_geodata = mn_geodata.astype({'GEOID': 'int32'}, copy=False)\n",
    "mn_geodata.drop(columns=['STATEFP', 'COUNTYFP', 'COUNTYNS', 'AFFGEOID', 'LSAD', 'ALAND', 'AWATER'], inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3116f789",
   "metadata": {},
   "source": [
    "To visualize county-level data, we add information to the GeoPandas dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f24a97a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_geodata = mn_geodata.merge(us_counties[['homes', 'uranium', 'FIPS']], \n",
    "                              how='inner', left_on='GEOID', right_on='FIPS')\n",
    "mn_geodata = mn_geodata.astype({'homes': 'int32'}, copy=False)\n",
    "mn_geodata.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "517bc531",
   "metadata": {},
   "source": [
    "Now we can use maps to visualize county-level information.\n",
    "The fill color indicates the soil-uranium level, dark for low levels, lighter for high levels.\n",
    "The text label is the number of observations per county."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b864eb67",
   "metadata": {},
   "outputs": [],
   "source": [
    "county_pts = mn_geodata.geometry.centroid.to_crs(mn_geodata.crs)\n",
    "(p9.ggplot()\n",
    " + p9.geom_map(data=mn_geodata, mapping=p9.aes(fill='uranium'))\n",
    " + p9.geom_text(mapping=p9.aes(x=county_pts.x, y=county_pts.y, label=list(mn_geodata.homes)),\n",
    "                color='white', size=10, ha='center')\n",
    " + p9.scale_fill_gradient()\n",
    " + p9.theme(figure_size=(10, 10))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3be912",
   "metadata": {},
   "source": [
    "**Preliminary Analysis of Minnesota Dataset Findings**\n",
    "\n",
    "* 83% of the data are measurements taken on the basement level.\n",
    "\n",
    "* 70% of the counties (60 out of 85) have observations from both floors 0 and 1, the remaining 30% only have observations from floor 0 (basement).\n",
    "\n",
    "* For most counties, there are fewer than 10 observations; 8 counties in metropolitan areas account for over half of the observations.\n",
    "\n",
    "* The counties with the highest soil uranium levels do not have a lot of observations per county.\n",
    "\n",
    "* Within each county, the range of radon measurements is very wide."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c36542c",
   "metadata": {},
   "source": [
    "## Best Practice: start with a simple model\n",
    "\n",
    "Starting from a simple model ensures that there is a good baseline\n",
    "against which to measure performance.\n",
    "\n",
    "**Two regular linear regression models**\n",
    "\n",
    "For this case study, \n",
    "the baseline model is a regular (non-multilevel) linear regression model,\n",
    "[see Appendix A](#regression),\n",
    "where the outcome _y_ is the home log radon level and the predictor _x_ is the floor\n",
    "on which the measurement was taken.\n",
    "We consider two possible models:   **complete pooling** and **no-pooling**.\n",
    "\n",
    "The **complete pooling** model estimates a single intercept term for the regression.\n",
    "\n",
    "$\n",
    "\\mathrm{log\\_radon}_i = \\alpha \\, + {\\beta}\\,\\mathrm{floor}_i + {\\epsilon}_i\n",
    "$\n",
    "\n",
    "The **no pooling** model estimates a per-county intercept term.\n",
    "The intercept term $\\alpha$ is a vector of size $\\mathrm{J}$, the number of counties.\n",
    "\n",
    "$\n",
    "\\mathrm{log\\_radon}_i = \\alpha_{j[i]} \\, + {\\beta}\\,\\mathrm{floor}_i + {\\epsilon}_i \\ \\ \\ \\\n",
    "$\n",
    "where $j = 1 \\ldots \\mathrm{J}$.\n",
    "\n",
    "Gelman and Hill use the notation $\\alpha_{j[i]}$ to denote the element of $j$ corresponding to observation $i$,\n",
    "arguing that this notation better reflects the structure of the data.\n",
    "\n",
    "Because the floors are coded\n",
    "$\\mathrm{floor}_{\\mathrm{basement}} = 0$,  $\\mathrm{floor}_{\\mathrm{ground}} = 1$,\n",
    "for basement measurements, the observed outcome $\\mathrm{log\\_radon}$ is just\n",
    "the intercept term plus measurement error.\n",
    "\n",
    "_Note:  we have a lot of options when it comes to coding floor information._\n",
    "\n",
    "We can use [plotnine.geom.geom_smooth(method=\"lm\")](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_smooth.html#plotnine.geoms.geom_smooth)\n",
    "to visualize both the complete pooling and no pooling-models.\n",
    "This adds a layer with the mean regression line, and grey margins which indicate the amount of variance.\n",
    "\n",
    "For the complete pooling model, we add this geom to the saved plot `plot_radon_floor`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3017bb5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "p1 = (plot_radon_floor\n",
    "      +  p9.geom_smooth(method='lm')\n",
    "      +  p9.ggtitle(\"regression log_radon on floor, all counties in MN\")\n",
    "     )      \n",
    "p1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b6f4b9",
   "metadata": {},
   "source": [
    "For the no-pooling model, we facet the radon by floor plot.\n",
    "When combined with `geom_smooth`, every facet has its own regression model.\n",
    "As noted above, most of the measurements are taken at the basement level and there are 25 counties without any measurements taken on floor 1. For these counties, in the faceted plot the geom_smooth trend line is absent; i.e., we cannot estimate the slope of the regression line. For counties with measurements on both floors, in a few cases, e.g., Todd and Carlton counties, the regression line between floor 0 and 1 has a positive slope, which a) goes against the estimate from the complete pooling model, and b) goes against what we know about how radon enters the home."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5714fb4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p9.ggplot(data=mn_radon, mapping=p9.aes('floor', 'log_radon'))\n",
    "     + p9.geom_jitter(width=0.05)\n",
    "     + p9.geom_smooth(method='lm')\n",
    "     + p9.facet_wrap('county')\n",
    "     + p9.ggtitle(\"per-county regression log_radon on floor, counties ordered by observations, ascending\")\n",
    "     + p9.scale_x_continuous(breaks=range(0,2), minor_breaks=[])\n",
    "     + p9.scales.ylim(-3, 4)  # same limits as complete pooling\n",
    "     + p9.theme(figure_size=(18,20))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b29eb5cd",
   "metadata": {},
   "source": [
    "## Linear regression in Stan\n",
    "\n",
    "The complete-pooling model corresponds to the simplest linear regression model in the\n",
    "[Stan User's Guide](https://mc-stan.org/docs/stan-users-guide/linear-regression.html).\n",
    "This model is in file [radon_cp.stan](stan/radon_cp.stan).\n",
    "It adds the following to the model in the Stan User's Guide:\n",
    "\n",
    "+ All model parameters have [weakly informative priors](https://github.com/stan-dev/stan/wiki/Prior-Choice-Recommendations).\n",
    "\n",
    "+ The `generated quantities` program block is used to generate a new sample `y_rep`, which is $\\widetilde{y}$, yet to be observed data.\n",
    "By using the model parameter estimates with Stan's [random number generator probability functions](https://mc-stan.org/docs/functions-reference/distributions-prng.html)\n",
    "we generate a new dataset which captures both sampler error and estimation error; see the User's Guide\n",
    "chapter on [Posterior Predictive Sampling](https://mc-stan.org/docs/stan-users-guide/posterior-prediction.html) for details.\n",
    "In the next section we show how to use `y_rep` to test model correctness.\n",
    "\n",
    "\n",
    "The complete-pooling model doesn't use county information.\n",
    "If there is no real difference between the radon levels observed across the different counties,\n",
    "this is model will provide useful estimates.\n",
    "\n",
    "```stan\n",
    "data {\n",
    "  int<lower=1> N;\n",
    "  vector[N] x;\n",
    "  vector[N] y;\n",
    "}\n",
    "parameters {\n",
    "  real alpha;\n",
    "  real beta;\n",
    "  real<lower=0> sigma;\n",
    "}\n",
    "model {\n",
    "  y ~ normal(alpha + beta * x, sigma);\n",
    "  alpha ~ normal(0, 10);\n",
    "  beta ~ normal(0, 10);\n",
    "  sigma ~ normal(0, 10);\n",
    "}\n",
    "generated quantities {\n",
    "  array[N] real y_rep = normal_rng(alpha + beta * x, sigma);\n",
    "}\n",
    "```\n",
    "\n",
    "The no-pooling model is in file [radon_np.stan](stan/radon_np.stan).\n",
    "It differs from the complete pooling model in that\n",
    "\n",
    "- The input data includes the number of counties `J` and a vector of county ids for each observation.\n",
    "- The intercept parameter `alpha` is coded as a vector of size `J`.\n",
    "\n",
    "```stan\n",
    "data {\n",
    "  int<lower=1> N;  // observations\n",
    "  int<lower=1> J;  // counties\n",
    "  array[N] int<lower=1, upper=J> county;\n",
    "  vector[N] x;     // floor\n",
    "  vector[N] y;     // radon\n",
    "}\n",
    "parameters {\n",
    "  vector[J] alpha;\n",
    "  real beta;\n",
    "  real<lower=0> sigma;\n",
    "}\n",
    "model {\n",
    "  y ~ normal(alpha[county] + beta * x, sigma);  \n",
    "  alpha ~ normal(0, 10);\n",
    "  beta ~ normal(0, 10);\n",
    "  sigma ~ normal(0, 10);\n",
    "}\n",
    "generated quantities {\n",
    "  array[N] real y_rep = normal_rng(alpha[county] + beta * x, sigma);\n",
    "}\n",
    "```\n",
    "\n",
    "**Stan Programming Language:  Multiple Indexes**\n",
    "\n",
    "The no-pooling model uses Stan's [multiple indexing](https://mc-stan.org/docs/reference-manual/language-multi-indexing.html) syntax to specify the likelihood of y and to compute y_rep\n",
    "\n",
    "```stan\n",
    "y ~ normal(alpha[county] + beta * x, sigma);\n",
    "```\n",
    "\n",
    "+ vector `alpha` has size `J`\n",
    "+ integer array `county` is size `N`\n",
    "+ the values of `county` are between 1 and `J`, inclusive, therefore `county` is a valid set of indexes into `alpha`\n",
    "+ the expression `alpha[county]` is a vector of size `N`\n",
    "+ the expression `alpha[county] + beta * x` is a vector of size `N`\n",
    "\n",
    "This satisfies the size constraints on the arguments to the vectorized normal distribution function.\n",
    "\n",
    "\n",
    "\n",
    "**Best Practice:  Use File Naming Conventions**\n",
    "\n",
    "In this case study\n",
    "\n",
    "+ suffix \"cp\" stands for \"complete pooling\"\n",
    "+ suffix \"np\" stands for \"no pooling\"\n",
    "+ suffix \"pp\" stands for \"partial pooling\"\n",
    "\n",
    "### Fitting Models in CmdStanPy\n",
    "\n",
    "We create [CmdStanModel](https://mc-stan.org/cmdstanpy/api.html#cmdstanmodel) objects for both the complete-pooling and no-pooling models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49da696c",
   "metadata": {},
   "outputs": [],
   "source": [
    "complete_pooling_model = CmdStanModel(stan_file=os.path.join('stan', 'radon_cp.stan'))\n",
    "no_pooling_model = CmdStanModel(stan_file=os.path.join('stan', 'radon_np.stan'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c1d3c5e",
   "metadata": {},
   "source": [
    "We assemble a Python dictionary which contains the definitions of the data block variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5124faf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_radon_data = {\"N\": len(mn_radon), \n",
    "              \"x\": mn_radon.floor.astype(float), \n",
    "              \"y\": mn_radon.log_radon,\n",
    "              \"J\":len(mn_counties),\n",
    "              \"uranium\" : mn_counties.log_uranium,\n",
    "              \"county\" : mn_radon.county_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc62fc0b",
   "metadata": {},
   "source": [
    "We call the model's [sample](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanModel.sample)\n",
    "method which runs Stan's NUTS-HMC sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe796abc-f142-4e6e-88f3-5ff9971443b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_complete_pooling = complete_pooling_model.sample(data=mn_radon_data, show_progress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054e6409",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mn_no_pooling = no_pooling_model.sample(data=mn_radon_data, show_progress=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd2f3ec",
   "metadata": {},
   "source": [
    "### Extracting fit information\n",
    "\n",
    "The `sample` method returns a [CmdStanMCMC](https://mc-stan.org/cmdstanpy/api.html#cmdstanmcmc) object which\n",
    "provides methods to summarize and diagnose the model fit and accessor methods to access the entire sample or individual items.\n",
    "Accessor functions allow the user\n",
    "to access the sample in whatever data format is needed for further analysis.\n",
    "\n",
    "The sample can be treated as a collection of named, structured variables\n",
    "\n",
    "+ methods [stan_variable](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanMCMC.stan_variable) and [stan_variables](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanMCMC.stan_variables) return a `numpy.ndarray` and a Python dictionary of `numpy.ndarray` objects, respectively, whose structure corresponds to the Stan variable.\n",
    "\n",
    "+ method [draws_xr](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanMCMC.draws_xr) returns an [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html) of all Stan variables.\n",
    "\n",
    "The sample can be extracted in tabular format, either as\n",
    "\n",
    "+ method [draws](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanMCMC.draws) returns a [numpy.ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray) over all columns in the output CSV file.\n",
    "\n",
    "+ method [draws_pd](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanMCMC.stan_variable) returns a [pandas.DataFrame](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html#pandas.DataFrame) over all columns in the output CSV file.  The argument `vars` can be used to restrict this to specified variables or columns."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3bbd9cc",
   "metadata": {},
   "source": [
    "    \n",
    "**Extracting model estimates as pandas.DataFrame**\n",
    "\n",
    "We use the `draws_pd` method to access the sample draws for variables `alpha`, `beta`, and `sigma`.\n",
    "The draws across all chains are flattened into a single dimension.\n",
    "In this example, the output 3-D array of 4 chains of 1000 draws over 3 variables which becomes\n",
    "a 2-D array of 4000 draws of 3 variables.\n",
    "\n",
    "There are two reasons why we use the `draws_pd` method here:\n",
    "\n",
    "+ pandas provides many useful statistical functions.\n",
    "+ plotnine is designed to work with pandas DataFrames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a91ce8d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_pd = mn_complete_pooling.draws_pd(vars=['alpha', 'beta', 'sigma'])\n",
    "pool_pd.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6dc3f6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_stats = pool_pd.describe()\n",
    "pool_stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0ce6ddf",
   "metadata": {},
   "source": [
    "**Extracting model estimates as structured variables**\n",
    "\n",
    "To see the difference in the estimates of the intercept term `alpha` for the\n",
    "complete-pooling (single intercept) and the no-pooling models,\n",
    "i.e. to compare the estimates for a single intercept term\n",
    "and a vector of per-county intercept terms,\n",
    "we use the accessor method `stan_variable` which returns these\n",
    "estimates as [numpy.ndarray](https://numpy.org/doc/stable/reference/arrays.ndarray.html) objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e4ed90",
   "metadata": {},
   "outputs": [],
   "source": [
    "complete_pool_alpha_mean = mn_complete_pooling.stan_variable('alpha').mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23b07939",
   "metadata": {},
   "source": [
    "Another, more computationally expensive way to get these summary statistics is to call the [CmdStanMCMC.summary](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanMCMC.summary) method.\n",
    "This returns a pandas.DataFrame of summary statistics for total joint log probability lp__ and all model variables,\n",
    "plus diagnostic statistics on the\n",
    "[effective sample size](https://mc-stan.org/docs/reference-manual/effective-sample-size.html)\n",
    "and [R_hat](https://mc-stan.org/docs/reference-manual/notation-for-samples-chains-and-draws.html#potential-scale-reduction), the potential scale reduction factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6ef8553",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_complete_pooling.summary().round(2)[1:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6095274b",
   "metadata": {},
   "source": [
    "### Excercise\n",
    "\n",
    "Fit the complete pooling model to the national level data.  How to account for floors coded as 2, 3, or 9?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efc5df5",
   "metadata": {},
   "source": [
    "### Visualizing model estimates with plotnine\n",
    "\n",
    "To check that the Stan model's estimates are in line with the `geom_smooth(method='lm')`  results shown above,\n",
    "we create the same plot from the fitted model.\n",
    "We need to overlay the jittered plot of the data with a trend line showing\n",
    "the mean estimated log radon level for floor 0 and 1.\n",
    "To do this, we plug the estimates of `alpha` and `beta` into\n",
    "the equation $y = \\alpha + \\beta \\, x$\n",
    "\n",
    "+ when x $= 0$, y $=$ `alpha`\n",
    "+ when x $= 1$, y $=$ `alpha` + `beta`.\n",
    "\n",
    "Then we add a trend line to the plot which connects this pair of (x, y) points.\n",
    "In plotnine, there are two ways to do this\n",
    "\n",
    "+ [plotnine.geoms.geom_line](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_line.html#plotnine.geoms.geom_line) draws a line through a set of connected points\n",
    "+ [plotnine.stats.stat_function](https://plotnine.readthedocs.io/en/stable/generated/plotnine.stats.stat_function.html#plotnine.stats.stat_function) superimposes a function on a plot.\n",
    "\n",
    "We use the `stats.stat_function` to draw the mean trend line and the `geoms.geom_line` function to plot\n",
    "a random sample of draws from the posterior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531bcf68",
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = pool_pd.alpha + pool_pd.beta  # y coord at floor 1 \n",
    "f0 = pd.Series(pool_pd.alpha.values)  # y coord at floor 0\n",
    "\n",
    "# \n",
    "sz = 100\n",
    "ys = pd.concat([f0, f1], axis=1)\n",
    "ys = ys.sample(sz).reset_index(drop=True)\n",
    "\n",
    "# add sample regression lines to plot_radon_floor (from earlier section)\n",
    "p2 = plot_radon_floor\n",
    "for i in range(sz):\n",
    "    p2 = p2 + p9.geom_line(data=ys.T, mapping=p9.aes(x=[0,1], y=ys.loc[i]),\n",
    "                           inherit_aes=False, color='grey', alpha=0.06)\n",
    "\n",
    "# add central regression line\n",
    "p2 = p2 + p9.stat_function(mapping=p9.aes(x=1),\n",
    "    fun=lambda x: pool_stats.alpha['mean'] + pool_stats.beta['mean']*x,\n",
    "    color='blue', size=1\n",
    ")\n",
    "p2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b66da1d",
   "metadata": {},
   "source": [
    "Like pandas, numpy provides [statistics routines](https://numpy.org/doc/stable/reference/routines.statistics.html?highlight=statistics).  We use these to plot the mean of the estimate of alpha from the complete-pooling model and the central 50% interval of the elements of vector alpha, the per-county estimates from the no-pooling model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ccda7d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "no_pool_alpha = mn_no_pooling.stan_variable('alpha')\n",
    "no_pool_alpha_mean = np.mean(no_pool_alpha, axis=0)  # axis=0 uses all rows, i.e., per-column mean\n",
    "no_pool_alpha_lower = np.quantile(no_pool_alpha, 0.16, axis=0)\n",
    "no_pool_alpha_upper = np.quantile(no_pool_alpha, 0.84, axis=0)\n",
    "no_pool_alpha_pd = pd.DataFrame(data={\n",
    "    \"mean\": no_pool_alpha_mean,\n",
    "    \"upper\": no_pool_alpha_upper, \n",
    "    \"lower\": no_pool_alpha_lower, \n",
    "    \"county\":mn_counties['county']\n",
    "})\n",
    "no_pool_alpha_pd.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec336ce1",
   "metadata": {},
   "source": [
    "To compare the complete-pooling and no-pooling estimates of `alpha` we use the `geoms.geom_line` function to visualize the central 50% interval of the per-county estimates of alpha.\n",
    "The x-axis shows the county labels, ordered by number of observations descending, and the y-axis shows the values of `alpha`.\n",
    "We use a \n",
    "[plotnine.geoms.geom_hline](https://plotnine.readthedocs.io/en/stable/generated/plotnine.geoms.geom_hline.html) to add a horizontal line which is the mean value of `alpha` from the complete-pooling model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4399ad73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get sort order\n",
    "pop_asc = mn_counties.sort_values(by='homes').reset_index(drop=True)\n",
    "pop_asc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bd5d552",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_no_pool = (p9.ggplot(data=no_pool_alpha_pd)\n",
    " + p9.geom_segment(\n",
    "     mapping=p9.aes(x='county', xend='county', y='lower', yend='upper'),\n",
    "     size=1.4, color='darkblue', alpha=0.5,\n",
    " )\n",
    " + p9.geom_point(mapping=p9.aes(x='county', y='mean'))\n",
    " + p9.geom_hline(yintercept=complete_pool_alpha_mean, color='darkorange', size=1.5)\n",
    " + p9.scale_x_discrete(limits=pop_asc['county'], expand=(0,1))\n",
    " + p9.ggtitle(\"No-pooling (blue) and complete-pooling (orange) estimates for alpha (basement log_radon level)\")\n",
    " + p9.xlab(\"observations per county\") + p9.ylab(\"central 67% interval\")\n",
    " + xlabels_90\n",
    " + p9.theme(figure_size=(20,6)) \n",
    ")\n",
    "p_no_pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebccfa96",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p9.ggplot(data=no_pool_alpha_pd)\n",
    " + p9.geom_boxplot(data=mn_radon, mapping=p9.aes(x='county', y='log_radon'),\n",
    "                   alpha=0.5, color='grey', outlier_alpha=0.5)\n",
    " # Range strip\n",
    " + p9.geom_segment(\n",
    "     mapping=p9.aes(x='county', xend='county', y='lower', yend='upper'),\n",
    "     size=1.4, color='darkblue', alpha=0.5,\n",
    " )\n",
    " + p9.geom_point(mapping=p9.aes(x='county', y='mean'))\n",
    " + p9.geom_hline(yintercept=complete_pool_alpha_mean, color='darkorange', size=1.5)\n",
    " + p9.scale_x_discrete(limits=pop_asc['county'])\n",
    " + p9.ggtitle(\"No-pooling (blue) and complete-pooling (orange) estimates, plus raw data boxplots\")\n",
    " + p9.xlab(\"observations per county\") + p9.ylab(\"central 67% interval\")\n",
    " + xlabels_90\n",
    " + p9.theme(figure_size=(20,6))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11881b90",
   "metadata": {},
   "source": [
    "#### Discussion\n",
    "\n",
    "These plots are for the intercept of the model, which corresponds to average radon level for home with basements across the dataset.\n",
    "\n",
    "* The complete pooling model finds a common intercept.\n",
    "* The no-pooling model has considerable uncertainty for counties with small amounts of data.\n",
    "* Plotting over the raw data shows spread of measurements across counties.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fde1cb88",
   "metadata": {},
   "source": [
    "## Best Practice: posterior predictive checks\n",
    "\n",
    "> Posterior predictive checks are the unit tests of probabilistic programming.\n",
    ">\n",
    "> -- <cite> Ben Goodrich </cite>\n",
    "\n",
    "Posterior predictive checks test how well the fitted model captures basic features of the data.\n",
    "\n",
    "In the generated quantities block we generate a new sample of \"replicated data\",\n",
    "by convention called `y_rep` by using the estimated model parameters as arguments to\n",
    "Stan's PRNG distribution functions.\n",
    "If a model captures the data well, summary statistics such as sample mean and standard deviation,\n",
    "should have similar values in the original and replicated data sets.\n",
    "See [the Stan Users Guide](https://mc-stan.org/docs/stan-users-guide/simulating-from-the-posterior-predictive-distribution.html) for further details.\n",
    "\n",
    "For the no-pooling model, the likelihood statement is\n",
    "\n",
    "```stan\n",
    "y ~ normal(alpha[county] + beta * x, sigma);\n",
    "```\n",
    "\n",
    "The posterior predictive check is\n",
    "\n",
    "\n",
    "```stan\n",
    "array[N] real y_rep = normal_rng(alpha[county] + beta * x, sigma);\n",
    "```\n",
    "\n",
    "To see how this works, we again use the `draws_pd` method with argument `vars='y_rep'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f185a5f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_rep_pp = mn_no_pooling.draws_pd(vars='y_rep')\n",
    "y_rep_pp.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "551ef487",
   "metadata": {},
   "source": [
    "Next, we estimate the per-county median value for each of the 85 observed counties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43c5a527",
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate the median y_rep for each county\n",
    "stat_median = []\n",
    "for i in range(1,len(mn_counties)+1):  # Stan uses 1-based indexing\n",
    "    idxs = mn_radon.index[mn_radon['county_id'] == i].tolist()\n",
    "    stat_median.append(np.median(y_rep_pp.iloc[:, idxs].to_numpy().flatten()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5648b812",
   "metadata": {},
   "source": [
    "We create a plot which combines  a boxplot layer and a layer of per-county median point estimates.\n",
    "For counties with a large number of observations, the median of `y` and `y_rep` are quite close."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bcce03a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np_ppc_median = (p9.ggplot()\n",
    "  + p9.geom_boxplot(data=mn_radon,\n",
    "                    mapping=p9.aes(x='county',y='log_radon'),\n",
    "                    color='orange', fatten=3, alpha=0.7, outlier_alpha=0.3)\n",
    "  + p9.geom_point(mapping=p9.aes(x=mn_counties.county, y=stat_median), color='purple')\n",
    "  + p9.scale_x_discrete(limits=pop_asc['county'], expand=(0,2))\n",
    "  + p9.ggtitle(\"No-pooling model, posterior predictive checks: median estimates for alpha\")\n",
    "  + p9.xlab(\"observations per county\")                \n",
    "  + xlabels_90\n",
    "  + p9.theme(figure_size=(16,6))\n",
    ")\n",
    "np_ppc_median"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c73ed36d",
   "metadata": {},
   "source": [
    "Another visualization is to plot the distribution of the actual data against a random sample of replicates.\n",
    "We plot 2% of the data - 80 replicates out of 4000."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95df4d8c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get a random sample of the draws\n",
    "sz = 80\n",
    "y_rep = mn_no_pooling.draws_pd(vars='y_rep')\n",
    "\n",
    "# each column is a replicate of the data, using estimates of alpha, beta\n",
    "y_rep_sample = y_rep.sample(sz).reset_index(drop=True).T\n",
    "\n",
    "# plot actual distribution of the data against predicted new data\n",
    "np_ppc = p9.ggplot()\n",
    "for i in range(sz):\n",
    "    np_ppc = np_ppc + p9.stat_density(mapping=p9.aes(x=y_rep_sample[i]), geom='line', color='lightblue', alpha=0.4)\n",
    "np_ppc = (np_ppc \n",
    "          + p9.stat_density(data=mn_radon, mapping=p9.aes(x='log_radon'), geom='line', color='darkblue', size=1.1)\n",
    "          + p9.ggtitle(\"No-pooling model, posterior predictive checks: density of y, y_rep\")\n",
    "          + p9.xlab(\"log radon levels\") + p9.ylab(\"density\") + p9.scales.xlim(-3,6)\n",
    "          + p9.theme(figure_size=(6,4))\n",
    "         )\n",
    "np_ppc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0e29df1",
   "metadata": {},
   "source": [
    "The replicates draws are more diffuse; wider at the base, not lining up at the peak."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b7adb6a",
   "metadata": {},
   "source": [
    "### Adding county uranium as a predictor\n",
    "\n",
    "Model [radon_np_2.stan](stan/radon_np_2.stan) adds county-level soil uranium as a predictor.\n",
    "\n",
    "```{stan}\n",
    "data {\n",
    "  int<lower=1> N;  // observations\n",
    "  int<lower=1> J;  // counties\n",
    "  array[N] int<lower=1, upper=J> county;\n",
    "  vector[N] y;     // radon\n",
    "  vector[N] x;     // floor\n",
    "  vector[J] urainum;     // uranium\n",
    "}\n",
    "transformed data {\n",
    "  vector[N] u = uranium[county];\n",
    "  matrix[N, 2] xs = [x', u']';\n",
    "}\n",
    "parameters {\n",
    "  vector[J] alpha;\n",
    "  vector[2] beta;\n",
    "  real<lower=0> sigma;\n",
    "}\n",
    "transformed parameters {\n",
    "  vector[N] alphas = alpha[county];\n",
    "}\n",
    "model {\n",
    "  y ~ normal_id_glm(xs, alphas, beta, sigma);  \n",
    "  alpha ~ normal(0, 2.5);\n",
    "  beta ~ normal(0, 2.5);\n",
    "  sigma ~ normal(0, 5);\n",
    "}\n",
    "generated quantities {\n",
    "  array[N] real y_rep = normal_rng(alpha[county] + beta[0] * x + beta[1] * u, sigma);\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34edaaab",
   "metadata": {},
   "outputs": [],
   "source": [
    "no_pooling_model_2 = CmdStanModel(stan_file=os.path.join('stan', 'radon_np_2.stan'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3506aef8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_no_pooling_2 = no_pooling_model_2.sample(data=mn_radon_data, show_progress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bafafc42",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_1 = mn_no_pooling.summary().round(2)\n",
    "summary_1.loc[[\"alpha[3]\", \"alpha[23]\", \"beta\", \"sigma\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ed8c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_2 = mn_no_pooling_2.summary().round(2)\n",
    "summary_2.loc[[\"alpha[3]\", \"alpha[23]\", \"beta[1]\", \"beta[2]\",\"sigma\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6357df",
   "metadata": {},
   "source": [
    "### Discussion\n",
    "\n",
    "What is going on here?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2298444b",
   "metadata": {},
   "source": [
    "## Multilevel Regression\n",
    "\n",
    "Multilevel regression models the dependency structures in the data\n",
    "in addition to the relation between outcome and predictors.\n",
    "If the data has a hierarchical structure, all levels of the hierarchy may be modeled\n",
    "and all of the data is used to estimate all parameters jointly.\n",
    "\n",
    "For this dataset, houses are nested within counties.\n",
    "An ordinary regression can either model all counties as being identical, the complete pooling model,\n",
    "or all counties as being different, the no-pooling model.\n",
    "In the no-pooling model, the resulting estimates for counties with sparse data\n",
    "were too diffuse.\n",
    "A multi-level regression accounts for the observed variation across counties\n",
    "by modeling the counties as being drawn from\n",
    "a common distribution, whose parameters are estimated jointly by the model.\n",
    "This allows for *partial pooling* of information.\n",
    "When there is very little variation across counties, the multi-level model approaches the complete-pooling model,\n",
    "at the other extreme, when the amount of variation is very large, it approaches the no-pooling model\n",
    "\n",
    "A simple linear regression model with a single predictor\n",
    "estimates two parameters:  the intercept and slope of the regression line.\n",
    "In a multilevel model, either or both of these parameters can be modeled.\n",
    "As the number of regression predictors increases, the modeling choices increase.\n",
    "\n",
    "###  Modeling the regression intercept term\n",
    "\n",
    "A first multilevel version of the radon model extends the no-pooling model\n",
    "by putting a distribution on the intercept term\n",
    "\n",
    "$\n",
    "\\alpha_j \\sim \\mathrm{N}(\\mu_\\alpha,\\, \\sigma_\\alpha),\\ \\ \\ \\ \\mathrm{for}\\ \\ j\\, = 1, \\ldots, \\mathrm{J}\n",
    "$\n",
    "\n",
    "In combination with the likelihood, the partial-pooling model is\n",
    "\n",
    "$\n",
    "y_i = \\alpha_{j[i]} \\, + {\\beta}\\,x_i + {\\epsilon}_i \\\\\n",
    "\\alpha_j \\sim \\mathrm{N}(\\mu_\\alpha,\\, \\sigma_\\alpha),\\ \\ \\ \\ \\mathrm{for}\\ \\ j\\, = 1, \\ldots, \\mathrm{J}\n",
    "$\n",
    "\n",
    "All parameters are also estimated jointly by the model.\n",
    "This model provides _partial pooling_ of information;\n",
    "it pulls the estimates of $\\alpha_j$ towards the mean level $\\mu_\\alpha$, to a greater or lesser degree.\n",
    "\n",
    "Partial pooling is a _soft constraint_ whose effect depends on the amount of group-level variance $\\sigma_\\alpha$.\n",
    "\n",
    "**As $\\sigma_\\alpha$ increases, the amount of pooling decreases so that when\n",
    "$\\sigma_\\alpha \\rightarrow \\infty$ there is no pooling and when $\\sigma_\\alpha \\rightarrow 0$\n",
    "there is complete pooling.**\n",
    "\n",
    "###  Writing the Stan model\n",
    "\n",
    "\n",
    "```stan\n",
    "data {\n",
    "  int<lower=1> N;  // observations\n",
    "  int<lower=1> J;  // counties\n",
    "  array[N] int<lower=1, upper=J> county;\n",
    "  vector[N] x;\n",
    "  vector[N] y;\n",
    "}\n",
    "parameters {\n",
    "  real mu_alpha;\n",
    "  real<lower=0> sigma_alpha;\n",
    "  vector<offset=mu_alpha, multiplier=sigma_alpha>[J] alpha;  // non-centered parameterization\n",
    "  real beta;\n",
    "  real<lower=0> sigma;\n",
    "}\n",
    "model {\n",
    "  y ~ normal(alpha[county] + beta * x, sigma);  \n",
    "  alpha ~ normal(mu_alpha, sigma_alpha); // partial-pooling\n",
    "  beta ~ normal(0, 10);\n",
    "  sigma ~ normal(0, 10);\n",
    "  mu_alpha ~ normal(0, 10);\n",
    "  sigma_alpha ~ normal(0, 10);\n",
    "}\n",
    "generated quantities {\n",
    "  array[N] real y_rep = normal_rng(alpha[county] + beta * x, sigma);\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3be9bf81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "partial_pooling_alpha_model = CmdStanModel(stan_file=os.path.join('stan', 'radon_pp_alpha.stan'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92eac02d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "partial_pooling_alpha_fit = partial_pooling_alpha_model.sample(data=mn_radon_data, show_progress=False)\n",
    "partial_pooling_alpha_fit.summary().round(2).head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d7fd500",
   "metadata": {},
   "source": [
    "### Visualizations\n",
    "\n",
    "To visualize the results, we plot the central 50% of the estimates for `alpha`, as we did for the no-pooling model above. \n",
    "We use the `stan_variable` method to compute the mean and central 50% interval of the elements of vector alpha."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bc5250c",
   "metadata": {},
   "outputs": [],
   "source": [
    "part_pool_mu_alpha = partial_pooling_alpha_fit.stan_variable('mu_alpha').mean()\n",
    "part_pool_alpha = partial_pooling_alpha_fit.stan_variable('alpha')\n",
    "part_pool_alpha_mean = np.mean(part_pool_alpha, axis=0)\n",
    "part_pool_alpha_lower = np.quantile(part_pool_alpha, 0.16, axis=0)\n",
    "part_pool_alpha_upper = np.quantile(part_pool_alpha, 0.84, axis=0)\n",
    "part_pool_alpha_pd = pd.DataFrame(\n",
    "    data={\n",
    "        \"mean\": part_pool_alpha_mean,\n",
    "        \"upper\": part_pool_alpha_upper, \n",
    "        \"lower\": part_pool_alpha_lower, \n",
    "        \"county\":mn_counties['county']\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d68156b3",
   "metadata": {},
   "source": [
    "We plot the per-county estimates of `alpha` just as we did above, and we keep the y-axis on the same scale\n",
    "as for the no-pooling model; which were in range $(0, 3.5)$ (roughly).   This shows how the hierarchical pools information\n",
    "among the intercept terms and helps shrink the variance of the estimates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f28f5bf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize\n",
    "p_partial_pool_intercept = (p9.ggplot(data=part_pool_alpha_pd)\n",
    " # Range strip\n",
    " + p9.geom_segment(\n",
    "     mapping=p9.aes(x='county', xend='county', y='lower', yend='upper'),\n",
    "     size=1.7, color='blue', alpha=0.7,\n",
    " )\n",
    " + p9.geom_point(mapping=p9.aes(x='county', y='mean'))\n",
    " + p9.geom_hline(yintercept=part_pool_mu_alpha, color='darkblue', size=1)\n",
    " + p9.scale_x_discrete(limits=pop_asc['county']) + p9.scales.ylim(0,3.5)\n",
    " + p9.ggtitle(\"multilevel varying intercept model estimates for alpha (basement log_radon level)\")\n",
    " + p9.xlab(\"ordered by observations per county\") + p9.ylab(\"central 67% interval\")\n",
    " + xlabels_90\n",
    " + p9.theme(figure_size=(20,6)) \n",
    ")\n",
    "p_partial_pool_intercept"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3d0e965",
   "metadata": {},
   "source": [
    "To compare the no-pooling and partial pooling model estimates of `alpha` directly, we overlay the above plot with the estimated from the no-pooling model in orange."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5884d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p_partial_pool_intercept\n",
    "    + p9.geom_segment(data=no_pool_alpha_pd,\n",
    "         mapping=p9.aes(x='county', xend='county', y='lower', yend='upper'),\n",
    "         size=1.4, color='orange', alpha=0.6,\n",
    "     )\n",
    "     + p9.geom_hline(yintercept=complete_pool_alpha_mean, color='darkorange', size=1)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60d05c93-d931-4fa7-9341-b3bcf9f6d447",
   "metadata": {},
   "source": [
    "Another visualization is to plot the multilevel model estimates with the boxplots of the raw data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2a24b99",
   "metadata": {},
   "outputs": [],
   "source": [
    "(p_partial_pool_intercept\n",
    " + p9.geom_boxplot(data=mn_radon, mapping=p9.aes(x='county',y='log_radon'), \n",
    "                   color='orange', alpha=0.4, outlier_alpha=0.3)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31892a51",
   "metadata": {},
   "source": [
    "### Posterior predictive checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55f86baf",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mn_counties)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcd0460e",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_rep_pp = partial_pooling_alpha_fit.draws_pd(vars='y_rep')\n",
    "\n",
    "# compute per-county medians, means\n",
    "pp_stat_median = []\n",
    "for i in range(1,len(mn_counties)+1):   # Stan's 1-based indices\n",
    "    idxs = mn_radon.index[mn_radon['county_id'] == i].tolist()\n",
    "    pp_stat_median.append(np.median(y_rep_pp.iloc[:, idxs].to_numpy().flatten()))\n",
    "\n",
    "# plot medians from sample against boxplot y\n",
    "(p9.ggplot()\n",
    "  + p9.geom_boxplot(data=mn_radon,\n",
    "                    mapping=p9.aes(x='county',y='log_radon'), \n",
    "                    color='orange', fatten=2, alpha=0.7, outlier_alpha=0.3)\n",
    "  + p9.geom_point(mapping=p9.aes(x=pop_asc['county'], y=pp_stat_median), color='purple')\n",
    "  + p9.scale_x_discrete(limits=pop_asc['county'])\n",
    "  + p9.ggtitle(\"Partial-pooling model, posterior predictive checks: median estimates for alpha\")\n",
    "  + xlabels_90\n",
    "  + p9.theme(figure_size=(16,6))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b37fa65b",
   "metadata": {},
   "source": [
    "The PPC density plots are also similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "883a69a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# each column is a replicate of the data, using estimates of alpha, beta\n",
    "y_rep_pp_sample = y_rep_pp.sample(sz).reset_index(drop=True).T\n",
    "\n",
    "# plot actual distribution of the data against predicted new data\n",
    "pp_ppc = p9.ggplot()\n",
    "for i in range(sz):\n",
    "    pp_ppc = pp_ppc + p9.stat_density(mapping=p9.aes(x=y_rep_pp_sample[i]),\n",
    "                                      geom='line', color='lightblue', alpha=0.4)\n",
    "pp_ppc = (pp_ppc \n",
    "          + p9.stat_density(data=mn_radon, mapping=p9.aes(x='log_radon'),\n",
    "                            geom='line', color='darkblue', size=1.1)\n",
    "          + p9.ggtitle(\"Partial-pooling model, posterior predictive checks: density for y, y_rep\")\n",
    "          + p9.xlab(\"log radon levels\") + p9.ylab(\"density\") + p9.scales.xlim(-3, 6)\n",
    "          + p9.theme(figure_size=(6,4))\n",
    "         )\n",
    "pp_ppc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5e3c3eb",
   "metadata": {},
   "source": [
    "We can use the patchworklib package to display these plots side by side.\n",
    "For the partial pooling model, the densities of the replicates in light blue better approximate the density plot of the raw data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f2be0a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# combine multiple plots\n",
    "import patchworklib as pw  # ignore warning about seaborn\n",
    "from plotnine.data import *\n",
    "\n",
    "g1 = pw.load_ggplot(pp_ppc)\n",
    "g1.savefig(quick=True)\n",
    "g2 = pw.load_ggplot(np_ppc)\n",
    "g2.savefig(quick=True)\n",
    "p12 = (g1 | g2)\n",
    "\n",
    "p12.savefig(quick=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb3ed573",
   "metadata": {},
   "source": [
    "### Next Steps\n",
    "\n",
    "Add the soil level uranium as a predictor to the varying intercept model.\n",
    "\n",
    "Modeling choices:\n",
    "\n",
    "- uranium is fixed effect (as before)\n",
    "- uranium is varying effect:  this leads to many more modeling choices for the hierarchical priors.\n",
    "\n",
    "Discussion and code in the Stan User's Guide chapter [Multivariate priors for hierarchical models](https://mc-stan.org/docs/stan-users-guide/multivariate-hierarchical-priors.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed4257a",
   "metadata": {},
   "source": [
    "## Discussion\n",
    "\n",
    "In CmdStanPy, fitting a model to data is straightforward.\n",
    "\n",
    "* Instantiate a CmdStanModel object\n",
    "\n",
    "* Assemble a data dictionary or JSON file which contains definitions for all data variables declared in the model's `data` block.\n",
    "\n",
    "* Run one of the available inference methods:\n",
    "\n",
    "    + `sample` to do exact Bayesian estimation\n",
    "    + `variational` to do approximate Bayesian estimation\n",
    "    + `optimize` to do penalized maximum likelihood estimation\n",
    "\n",
    "* Extract the estimates for parameters and quantities of interest.  CmdStanPy's accessor functions make it easy to get the outputs in whatever data format is appropriate for downstream analysis.\n",
    "\n",
    "Fitting the model to the data is just one component of the analysis.\n",
    "The bulk of the code in this notebook is devoted to visualizing both the raw data and the model estimates;\n",
    "data visualization drives both the model specification process and the model testing process.\n",
    "With the plotnine package we can create multi-layer plots which overly the raw data with the model estimates\n",
    "and predictions or multi-layer plots which demonstrate the behaviors of different models.\n",
    "These activities play a central role in developing trustworthy data analysis pipelines:\n",
    "\n",
    "- Plotting the raw or simulated data before modeling informs the model design process\n",
    "\n",
    "- Plotting prior and posterior predictive checks drives testing\n",
    "\n",
    "- Plotting model estimates and predictions drives documentation and dissemination of results.\n",
    "\n",
    "\n",
    "In this notebook we work through the multilevel regression model from Gelman and Hill, chapter 12,\n",
    "restricting our analysis to just the data from the state of Minnesota.\n",
    "Even if we weren't writing this up as part of an introduction to multi-level modeling,\n",
    "that is, if we were tasked with this analysis qua analysis, we would proceed similarly -\n",
    "starting with a very simple model.\n",
    "There are many possible next steps:\n",
    "\n",
    "- Expanding the analysis to other states or the entire US.\n",
    "- Expanding the model to use the county-level soil uranium level measurements\n",
    "   + as a varying-intercept / varying-slope model (Chapter 12)\n",
    "   + as a spatial-effects model\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## References and Resources\n",
    "\n",
    "This notebook is based on Chris Fonnesbeck's excellent [A Primer on Bayesian Multilevel Modeling using PyStan](https://mc-stan.org/users/documentation/case-studies/radon.html), which was developed as part of a [Stan workshop](https://statmoddev.stat.columbia.edu/2016/06/09/a-primer-on-bayesian-multilevel-modeling-using-pystan/)  for biomedical statisticians at Vanderbilt University.\n",
    "\n",
    "[Stan Tutorials YouTube Playlist](https://www.youtube.com/playlist?list=PLCrWEzJgSUqwL85xIj1wubGdY15C5Gf7H) Maggie Lieu - a series of introductory videos on Bayesian modeling with Stan\n",
    "\n",
    "[Stan User's Guide](https://mc-stan.org/docs/stan-users-guide/index.html)\n",
    "\n",
    "[Visualization in Bayesian workflow](https://rss.onlinelibrary.wiley.com/doi/pdfdirect/10.1111/rssa.12378)\n",
    "Jonah Gabry, Daniel Simpson, Aki Vehtari, Michael Betancourt, and Andrew Gelman, 2019\n",
    "\n",
    "[Designing for interactive exploratory data analysis requires theories of graphical inference](http://www.stat.columbia.edu/~gelman/research/published/01640195064189.pdf) Jessica Hullman and Andrew Gelman, 2021\n",
    "\n",
    "[Data analysis using regression and multilevel/hierarchical models](http://www.stat.columbia.edu/~gelman/arm/) Andrew Gelman and Jennifer Hill, 2007\n",
    "\n",
    "[Statistical rethinking](https://xcelab.net/rm/statistical-rethinking/) by Richard McElreath -\n",
    "an intro-stats/linear models course taught from a Bayesian perspective.\n",
    "\n",
    " - [GitHub page](https://github.com/rmcelreath/rethinking)\n",
    " - [Course page](https://www.youtube.com/playlist?list=PLDcUM9US4XdMROZ57-OIRtIK0aOynbgZN)\n",
    "\n",
    "[Making Plots With plotnine](https://datacarpentry.org/python-ecology-lesson/07-visualization-ggplot-python/) - plotnine tutorial notebook.\n",
    "\n",
    "Papers by Price and Gelman on the radon data and multilevel models:\n",
    "\n",
    "[Centralized analysis of local data, with dollars and lives on the line: Lessons from the home radon experience.](http://www.stat.columbia.edu/~gelman/research/published/gelman_price_radon_book_chapter.pdf) Phillip N. Price and Andrew Gelman, 2015\n",
    "\n",
    "[Analysis of local decisions using hierarchical modeling, applied to home radon measurement and remediation (with discussion)](http://www.stat.columbia.edu/~gelman/research/published/lin.pdf) Phillip N. Price and Andrew Gelman, 1999\n",
    "\n",
    "<span id=\"regression\"> </span>\n",
    "\n",
    "## Appendix A: Linear regression review (chapters 2 and 3 of Gelman and Hill)\n",
    "\n",
    "\n",
    "[Linear regression](https://en.wikipedia.org/wiki/Linear_regression#Formulation) models the relationship between a scalar response and one or more prediexplanatory variables.\n",
    "\n",
    "<a title=\"Krishnavedala, CC BY-SA 3.0 &lt;https://creativecommons.org/licenses/by-sa/3.0&gt;, via Wikimedia Commons\" href=\"https://commons.wikimedia.org/wiki/File:Linear_least_squares_example2.svg\">\n",
    "    <img width=\"256\" alt=\"Linear least squares example2\"\n",
    "             style=\"vertical-align:middle;margin:0px 50px\" src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/b/b0/Linear_least_squares_example2.svg/256px-Linear_least_squares_example2.svg.png\">\n",
    "</a>\n",
    "\n",
    "A simple linear regression where the observed data (red) are assumed to be\n",
    "the result of random deviations (green) from an underlying relationship (blue)\n",
    "between the dependent variable (y) and an independent variable (x).\n",
    "The model fits the data by finding the linear function (a non-vertical straight line)\n",
    "which minimizes the distance between \"$y_i$ at offset $x_i$   and the line at offset $x_i$.\n",
    "\n",
    "$y_i = \\alpha \\, + \\beta\\,x_i + {\\epsilon}_i$\n",
    "\n",
    "* $\\alpha$ is the _intercept_, the offset from zero on the x-axis\n",
    "* $\\beta$ is the _slope_, the multiplier applied to x.\n",
    "* ${\\epsilon}_i$ is the error term, assuming independent errors drawn from a normal distribution with mean $0$, standard deviation $\\sigma$.\n",
    "\n",
    "In a simple linear regression, there is only a single predictor.  In an ordinary linear regression,\n",
    "there are many predictors, i.e., the model fits a vector of coefficients $\\beta$ to the vector\n",
    "of independent variables $X$.\n",
    "\n",
    "**Interpretation**\n",
    "\n",
    "In Chapter 3, Gelman and Hill write:\n",
    "\n",
    "> Linear regression is a method that summarizes how the average values of a numerical outcome variable vary over subpopulations defined by linear functions of predictors. ... Regression can be used to predict an outcome given a linear function of these predictors, and regression coefficients can be thought of as comparisons across predicted values or as comparisons among averages in the data.\n",
    "\n",
    "\n",
    "**Linear regression:  two ways to write the model**\n",
    "\n",
    "The goal of inference is to learn from incomplete or imperfect data.\n",
    "In the simple linear regression model, the error term $\\epsilon$ accounts for imperfect measurements of the data.\n",
    "\n",
    "$\n",
    "y_i = \\alpha \\, + {\\beta}\\,x_i + {\\epsilon}_i\n",
    "$\n",
    "where the errors ${\\epsilon}_i$ have independent normal distributions with mean $0$ and standard deviation $\\sigma$.\n",
    "\n",
    "An equivalent representation is\n",
    "\n",
    "$\n",
    "y_i ∼ \\mathrm{N}(\\alpha + \\beta\\,\\mathrm{X}_i,\\, \\sigma^2), \\ \\ \\ \\mathrm{for}\\ i=1, ..., n\n",
    "$\n",
    "\n",
    "The corresponding Stan statement is\n",
    "\n",
    "```stan\n",
    "y ~ normal(alpha + beta * x, sigma);\n",
    "```\n",
    "\n",
    "Stan provides [vectorized](https://mc-stan.org/docs/functions-reference/vectorization.html#vectorization) versions of all univariate probability distrions.\n",
    "This statement is far more efficient than using a `for` loop over all $x_i$ and $y_i$ pairs.\n",
    "\n",
    "**Simple linear regression model in Stan**\n",
    "\n",
    "The simple linear regression with a single predictor and a slope and intercept coefficient and normally distributed noise is the first model discussed in the [Stan User's Guide Regression Models chapter](https://mc-stan.org/docs/stan-users-guide/linear-regression.html).\n",
    "\n",
    "```stan\n",
    "data {\n",
    "  int<lower=0> N;\n",
    "  vector[N] x;\n",
    "  vector[N] y;\n",
    "}\n",
    "parameters {\n",
    "  real alpha;\n",
    "  real beta;\n",
    "  real<lower=0> sigma;\n",
    "}\n",
    "model {\n",
    "  y ~ normal(alpha + beta * x, sigma);\n",
    "}\n",
    "```\n",
    "\n",
    "This model is the minimal possible model.\n",
    "It consists of three named [program blocks](https://mc-stan.org/docs/reference-manual/overview-of-stans-program-blocks.html):\n",
    "the data and parameters are declared in the respectively named data and parameters block;\n",
    "the model block specifies the likelihood, i.e., the probability of the data given the model.\n",
    "Because no priors are specified on the model parameters, they are given the default prior\n",
    "distribution, which is uniform from  $-\\infty$ to $+\\infty$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a539568e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": true,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}