Tutorial 4.1 Shapley.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "abed0ec8-aebc-4f86-b6d4-27aa5bf09b19",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 4.1. SHAP Values: Origins and Applications \n",
    "### Alex Gagliano (gaglian2@mit.edu)\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/alexandergagliano/InterpretabilityDemos/blob/main/Tutorial%204.1%20Shapley.ipynb)\n",
    "\n",
    "References and resources for additional reading:\n",
    "* [A Unified Approach to Interpreting Model Predictions](https://github.com/shap/shap) (Lundberg & Lee, 2017)\n",
    "* [Explaining Predictive Uncertainty with Information Theoretic Shapley Values](https://arxiv.org/pdf/2306.05724.pdf) (Watson et al., 2023)\n",
    "* [Interpretable Machine Learning](https://christophm.github.io/interpretable-ml-book/), Christoph Molnar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf761308-bc42-46c4-b4be-31e32b9a9bf7",
   "metadata": {},
   "source": [
    "## Interpretability is deeply subjective and problem-dependent.\n",
    "Is understanding how the model is constructed enough, or do we want the construction to be intuitive (e.g., through regularization of the latent space)? Is it sufficient that a model is _robust_ (generalizes well to unseen data), or would we like to make causal statements linking features to the model's predictions?   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fadb9cd7-5883-4549-a7b8-038d1f0fe0f9",
   "metadata": {},
   "source": [
    "However you define it, interpretability often decreases as model complexity increases; we build a deep neural network for its expressivity, but then fail to understand how input leads to its output. In these cases, it can be useful to build a simpler _explanation_ model, a highly interpretable approximation of our original model. There are many ways to build explanation models; we can build a random forest and plot individual trees, or a linear regression model and examine its coefficients. \n",
    "\n",
    "Additive Feature Attribution Models are explanation models that take the following functional form:\n",
    "$$\n",
    "g(z') = \\phi_0 + \\sum_{i=1}^N \\phi_i z'_i\n",
    "$$\n",
    "\n",
    "Here $\\phi$ are the coefficients of the linear model we'd like to find, and $z'$ is a binary encoding of $N$ features in the model. The features are binarized such that, if the output approximates our more complex model, the derived coefficients $\\phi$ (and their sign) become a proxy for how much each feature contributes to our final output.\n",
    "\n",
    "These types of explanation models have been applied in many different forms through the literature: some interpretability tools using them include LIME [(Ribeiro et al., 2016)](https://arxiv.org/abs/1602.04938), DeepLIFT [(Shrikumar et al., 2017)](https://arxiv.org/abs/1704.02685), and Layer-Wise Relevance Propagation [(Binder et al., 2016)](https://arxiv.org/abs/1604.00825). Another explanation of these from cooperative game theory is the _Shapley value_. This tutorial provides a basic description for these below and then dives into some concrete applications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195513a2-3267-442f-a298-0bcc0d01e571",
   "metadata": {},
   "source": [
    "## Question: Let's say we're designing an additive feature attribution model $g(z')$. What properties should the model satisfy?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65e2302b-50d6-4af4-a031-fbb0ca7c1cab",
   "metadata": {},
   "source": [
    "1. *Local accuracy*: The function should approximate the model output for the same input features $x$:\n",
    "$$\n",
    "f(x) \\approx g(z')\n",
    "$$\n",
    "We only require local accuracy because, if $g$ could approximate our more complex $f$ globally, we wouldn't need it to begin with."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0109ad03-d4bc-45dc-9b21-2b12e1a9435e",
   "metadata": {},
   "source": [
    "2. *Missingness*: An absent feature should not influence the output:\n",
    "$$\n",
    "z_i' = 0 \\rightarrow \\phi_i = 0$$\n",
    "Technically, if $z_i' = 0$ then _any_ value of $\\phi_i$ would be possible (the term would have no impact on our output value $g(z')$). Constraining $\\phi_i=0$ in these cases matches our intuition: if we don't include feature $i$, it has no bearing on our model output."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb713eb-04ed-4bab-bf5d-05ae82769057",
   "metadata": {},
   "source": [
    "3. *Consistency*: If model 1's output increases with feature *i* more than in model 2, then its coefficient for feature $i$ in the linear model should be larger; in other words,\n",
    "$$f_1(x \\cup \\{i\\}) - f_1(x) > f_2(x \\cup \\{i\\}) - f_2(x) $$\n",
    "$$\\rightarrow \\phi_{i, f_1} > \\phi_{i, f_2}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a19b5125-1d20-4a69-a055-ff5ed119fca6",
   "metadata": {},
   "source": [
    "Demanding these three constraints, only a *single* function satisfies all these properties. \n",
    "$$\n",
    "\\phi_i = \\frac{1}{N!}\\sum_{R} \\left[f(x_R \\cup \\{i\\}) - f(x_R) \\right]\n",
    "$$\n",
    "Where $R$ is an order-specific subset (called a _coalition_) of features from set $S$ and $N$ is the total number of features.\n",
    "\n",
    "\n",
    "This function calculates Shapley values, first introduced by Lloyd Shapley in 1951. Shapley values are the average expected marginal contributions of each feature (or each player in a cooperative game) after all possible combinations of features (players) have been considered.\n",
    "\n",
    "Very confusingly, the paper that unified these additive feature attribution models under the Shapley equation called their approach SHAP (SHapley Additive exPlanation; Lundberg & Lee, 2017). Shapley values are the exact maginal contributions, and SHAP approaches are those that estimate them in various way (calculating them precisely is difficult). One popular technique is samping permutations of features using monte carlo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0089eab-9c44-4293-acf7-76b832aecf68",
   "metadata": {},
   "source": [
    "Let's start by installing and importing some necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dc2203f-eaf0-4607-ac17-ad11981f000a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install xgboost shap torchvision torch seaborn pyarrow gdown rfpimp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb93440-d1f1-4b7c-ae95-4879f7aa2f72",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import xgboost\n",
    "import numpy as np \n",
    "import shap\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import subprocess\n",
    "from bisect import bisect\n",
    "import itertools \n",
    "import rfpimp\n",
    "import gdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf335983-e4ca-43f4-9419-d37a209d0775",
   "metadata": {},
   "outputs": [],
   "source": [
    "#make our plots pretty \n",
    "sns.set_context(\"notebook\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "637c8d14-7ac0-4984-8dfb-917a412a95d9",
   "metadata": {},
   "source": [
    "## 4.1.1. Photometric redshifts.\n",
    "Redshift estimation is a field in which machine learning methods are relatively mature, owing to the use of features with observational uncertainties and the significant and unavoidable degeneracies involved (a red, dim galaxy at low-redshift looks like a blue, bright galaxy at high-redshift). In this example, we're going to take a look at simulated galaxy photometry from the [CosmoDC2](https://arxiv.org/abs/1907.06530) catalog. The photometry is generated using an image simulation of the upcoming Vera Rubin Observatory, and includes realistic models for dust from both the Milky Way and the host galaxy, atmospheric distortion, and the optical response system of the telescope. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd72b662-7d17-44e4-845a-7ffb73b7ac1d",
   "metadata": {},
   "source": [
    "First, let's download the data we'll need for all the tutorials and open the redshift example. We'll also download some archival models, in case we can't get anything to load in real-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10612186-8b15-4262-857b-3e9fa24ce8da",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gdown 1s42ri7tpvBHN-kneC0MHXHsd2nB2BNTE\n",
    "!gdown 1CC07axEZravXcIkq0w2gnkFGGhO4vbnx\n",
    "\n",
    "for file in ['data', 'models']:\n",
    "    subprocess.run([\"tar\", \"-xf\", \"%s.tar.gz\" % file]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b54f0bde-7742-4224-90aa-23963fe17431",
   "metadata": {},
   "outputs": [],
   "source": [
    "traindf = pd.read_parquet('data/redshift/dc2_gold_training_9816.pq')\n",
    "testdf = pd.read_parquet('data/redshift/dc2_gold_test_9816.pq')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbcd3649-d51c-47e1-95d3-b22f0bcb9e6b",
   "metadata": {},
   "source": [
    "Let's look at the features in the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff56eb52-cc5b-476b-95fa-687f3fd09d9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "traindf.drop(columns=['redshift'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "675dad98-ed87-4952-a5d3-07ac84310ac7",
   "metadata": {},
   "source": [
    "Now let's train a gradient boosted decision tree model with XGBoost:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed00565f-2fb2-42d1-abc5-de138d7ba9f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = traindf[['ug', 'gr', 'ri', 'iz', 'zy']]\n",
    "y_train = traindf['redshift']\n",
    "model = xgboost.XGBRegressor().fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2660b1a5-feea-4f05-8d76-8d75b9adbacf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#is it a good model? \n",
    "X_test = testdf[['ug', 'gr', 'ri', 'iz', 'zy']]\n",
    "y_test = testdf['redshift']\n",
    "y_pred = model.predict(X_test)\n",
    "plt.plot(y_test, y_pred, 'o', ms=1);\n",
    "plt.ylabel(r\"$z_{\\rm{phot}}$\");\n",
    "plt.xlabel(r\"$z_{\\rm{true}}$\");\n",
    "plt.plot(np.linspace(0, 3), np.linspace(0, 3), c='k', ls='--', lw=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2df5d829-2d3c-43ba-9198-4e77e3c674d2",
   "metadata": {},
   "source": [
    "## Question: How do I estimate the uncertainties on my predicted redshifts? "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0375dedd-3d6f-42c3-a73c-6981fa5e33b9",
   "metadata": {},
   "source": [
    "One might guess that we can just monte-carlo our estimates, assuming our errors are normally distributed and independent. \n",
    "\n",
    "## Challenge: Translate the photometric uncertainties on our observations into color uncertainties, and generate a basic noise model for the first 100 galaxies.\n",
    "Assume our uncertainties are drawn from a multivariate Normal, and draw 1000 samples of the feature set for each galaxy. Then, apply our XGBoost model to generate 1000 redshift estimates for each galaxy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f829c99-f637-42fa-a155-b173a0c2f132",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Translate photometric uncertainties (mag_err_g_lsst, mag_err_r_lsst, etc) to color uncertainties in testdf here.\n",
    "\n",
    "true_vals = []\n",
    "redshift_samples = []\n",
    "for idx in np.arange(10000):\n",
    "    means =  testdf.iloc[idx][['ug', 'gr', 'ri', 'iz', 'zy']]\n",
    "    errs = testdf.iloc[idx][['ug_err', 'gr_err', 'ri_err', 'iz_err', 'zy_err']]\n",
    "    \n",
    "    num_samples = 1000\n",
    "    flat_means = means.ravel()\n",
    "    \n",
    "    feature_samples = #TODO: draw num_samples from a multivariate Gaussian to generate additional noisy features\n",
    "    # Store the predicted redshifts and the true redshifts\n",
    "    redshift_samples.append(model.predict(feature_samples))\n",
    "    true_vals.append(testdf.iloc[idx]['redshift'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f613eed-8fc8-4281-9327-1b1f865a7aa7",
   "metadata": {},
   "source": [
    "That would give us the following prediction for a single galaxy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "894c3461-7a67-49f0-a769-eb55552bb3af",
   "metadata": {},
   "outputs": [],
   "source": [
    "point =  model.predict(testdf.iloc[[idx]][['ug', 'gr', 'ri', 'iz', 'zy']])\n",
    "err = np.std(redshift_samples[idx])\n",
    "print(r\"Prediction: z = %.2f +/- %.2f.\"%(point, err))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2affb3d-535f-41a3-b737-887f4bba0e74",
   "metadata": {},
   "source": [
    "But is our guess at the noise model reasonable? Let's plot the CDF for a single galaxy and see where our true redshift falls along the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8a1cbab-a63d-4667-93ae-c9a46d7015a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "CDF_rank = bisect(np.sort(redshift_samples[0]), true_vals[0])\n",
    "\n",
    "plt.plot(np.sort(redshift_samples[0]), np.linspace(0, 1, len(redshift_samples[0]), endpoint=False))\n",
    "plt.plot(true_vals[0], CDF_rank/len(redshift_samples[0]), 'o')\n",
    "plt.xlabel(\"Redshift\");\n",
    "plt.ylabel(\"Empirical CDF\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df0943eb-11aa-44ea-a1a5-0fd081ec4524",
   "metadata": {},
   "source": [
    "Clearly our guess for the first galaxy is a bad one. If we plot the rank for all of our galaxies, we would expect them to be uniformly distributed if our noise model was well-calibrated _globally_. That doesn't mean that it can't be locally wrong! Let's look at a larger sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73ef0545-9588-4ca7-8029-ef54931fdc41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Generalize the above routine to store the CDF ranks of all galaxies \n",
    "# for which we generated samples, and plot a histogram.\n",
    "\n",
    "plt.xlabel(\"True Redshift Rank\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c67434ae-fe8a-40cd-8b46-f8c1686ebd97",
   "metadata": {},
   "source": [
    "## Question: Does the result look like a uniform distribution? If not, what features do you notice? Is your noise model biased? Overestimated?\n",
    "\n",
    "Propagating the uncertainties directly through a linear model is fine in the Taylor expansion limit with small errors relative to the measurement. Unfortunately, XGBoost uses a softmax objective function to predict a mean point estimate and \"double-counts\" the measurement error to generate variance in that estimate. Sometimes this is get close to the right variance, but sometimes it doesn't. Be careful!\n",
    "\n",
    "In practice, it's common to determine a scale factor on all variances that pushes this distribution closer to uniform. But an individual noise estimate can still be wrong. Another way to understand investigate these is to generate a full posterior from a well-trained normalizing flow model, and compare the uncertainties between approaches. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc38440f-6568-4574-9417-252039a85bf8",
   "metadata": {},
   "source": [
    "Next, let's explain the model's predictions using SHAP.\n",
    "\n",
    "How informative are our colors? Let's check it out for a low-redshift and then a high-redshift event:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db928edf-ef83-468a-8b12-a9e6e825d725",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that this same syntax works for LightGBM, CatBoost, scikit-learn, transformers, Spark, etc!\n",
    "shap.initjs()\n",
    "\n",
    "explainer = shap.Explainer(model)\n",
    "shap_values = explainer(X_test)\n",
    "\n",
    "shap.plots.waterfall(shap_values[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d95f4a0-2b23-4bcf-9478-e8f62c9b0ca6",
   "metadata": {},
   "source": [
    "## Question: Which features have the largest marginal contribution toward a final prediction for this low-redshift event, and which have the lowest?\n",
    "\n",
    "What about for a high-redshift event?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7238b18-b1b0-4390-afd7-fcefbc3d80d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "shap.initjs()\n",
    "\n",
    "shap.plots.waterfall(shap_values[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "697456ac-9124-4029-9560-2bd93f4ee498",
   "metadata": {},
   "source": [
    "Do you notice any trends between the high-redshift events and the low-redshift ones?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "629d04a6-4eb6-49dd-b5e3-5b91465f570a",
   "metadata": {},
   "source": [
    "How do we observe the shapley values for an ensemble of predictions, and not just one? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609fdb0a-4cf0-4d3d-b667-23a9bb0864c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "shap.initjs()\n",
    "\n",
    "shap.plots.force(shap_values[0:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35b00842-f3bf-442e-a42e-1e505896cf06",
   "metadata": {},
   "source": [
    "Here we see that most colors suggest that nearby events are low-redshift. It looks like the feature with the most leverage at low-$z$ (on average) is $g-r$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07025f1e-8132-4942-bf58-bbffff847f1f",
   "metadata": {},
   "source": [
    "How does the entire distribution of SHAP values look for each parameter? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3fd6d6c-fa52-4f95-bbd5-315b66368109",
   "metadata": {},
   "outputs": [],
   "source": [
    "shap.initjs()\n",
    "\n",
    "shap.plots.beeswarm(shap_values);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc232ebd-9f39-420b-9430-f75c5c937fe7",
   "metadata": {},
   "source": [
    "Another way of seeing this is, for a simple model and a set of features, _directly visualizing_ the marginal contribution to the expectation value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c3ec4cd-ecf9-4b9e-bfaa-56d9a46eb744",
   "metadata": {},
   "outputs": [],
   "source": [
    "shap.initjs()\n",
    "\n",
    "idx = 2000\n",
    "\n",
    "#assume some background distribution of 100 values\n",
    "X100 = shap.utils.sample(X_test, 100)\n",
    "\n",
    "explainer_xgb = shap.Explainer(model, X100)\n",
    "shap_values_xgb = explainer_xgb(X_test)\n",
    "\n",
    "shap.plots.scatter(shap_values_xgb[:,\"gr\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71719d38-63d4-4fa9-bd3b-001521764f5c",
   "metadata": {},
   "source": [
    "## Question: What features do we observe? \n",
    "Remember what SHAP values indicate - the direction the features push the specific expectation value away from the mean expectation value. Can you think of a physically motivated reason for these particular color features? Consider what happens to a galaxy's position in color-color space as it gets redshifted:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed22cbe9-834d-4ec3-afe3-4409873b569a",
   "metadata": {},
   "source": [
    "<img src=\"images/ColorZ.gif\" alt=\"ColorZGif\" width=\"800\"/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0374ff35-6ca2-40ab-9362-40ed4ac6e6d2",
   "metadata": {},
   "source": [
    "Before we move on, we should note that we're using gradient boosted trees as our model, and these have some useful model-specific interpretability tricks themselves. These will be _global_, meaning they speak to the overall model across a wide range of features. \n",
    "\n",
    "First, because the model constructs a series of decision trees, we can visualize the trees directly to better understand how the model's decision-making process takes place. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3339ad54-b639-41a4-90a9-020e65ee0a97",
   "metadata": {},
   "outputs": [],
   "source": [
    "#on macOS, run \"brew install graphviz\" before executing this cell\n",
    "fig, ax = plt.subplots(figsize=(20, 20))\n",
    "xgboost.plot_tree(model, num_trees=1, ax=ax, rankdir='LR')\n",
    "plt.savefig(\"DecisionTree.png\", dpi=500, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c111b92-6715-4c06-9d98-b41d60cfecea",
   "metadata": {},
   "source": [
    "This tells us _exactly_ what the regressor is doing. Further, the fact that $g-r$ is the first split in this tree suggests it might be very important in the model overall. Unfortunately (at least for interpretability), there can be many trees constructed by the model (XGBoost uses 100 by default, albeit small by neural network standards), making it a bit too onerous to inspect them one-by-one. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb05221a-ecf7-45bc-90bc-69724d735d8e",
   "metadata": {},
   "source": [
    "We can also calculate the _feature importances_ of the model. The default here is the \"gain\", the average increase in accuracy of a tree after introducing the feature in question. Because this is averaged across all trees (ensembling), it is more a more unbiased estimate than what you might get from a single decision tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d58827a3-298e-4441-940b-568ad8108317",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_idx = model.feature_importances_.argsort()\n",
    "plt.barh(X_test.columns[sorted_idx], model.feature_importances_[sorted_idx])\n",
    "plt.xlabel(\"Xgboost Feature Importance\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "747efc71-a805-410b-a405-dcc5eea246e5",
   "metadata": {},
   "source": [
    "Our guess from the first tree was right - it looks like $g-r$ is the most important feature in our model. It is important to note that there are _many_ ways to calculate feature importance, and [some of them produce biased results](https://explained.ai/rf-importance/#7). To be safe, always explore multiple feature importance estimators:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af84f725-d951-4797-b537-f8ca78fa57e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nrows=1, ncols=6, figsize=(20, 8))\n",
    "axs = axs.ravel()\n",
    "measures = ['gain','total_gain', 'cover',  'total_cover', 'weight']\n",
    "palette = itertools.cycle(sns.color_palette('Dark2'))\n",
    "\n",
    "for i, measure in enumerate(measures):\n",
    "    importances = model.get_booster().get_score(importance_type=measure)\n",
    "    axs[i].barh(list(importances.keys()), list(importances.values()), color=next(palette))\n",
    "    axs[i].set_xlabel(\"Feat. Importance\");\n",
    "    axs[i].set_title(\"Importance = %s\"%measure)\n",
    "\n",
    "importance_df = rfpimp.importances(model, X_test, y_test).reindex(importances.keys())\n",
    "axs[5].barh(importance_df.index, importance_df['Importance'])\n",
    "axs[5].set_xlabel(\"Feat. Importance\");\n",
    "axs[5].set_title(\"Permutation Importance\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db91840a-f035-432f-9071-ff1b88d831ad",
   "metadata": {},
   "source": [
    "Total gain is the increase in accuracy summed across all decision trees (instead of averaged), while coverage (and total coverage) is the mean (and total) number of observations included within the splits determined by a feature. The permutation importance is defined as the decrease in model score caused by randomly shuffling the values of that feature, and is believed to resolve several biases introduced by other estimators. \n",
    "\n",
    "We can see how much variability there is between importance estimates. Nonetheless, some trends emerge. In many estimators, $z-y$ and $u-g$ are of low importance to the model - $u$ and $y$ have low overall transmission and the photometry in these bands is typically noisy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3883b5dd-bab4-4f42-aa16-16c3e5f852af",
   "metadata": {},
   "source": [
    "One final point: Look at the SHAP values for a single galaxy versus the gain importance for the model overall:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ea97d4a-9b77-4eee-9ed5-636def68d667",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(10, 8))\n",
    "axs = axs.ravel()\n",
    "palette = itertools.cycle(sns.color_palette('Dark2'))\n",
    "\n",
    "axs[0].barh(X_test.columns, model.feature_importances_)\n",
    "axs[0].set_title(\"Gain Importance\");\n",
    "\n",
    "shap_df = pd.DataFrame({'Shap':shap_values[-1].values}, index=X_test.columns).reindex(importances.keys())\n",
    "axs[1].barh(shap_df.index, shap_df['Shap'])\n",
    "axs[1].set_title(\"SHAP values for z = %.2f galaxy\"%y_test.values[-1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f5d3f49-d971-4308-85f8-81188ae9830d",
   "metadata": {},
   "source": [
    "We see some similarities, but also some important differences - a macroscopic and microscopic view of the same model!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cff3eca3-9aee-4195-b9b0-454d353fc6e3",
   "metadata": {},
   "source": [
    "## 4.1.2. Now let's  use Shapley values to better understand a more complex model. \n",
    "(Modified from https://github.com/pmocz/artificialneuralnetwork-python) \n",
    "\n",
    "Here we're going to attempt to classify galaxies according to their morphology. We consider three classes: spiral, elliptical, and irregular. These are reasonably straightforward to distinguish by eye: Spiral galaxies have a wound spiral structure, elliptical galaxies have a smooth light profile with a bright core, and irregular galaxies have some messy structure (this is a catchall category for the oddballs - irregular galaxies likely formed from some merger, close encounter, or chaotic internal activity):"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "535e0995-fde5-4db1-b716-38c7e8317c0e",
   "metadata": {},
   "source": [
    "![](Images/GalaxyTypes.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eec7f78-561f-4bf9-9d76-e5475d26cf7e",
   "metadata": {},
   "source": [
    "The images we're using are from the Sloan Digital Sky Survey. Let's start off by loading some required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8319c07a-b3a5-4ce5-9529-0fa1d77a53cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, torchvision\n",
    "from torchvision import datasets, transforms\n",
    "from torch import nn, optim\n",
    "import subprocess\n",
    "import os\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f70c993-c922-41d2-92b5-583d20888923",
   "metadata": {},
   "outputs": [],
   "source": [
    "galaxyPath = \"./data/galaxy/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ffaf3e8-4e41-4ea5-9db9-6f688c47380c",
   "metadata": {},
   "source": [
    "We specify a batch size for our training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05396d81-7279-4ee2-89ee-16d331e252fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e61de1-da56-44fa-82f4-5c1b4c1a14ce",
   "metadata": {},
   "source": [
    "Now we create datasets for the train and test sets, making sure to convert to grayscale and represent the image as tensors (so that pytorch can use them). We also use a pytorch generator object to ensure that our random batches are reproducible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f580bcc-5e92-4559-8d70-c36abc324d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_worker(worker_id):\n",
    "    worker_seed = torch.initial_seed() % 2**32\n",
    "    numpy.random.seed(worker_seed)\n",
    "    random.seed(worker_seed)\n",
    "\n",
    "g = torch.Generator()\n",
    "g.manual_seed(0)\n",
    "\n",
    "data_transform = transforms.Compose([transforms.Grayscale(num_output_channels=1),\n",
    "                                     transforms.ToTensor()])\n",
    "\n",
    "train = torchvision.datasets.ImageFolder(os.path.join(galaxyPath, 'train'), transform = data_transform)\n",
    "train_loader = DataLoader(train, shuffle=True, batch_size=batch_size, worker_init_fn=seed_worker, generator=g)\n",
    "\n",
    "test = torchvision.datasets.ImageFolder(os.path.join(galaxyPath, 'test'), transform = data_transform)\n",
    "test_loader = DataLoader(test, shuffle=True, batch_size=batch_size, worker_init_fn=seed_worker, generator=g)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5de01d42-cd83-45e3-b04d-fb2cf4ac848e",
   "metadata": {},
   "source": [
    "Our classification scheme is encoded as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9244004-871e-4707-9547-40186c41d4e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "GalaxyClasses = {0:'Ellip.', 1:'Spiral', 2:'Irreg.'}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5542f908-68d9-4fe3-b4d7-706a0ed224c8",
   "metadata": {},
   "source": [
    "Next, we create a basic convolutional neural network for classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e8a0d69-619e-42f0-a351-90baf20028e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cpu')\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Conv2d(1, 3, kernel_size = 3, stride = 1, padding = 1),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(3, 20, kernel_size=5),\n",
    "            nn.Dropout(0.3), #TODO: Modify the dropout fraction here.\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.fc_layers = nn.Sequential(\n",
    "            nn.Linear(720, 50),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.Linear(50, 3),\n",
    "            nn.Softmax(dim=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layers(x)\n",
    "        x = x.view(-1, 3*16*15)\n",
    "        x = self.fc_layers(x)\n",
    "        return x\n",
    "        \n",
    "model = Net().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ebfc2fa-753c-4834-8f91-ef7d1d8ce12b",
   "metadata": {},
   "source": [
    "Our network includes _dropout_, in which random subsets of neurons are deactivated during the training. The number in parentheses sets the dropout fraction. Dropout is an important tool in ensuring efficient training of neural networks, as it prevents a few neurons from overtraining while the majority of neurons learn nothing. This is similar to ensembling - training multiple random subsets of the network at a time and combining the results is often more powerful than training the full network at once."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2795bb2-18f5-466e-b6bd-5b6973372a47",
   "metadata": {},
   "source": [
    "We now train for 50 epochs, evaluating the training loss every epoch to see how we're improving.\n",
    "\n",
    "## Challenge: Change the dropout fraction, and evaluate its impact on the training routine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda3c6cc-6e86-49fe-9b49-1ee43ea3036a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "criterion = nn.CrossEntropyLoss() \n",
    "\n",
    "for epoch in trange(50):  # loop over the dataset multiple times\n",
    "    for i, data in enumerate(train_loader, 0):\n",
    "            # get the inputs; data is a list of [inputs, labels]\n",
    "            inputs, labels = data\n",
    "\n",
    "            # zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward + backward + optimize\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "print('Finished Training')\n",
    "\n",
    "#save the model \n",
    "torch.save(model.state_dict(), './models/galaxyCNN.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902a24f2-aba3-4954-bdd3-afa84bfeadd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you're having trouble training, pre-load the weights from one that I trained earlier:\n",
    "#model.load_state_dict(torch.load('./models/galaxyCNN_legacy.pth'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7b344b1-9037-4760-a927-5283b3a0fa94",
   "metadata": {},
   "source": [
    "How accurately do we distinguish between our classes after this training? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accd46cb-18cd-44c6-97fd-ab3bb51cca3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "correct_pred = {name: 0 for name in GalaxyClasses.values()}\n",
    "total_pred = {name: 0 for name in GalaxyClasses.values()}\n",
    "\n",
    "for x, y in test_loader:\n",
    "    images, labels = x.to(device), y.to(device)\n",
    "    outputs = model(images)\n",
    "    _, predictions = torch.max(outputs, 1)\n",
    "    for label, prediction in zip(labels, predictions):\n",
    "        if label == prediction:\n",
    "            correct_pred[GalaxyClasses[label.item()]] += 1\n",
    "        total_pred[GalaxyClasses[label.item()]] += 1\n",
    "\n",
    "for name, correct_count in correct_pred.items():\n",
    "    accuracy = 100 * float(correct_count) / total_pred[name]\n",
    "    print(\"Accuracy for class {} is: {:.1f} %\".format(name, accuracy))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "508c2672-1945-4c74-a9ee-900703a1882d",
   "metadata": {},
   "source": [
    "## Question: How do we add additional regularization? \n",
    "The stoschastic gradient descent (SGD) optimizer's `weight_decay` parameter implements L2-regularization on the network weights, which we learned on Monday can improve generalizability. Try re-running the above network with the new optimizer and the `weight_decay` parameter set. Compute the L2-norm of the weights in each case and compare them. Do the same for multiple values of `weight_decay`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39985f1f-cd30-4450-a651-2c165292486b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: compute L2 norm of earlier trained weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f5700d2-7709-4e37-967d-174799a6d925",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = #TODO: Implement stochastic gradient descent optimizer, varying weight_decay.\n",
    "criterion = nn.CrossEntropyLoss() \n",
    "\n",
    "for epoch in trange(50):  # loop over the dataset multiple times\n",
    "    for i, data in enumerate(train_loader, 0):\n",
    "            # get the inputs; data is a list of [inputs, labels]\n",
    "            inputs, labels = data\n",
    "\n",
    "            # zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward + backward + optimize\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "print('Finished Training')\n",
    "\n",
    "#save the model \n",
    "torch.save(model.state_dict(), './models/galaxyCNN.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1571bba7-2f02-40f2-9ad2-8197eb0031f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: compute L2 norm of the new model's trained weights. How does it compare?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476f5580-2b4d-4470-b976-36b44d55de32",
   "metadata": {},
   "source": [
    "Our results aren't terrible for a quick training process. Now let's evaluate our model on a subset of the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9178658a-363b-48bc-a3a1-ed0924ea00b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# snag one batch of the test set\n",
    "# TODO: Re-evaluate this cell multiple times to see how the results change. Where does the model struggle?\n",
    "dataiter = iter(test_loader)\n",
    "images, labels = next(dataiter)\n",
    "\n",
    "_, axs = plt.subplots(1, batch_size, figsize=(12, 12))\n",
    "axs = axs.flatten()\n",
    "\n",
    "# show all images in greyscale\n",
    "for i, ax in enumerate(axs):\n",
    "    ax.imshow(images[i].squeeze(0, 1), cmap='gray')\n",
    "    \n",
    "# apply the model to the images, and convert probability scores to classifications    \n",
    "outputs = model(images)\n",
    "_, predicted = torch.max(outputs, 1)\n",
    "\n",
    "# print the results alongside the images\n",
    "print('Ground Truth:     ', '      '.join(f'{GalaxyClasses[labels.numpy()[j]]:5s}' for j in range(batch_size)))\n",
    "print('Model Prediction: ', '      '.join(f'{GalaxyClasses[predicted.numpy()[j]]:5s}' for j in range(batch_size)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f9d566-67a8-409e-aaf1-d363e8bafbb0",
   "metadata": {},
   "source": [
    "## Challenge: Print a few batches of test data and the corresponding model predictions. Which images does the model correctly classify? Do they have anything in common (either classes or features)? \n",
    "The first tip for interpreting your model is _looking at your data._"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07dcb3cc-245f-4177-9e0b-d4e1cf585d7d",
   "metadata": {},
   "source": [
    "How can we use SHAP values to understand how the individual pixels in the image contribute to the final classification? If we consider each pixel as a feature, we can use our SHAP framework as before. This Shapley estimator is modified from the Deep Learning Important FeaTures (DeepLIFT) method (Shrikumar et al., 2017)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c619004f-c398-45ff-aba2-cd1dfbbb8465",
   "metadata": {},
   "outputs": [],
   "source": [
    "# since shuffle=True, this is a random sample of test data. We increase the batch_size so that we have more images in our background\n",
    "background_size = # TODO: Choose a background batch size, and then vary it.  \n",
    "background_loader = DataLoader(test, shuffle=True, batch_size=background_size)\n",
    "background_batch = next(iter(background_loader))\n",
    "background_images, background_labels = background_batch\n",
    "\n",
    "test_batch = next(iter(test_loader))\n",
    "test_images, test_labels = test_batch\n",
    "\n",
    "e = shap.DeepExplainer(model, background_images);\n",
    "shap_values = e.shap_values(test_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "818fd505-eb72-4a03-9c05-d00c659c7ff4",
   "metadata": {},
   "source": [
    "It's important to note here that a \"background\" estimate is required across multiple images. As before, SHAP values describe how individual values push an estimate _away from the overall expectation value of the model._ Here, the expectation value is approximated using a finite sample of images. As we increase the number of background images, we more closely approximate the true shapley values for each pixel, but our method will be more computationally expensive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b280886f-b118-4e68-bb7d-4ee30e5573be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We re-arrange our grid of shapley values and test images so that we can show them side-by-side\n",
    "shap_numpy = [np.swapaxes(np.swapaxes(s, 1, -1), 1, 2) for s in shap_values]\n",
    "test_numpy = np.swapaxes(np.swapaxes(test_images.numpy(), 1, -1), 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e21273-30e5-42a7-8d6c-2894a6561e39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, plot them\n",
    "title_row = list(GalaxyClasses.values())\n",
    "titles = np.vstack([title_row]*len(test_images))\n",
    "\n",
    "shap.image_plot(shap_numpy, test_numpy, cmap='coolwarm', labels=titles)\n",
    "\n",
    "_, predicted = torch.max(model(test_images), 1)\n",
    "print('Ground Truth:     ', '    '.join(f'{GalaxyClasses[test_labels.numpy()[j]]:5s}' for j in range(len(test_images))))\n",
    "print('Model Prediction: ', '    '.join(f'{GalaxyClasses[predicted.numpy()[j]]:5s}' for j in range(len(test_images))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "747fec6c-b209-4422-8d2d-702b80627c5a",
   "metadata": {},
   "source": [
    "## Challenge: Play around with different background data sizes (and different images) and see how this changes the resultant SHAP values for each image. How many background images do you need for reasonable Shapley estimates?\n",
    "Think for a bit, then check the solutions notebook for the answer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f332a195-992a-459b-a7d9-a898385dd1b5",
   "metadata": {},
   "source": [
    "### Does this have anything to do with saliency maps? \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a137b0b-8b5c-4e5e-b27a-935d09e71092",
   "metadata": {},
   "source": [
    "Both saliency maps and pixel-based SHAP are 'pixel attribution methods', although saliency maps use image gradients whereas SHAP uses perturbations from the original pixels in each image. We'll talk more about saliency maps later on in the tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73d29ba6-4981-490c-a7fc-edc0c849120d",
   "metadata": {},
   "source": [
    "## Conclusion: There are pros and cons to SHAP values. \n",
    "On the one hand, it can be useful to estimate the contribution of _a specific feature_ toward the models where it is used relative to _all other possible models_ that don't use it. On the other hand, it is difficult (but not impossible) to use SHAP values to explore correlated features, and SHAP values are always measured relative to the average prediction of the model. It can also be slow for many features.\n",
    "\n",
    "An open question is how to adapt the SHAP framework for interpreting features of temporal datasets, since with e.g., Recurrent Neural Networks, each recurrent unit uses information from elsewhere in the time-series sample. This could be a useful summer school hack!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdbac1c2-e665-4a2f-a3a5-87278632b80a",
   "metadata": {},
   "source": [
    "### Does this have anything to do with uncertainty quantification?\n",
    "Let's return for a second to our equation for Shapley values:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1a7f64b-af12-4af6-a7ea-1a3f1330fcee",
   "metadata": {},
   "source": [
    "$$\n",
    "\\phi_i = \\frac{1}{N!}\\sum_{R} \\left[f(x_R \\cup \\{i\\}) - f(x_R) \\right]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09de1336-1e34-4d7a-bf54-68b80fe19152",
   "metadata": {},
   "source": [
    "This function calculates the average marginal contribution of each feature (or player) $i$ to a value function $f$, which in our case was our model. What if we replaced $f$ with some other value function $v$? \n",
    "\n",
    "Two months ago, [a paper was released](https://arxiv.org/pdf/2306.05724.pdf) which defined $v$ as the conditional entropy of a given model. Under this framework, it was proposed, one could estimate the contribution of each feature not to the output of the model but to _its uncertainty_. Let's define a new network, very similar to our CNN above but trained to predict its entropy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cf1b28-af8f-4e6b-87d2-d81759ca7e2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntropyNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EntropyNet, self).__init__()\n",
    "\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Conv2d(1, 3, kernel_size = 3, stride = 1, padding = 1),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(3, 20, kernel_size=5),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.fc_layers = nn.Sequential(\n",
    "            nn.Linear(720, 50),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(50, 3),\n",
    "            nn.Softmax(dim=1)\n",
    "        )\n",
    "\n",
    "    # we now define the entropy of the model:\n",
    "    def entropy(self, x):\n",
    "        _x = x\n",
    "        logx = torch.log(_x)\n",
    "        out = _x * logx\n",
    "        out = torch.sum(out, 1)\n",
    "        out = out[:, None]\n",
    "        return -out\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv_layers(x)\n",
    "        x = x.view(-1, 3*16*15)\n",
    "        x = self.fc_layers(x)\n",
    "        #the only change - calculate the entropy of the prediction:\n",
    "        x = self.entropy(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2b836db-7761-43db-8fb8-83ae55cfccd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "entropyModel = EntropyNet().to(device)\n",
    "entropyModel.load_state_dict(torch.load('./models/galaxyCNN.pth'))\n",
    "#to instead use the pre-trained weights:\n",
    "#entropyModel.load_state_dict(torch.load('./models/galaxyCNN_legacy.pth')) \n",
    "entropyModel.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3f0df1f-919f-4dd9-933c-1d56a3e07f93",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_batch = next(iter(test_loader))\n",
    "test_images, test_labels = test_batch\n",
    "\n",
    "entropies = entropyModel(test_images)\n",
    "entropy_explainer = shap.DeepExplainer(entropyModel, background_images)\n",
    "entropy_shap = entropy_explainer.shap_values(test_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96958817-b054-4163-b047-d8ef9ecb93b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We re-arrange our grid of shapley values and test images so that we can show them side-by-side\n",
    "entropy_numpy = [np.swapaxes(np.swapaxes(entropy_shap, 1, -1), 1, 2)]\n",
    "test_numpy = np.swapaxes(np.swapaxes(test_images.numpy(), 1, -1), 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4f5beb4-0a9a-4ec2-a0a0-cd33e07892f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, plot them\n",
    "title_row = 'Uncertainty'\n",
    "titles = np.vstack([title_row]*len(test_images))\n",
    "print([GalaxyClasses[x.item()] for x in test_labels])\n",
    "shap.image_plot(entropy_numpy, test_numpy, cmap='coolwarm', labels=titles)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "iaifi_summerschool",
   "language": "python",
   "name": "iaifi_summerschool"
  },
  "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}