notebooks/FastTreeSHAP_Census_Income.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FastTreeSHAP in Census Income Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains usages and detailed comparisons of FastTreeSHAP v1, FastTreeSHAP v2 and the original TreeSHAP in **binary classification** problems using scikit-learn, XGBoost and LightGBM. It also contains the discussions of automatic algorithm selection. It may take a few minutes to run through all code in this notebook. The source of census income data is https://archive.ics.uci.edu/ml/datasets/census+income."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Python libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import roc_auc_score, accuracy_score\n",
    "import xgboost as xgb\n",
    "import lightgbm as lgb\n",
    "import fasttreeshap\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-process training and testing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# source of data: https://archive.ics.uci.edu/ml/datasets/census+income\n",
    "train = pd.read_csv(\"../data/adult_data.txt\", sep = \",\\s+\", header = None, engine = \"python\")\n",
    "test = pd.read_csv(\"../data/adult_test.txt\", sep = \",\\s+\", header = None, skiprows = 1, engine = \"python\")\n",
    "label_train = train[14].map({\"<=50K\": 0, \">50K\": 1}).tolist()\n",
    "label_test = test[14].map({\"<=50K.\": 0, \">50K.\": 1}).tolist()\n",
    "train = train.iloc[:, :-2]\n",
    "test = test.iloc[:, :-2]\n",
    "\n",
    "# one-hot-encoding on categorical features\n",
    "feature_names = [\"age\", \"workclass\", \"fnlwgt\", \"education\", \"education-num\", \"marital-status\", \"occupation\", \n",
    "                 \"relationship\", \"race\", \"sex\", \"capital-gain\", \"capital-loss\", \"hours-per-week\"]\n",
    "train.columns = feature_names\n",
    "test.columns = feature_names\n",
    "categorical_feature_names = [\"workclass\", \"education\", \"marital-status\", \"occupation\", \"relationship\", \"race\", \"sex\"]\n",
    "def dummy_transform(df):\n",
    "    for name in categorical_feature_names:\n",
    "        dummy_df = pd.get_dummies(df[name])\n",
    "        if \"?\" in dummy_df.columns.values:\n",
    "            dummy_df.drop(\"?\", axis=1, inplace=True)\n",
    "        df = pd.concat([df, dummy_df], axis=1)\n",
    "        df.drop(name, axis=1, inplace=True)\n",
    "    return df\n",
    "train = dummy_transform(train)\n",
    "test = dummy_transform(test)\n",
    "print(\"Training data has {} rows and {} columns.\".format(train.shape[0], train.shape[1])) \n",
    "print(\"Testing data has {} rows and {} columns.\".format(test.shape[0], test.shape[1])) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a random forest model using scikit-learn and compute SHAP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_estimators = 200  # number of trees in random forest model\n",
    "max_depth = 8  # maximum depth of any trees in random forest model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train a random forest model\n",
    "rf_model = RandomForestClassifier(n_estimators = n_estimators, max_depth = max_depth, random_state = 0)\n",
    "rf_model.fit(train, label_train)\n",
    "print(\"AUC on testing set is {:.2f}.\".format(roc_auc_score(label_test, rf_model.predict_proba(test)[:, 1])))\n",
    "print(\"Accuracy on testing set is {:.2f}.\".format(accuracy_score(label_test, rf_model.predict(test))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obtain total number of leaves\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model)\n",
    "num_leaves = sum(shap_explainer.model.num_nodes) - sum(sum(shap_explainer.model.children_left > 0))\n",
    "print(\"Total number of leaves is {}.\".format(num_leaves))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate memory usage of FastTreeSHAP v2 since FastTreeSHAP v2 has a stricter memory constraint than\n",
    "# TreeSHAP and FastTreeSHAP v1\n",
    "# derivation of the memory estimation can be found in Deep Dive Section in this notebook\n",
    "def memory_estimate_v2(shap_explainer, num_sample, num_feature, n_jobs):\n",
    "    max_node = max(shap_explainer.model.num_nodes)\n",
    "    max_leaves = (max_node + 1) // 2\n",
    "    max_combinations = 2**int(shap_explainer.model.max_depth)\n",
    "    phi_dim = num_sample * (num_feature + 1) * shap_explainer.model.num_outputs\n",
    "    n_jobs = os.cpu_count() if n_jobs == -1 else n_jobs\n",
    "    memory_1 = (max_leaves * max_combinations + phi_dim) * 8 * n_jobs\n",
    "    memory_2 = max_leaves * max_combinations * shap_explainer.model.values.shape[0] * 8\n",
    "    memory = min(memory_1, memory_2)\n",
    "    if memory < 1024:\n",
    "        print(\"Memory usage of FastTreeSHAP v2 is around {:.2f}B.\".format(memory))\n",
    "    elif memory / 1024 < 1024:\n",
    "        print(\"Memory usage of FastTreeSHAP v2 is around {:.2f}KB.\".format(memory / 1024))\n",
    "    elif memory / 1024**2 < 1024:\n",
    "        print(\"Memory usage of FastTreeSHAP v2 is around {:.2f}MB.\".format(memory / 1024**2))\n",
    "    else:\n",
    "        print(\"Memory usage of FastTreeSHAP v2 is around {:.2f}GB.\".format(memory / 1024**3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute SHAP values via different versions of TreeSHAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 10000  # number of samples to be explained\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v0 (i.e., original TreeSHAP)\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v0\", n_jobs = n_jobs)\n",
    "shap_values_v0 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v1\", n_jobs = n_jobs)\n",
    "shap_values_v1 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v1\n",
    "print(\"Maximum difference of SHAP values between v1 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v1 - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate memory usage of FastTreeSHAP v2 since FastTreeSHAP v2 has a stricter memory constraint than\n",
    "# TreeSHAP and FastTreeSHAP v1\n",
    "memory_estimate_v2(shap_explainer, num_sample, test.shape[1], n_jobs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v2\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v2\", n_jobs = n_jobs)\n",
    "shap_values_v2 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v2\n",
    "print(\"Maximum difference of SHAP values between v2 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v2 - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via automatic TreeSHAP algorithm selection\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"auto\", n_jobs = n_jobs)\n",
    "shap_values_auto = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_auto.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of automatically selected TreeSHAP algorithm\n",
    "# it turns out that \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm\n",
    "print(\"Maximum difference of SHAP values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare running times of different versions of TreeSHAP in computing SHAP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values/SHAP interaction values via TreeSHAP algorithm with version \"algorithm_version\"\n",
    "# (parallel on \"n_jobs\" threads)\n",
    "def run_fasttreeshap(model, sample, interactions, algorithm_version, n_jobs, num_round, num_sample, shortcut = False):\n",
    "    shap_explainer = fasttreeshap.TreeExplainer(\n",
    "        model, algorithm = algorithm_version, n_jobs = n_jobs, shortcut = shortcut)\n",
    "    run_time = np.zeros(num_round)\n",
    "    for i in range(num_round):\n",
    "        start = time.time()\n",
    "        shap_values = shap_explainer(sample.iloc[:num_sample], interactions = interactions).values\n",
    "        run_time[i] = time.time() - start\n",
    "        print(\"Round {} takes {:.3f} sec.\".format(i + 1, run_time[i]))\n",
    "    print(\"Average running time of {} is {:.3f} sec (std {:.3f} sec){}.\".format(\n",
    "        algorithm_version, np.mean(run_time), np.std(run_time), \" (with shortcut)\" if shortcut else \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 10000  # number of samples to be explained\n",
    "num_round = 3  # number of rounds to record mean and standard deviation of running time\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v0 (i.e., original TreeSHAP) multiple times and record its average running time\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = False, algorithm_version = \"v0\", n_jobs = n_jobs,\n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v1 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = False, algorithm_version = \"v1\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v2 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = False, algorithm_version = \"v2\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run automatically selected TreeSHAP algorithm multiple times and record its average running time\n",
    "# it turns out that \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = False, algorithm_version = \"auto\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute SHAP interaction values via different versions of TreeSHAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 100  # number of samples to be explained\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via FastTreeSHAP v0 (i.e., original TreeSHAP)\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v0\", n_jobs = n_jobs)\n",
    "shap_interaction_values_v0 = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_v0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v1\", n_jobs = n_jobs)\n",
    "shap_interaction_values_v1 = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_v1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v1\n",
    "print(\"Maximum difference of SHAP interaction values between v1 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_interaction_values_v1 - shap_interaction_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via automatic TreeSHAP algorithm selection\n",
    "# v1 is always preferred to v0 in any use cases, and v2 does not support interactions\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"auto\", n_jobs = n_jobs)\n",
    "shap_interaction_values_auto = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_auto.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of automatically selected TreeSHAP algorithm\n",
    "print(\"Maximum difference of SHAP interaction values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_interaction_values_auto - shap_interaction_values_v0))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare running times of different versions of TreeSHAP in computing SHAP interaction values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 100  # number of samples to be explained\n",
    "num_round = 3  # number of rounds to record mean and standard deviation of running time\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v0 (i.e., original TreeSHAP) multiple times and record its average running time\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = True, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v1 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = True, algorithm_version = \"v1\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run automatically selected TreeSHAP algorithm multiple times and record its average running time\n",
    "# v1 is always preferred to v0 in any use cases, and v2 does not support interactions\n",
    "run_fasttreeshap(\n",
    "    model = rf_model, sample = test, interactions = True, algorithm_version = \"auto\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train an XGBoost model and compute SHAP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_estimators = 200  # number of trees in XGBoost model\n",
    "max_depth = 8  # maximum depth of any trees in XGBoost model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train an XGBoost model\n",
    "xgb_model = xgb.XGBClassifier(\n",
    "    max_depth = max_depth, n_estimators = n_estimators, learning_rate = 0.1, n_jobs = -1,\n",
    "    use_label_encoder = False, eval_metric = \"logloss\", random_state = 0)\n",
    "xgb_model.fit(train, label_train)\n",
    "print(\"AUC on testing set is {:.2f}.\".format(roc_auc_score(label_test, xgb_model.predict_proba(test)[:, 1])))\n",
    "print(\"Accuracy on testing set is {:.2f}.\".format(accuracy_score(label_test, xgb_model.predict(test))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obtain total number of leaves\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model)\n",
    "num_leaves = sum(shap_explainer.model.num_nodes) - sum(sum(shap_explainer.model.children_left > 0))\n",
    "print(\"Total number of leaves is {}.\".format(num_leaves))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute SHAP values via different versions of TreeSHAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 10000  # number of samples to be explained\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via \"shortcut\" (i.e., original TreeSHAP in XGBoost package)\n",
    "# by default, parallel computing on all available cores is enabled in \"shortcut\"\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = True)\n",
    "shap_values_shortcut = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_shortcut.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v0 (i.e., original TreeSHAP in SHAP package)\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_v0 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v0\n",
    "print(\"Mean and maximum differences of SHAP values between v0 and shortcut is {:.2e} and {:.2e}.\".format(\n",
    "    np.mean(abs(shap_values_v0 - shap_values_shortcut)), np.max(abs(shap_values_v0 - shap_values_shortcut))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v1\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_v1 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v1\n",
    "print(\"Maximum difference of SHAP values between v1 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v1 - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate memory usage of FastTreeSHAP v2 since FastTreeSHAP v2 has a stricter memory constraint than\n",
    "# TreeSHAP and FastTreeSHAP v1\n",
    "memory_estimate_v2(shap_explainer, num_sample, test.shape[1], n_jobs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v2\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v2\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_v2 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v2\n",
    "print(\"Maximum difference of SHAP values between v2 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v2 - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via automatic TreeSHAP algorithm selection\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"auto\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_auto = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_auto.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of automatically selected TreeSHAP algorithm\n",
    "# it turns out that \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm\n",
    "print(\"Maximum difference of SHAP values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare running times of different versions of TreeSHAP in computing SHAP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 10000  # number of samples to be explained\n",
    "num_round = 3  # number of rounds to record mean and standard deviation of running time\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run \"shortcut\" version of TreeSHAP multiple times and record its average running time\n",
    "# by default, parallel computing on all available cores is enabled in \"shortcut\"\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = False, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v0 (i.e., original TreeSHAP) multiple times and record its average running time\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = False, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v1 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = False, algorithm_version = \"v1\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v2 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = False, algorithm_version = \"v2\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run automatically selected TreeSHAP algorithm multiple times and record its average running time\n",
    "# it turns out that \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = False, algorithm_version = \"auto\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute SHAP interaction values via different versions of TreeSHAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 100  # number of samples to be explained\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via \"shortcut\" (i.e., original TreeSHAP in XGBoost package)\n",
    "# by default, parallel computing on all available cores is enabled in \"shortcut\"\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = True)\n",
    "shap_interaction_values_shortcut = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_shortcut.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via FastTreeSHAP v0 (i.e., original TreeSHAP in SHAP package)\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_interaction_values_v0 = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_v0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v0\n",
    "print(\"Mean and maximum differences of SHAP values between v0 and shortcut is {:.2e} and {:.2e}.\".format(\n",
    "    np.mean(abs(shap_interaction_values_v0 - shap_interaction_values_shortcut)), \n",
    "    np.max(abs(shap_interaction_values_v0 - shap_interaction_values_shortcut))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"v1\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_interaction_values_v1 = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_v1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v1\n",
    "print(\"Maximum difference of SHAP interaction values between v1 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_interaction_values_v1 - shap_interaction_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via automatic TreeSHAP algorithm selection\n",
    "# v1 is always preferred to v0 in any use cases, and v2 does not support interactions\n",
    "shap_explainer = fasttreeshap.TreeExplainer(xgb_model, algorithm = \"auto\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_interaction_values_auto = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_auto.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of automatically selected TreeSHAP algorithm\n",
    "print(\"Maximum difference of SHAP interaction values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_interaction_values_auto - shap_interaction_values_v0))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare running times of different versions of TreeSHAP in computing SHAP interaction values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 100  # number of samples to be explained\n",
    "num_round = 3  # number of rounds to record mean and standard deviation of running time\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run \"shortcut\" version of TreeSHAP multiple times and record its average running time\n",
    "# by default, parallel computing on all available cores is enabled in \"shortcut\"\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = True, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v0 (i.e., original TreeSHAP) multiple times and record its average running time\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = True, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v1 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = True, algorithm_version = \"v1\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run automatically selected TreeSHAP algorithm multiple times and record its average running time\n",
    "# v1 is always preferred to v0 in any use cases, and v2 does not support interactions\n",
    "run_fasttreeshap(\n",
    "    model = xgb_model, sample = test, interactions = True, algorithm_version = \"auto\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a LightGBM model and compute SHAP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_estimators = 500  # number of trees in LightGBM model\n",
    "max_depth = 8  # maximum depth of any trees in LightGBM model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train a LightGBM model\n",
    "lgb_model = lgb.LGBMClassifier(\n",
    "    max_depth = max_depth, n_estimators = n_estimators, learning_rate = 0.1, n_jobs = -1, random_state = 0)\n",
    "lgb_model.fit(train, label_train)\n",
    "print(\"AUC on testing set is {:.2f}.\".format(roc_auc_score(label_test, lgb_model.predict_proba(test)[:, 1])))\n",
    "print(\"Accuracy on testing set is {:.2f}.\".format(accuracy_score(label_test, lgb_model.predict(test))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obtain total number of leaves\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model)\n",
    "num_leaves = sum(shap_explainer.model.num_nodes) - sum(sum(shap_explainer.model.children_left > 0))\n",
    "print(\"Total number of leaves is {}.\".format(num_leaves))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute SHAP values via different versions of TreeSHAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 10000  # number of samples to be explained\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via \"shortcut\" (i.e., original TreeSHAP in LightGBM package)\n",
    "# by default, parallel computing on all available cores is enabled in \"shortcut\"\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = True)\n",
    "shap_values_shortcut = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_shortcut.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v0 (i.e., original TreeSHAP in SHAP package)\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_v0 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v0\n",
    "print(\"Mean and maximum differences of SHAP values between v0 and shortcut is {:.2e} and {:.2e}.\".format(\n",
    "    np.mean(abs(shap_values_v0 - shap_values_shortcut[:,:,1])),\n",
    "    np.max(abs(shap_values_v0 - shap_values_shortcut[:,:,1]))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"v1\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_v1 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v1\n",
    "print(\"Maximum difference of SHAP values between v1 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v1 - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimate memory usage of FastTreeSHAP v2 since FastTreeSHAP v2 has a stricter memory constraint than\n",
    "# TreeSHAP and FastTreeSHAP v1\n",
    "memory_estimate_v2(shap_explainer, num_sample, test.shape[1], n_jobs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v2\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"v2\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_v2 = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_v2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v2\n",
    "print(\"Maximum difference of SHAP values between v2 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v2 - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via automatic TreeSHAP algorithm selection\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"auto\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_values_auto = shap_explainer(test.iloc[:num_sample]).values\n",
    "shap_values_auto.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of automatically selected TreeSHAP algorithm\n",
    "# it turns out that \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm\n",
    "print(\"Maximum difference of SHAP values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v0))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare running times of different versions of TreeSHAP in computing SHAP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 10000  # number of samples to be explained\n",
    "num_round = 3  # number of rounds to record mean and standard deviation of running time\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run \"shortcut\" version of TreeSHAP multiple times and record its average running time\n",
    "# by default, parallel computing on all available cores is enabled in \"shortcut\"\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = False, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v0 (i.e., original TreeSHAP) multiple times and record its average running time\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = False, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v1 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = False, algorithm_version = \"v1\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v2 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = False, algorithm_version = \"v2\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run automatically selected TreeSHAP algorithm multiple times and record its average running time\n",
    "# it turns out that \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = False, algorithm_version = \"auto\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute SHAP interaction values via different versions of TreeSHAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 100  # number of samples to be explained\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via FastTreeSHAP v0 (i.e., original TreeSHAP in SHAP package)\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "# \"shortcut\" of SHAP interaction values is not enabled for LightGBM in SHAP package\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"v0\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_interaction_values_v0 = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_v0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"v1\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_interaction_values_v1 = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_v1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of FastTreeSHAP v1\n",
    "print(\"Maximum difference of SHAP interaction values between v1 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_interaction_values_v1 - shap_interaction_values_v0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP interaction values via automatic TreeSHAP algorithm selection\n",
    "# v1 is always preferred to v0 in any use cases, and v2 does not support interactions\n",
    "shap_explainer = fasttreeshap.TreeExplainer(lgb_model, algorithm = \"auto\", n_jobs = n_jobs, shortcut = False)\n",
    "shap_interaction_values_auto = shap_explainer(test.iloc[:num_sample], interactions = True).values\n",
    "shap_interaction_values_auto.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# justify the correctness of automatically selected TreeSHAP algorithm\n",
    "print(\"Maximum difference of SHAP interaction values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_interaction_values_auto - shap_interaction_values_v0))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare running times of different versions of TreeSHAP in computing SHAP interaction values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sample = 100  # number of samples to be explained\n",
    "num_round = 3  # number of rounds to record mean and standard deviation of running time\n",
    "n_jobs = -1  # number of parallel threads (-1 means utilizing all available cores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v0 (i.e., original TreeSHAP) multiple times and record its average running time\n",
    "# parallel computing is not enabled in original TreeSHAP in SHAP package, but here we enable it for a fair comparison\n",
    "# on execution time\n",
    "# \"shortcut\" of SHAP interaction values is not enabled for LightGBM in SHAP package\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = True, algorithm_version = \"v0\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run FastTreeSHAP v1 multiple times and record its average running time\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = True, algorithm_version = \"v1\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run automatically selected TreeSHAP algorithm multiple times and record its average running time\n",
    "# v1 is always preferred to v0 in any use cases, and v2 does not support interactions\n",
    "run_fasttreeshap(\n",
    "    model = lgb_model, sample = test, interactions = True, algorithm_version = \"auto\", n_jobs = n_jobs, \n",
    "    num_round = num_round, num_sample = num_sample, shortcut = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deep dive into automatic algorithm selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The default value of the argument `algorithm` in the class `TreeExplainer` is `auto`, indicating that the TreeSHAP algorithm is automatically selected from `\"v0\"`, `\"v1\"` and `\"v2\"` according to the number of samples to be explained and the constraint on the allocated memory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specifically, `\"v1\"` is always perferred to `\"v0\"` in any use cases, and `\"v2\"` is perferred to `\"v1\"` when the number of samples to be explained is sufficiently large: <img src=\"https://latex.codecogs.com/svg.latex?M>2^{D+1}/D,\"/> and the memory constraint is also satisfied: <img src=\"https://latex.codecogs.com/svg.latex?min\\{(MN+L2^D)\\cdot C,\\;TL2^D\\}\\cdot8Byte<0.25\\cdot Total\\,Memory.\"/> Here *M* is the number of samples to be explained, *D* is the maximum depth of any tree, *N* is the number of features, *L* is the maximum number of leaves in any tree, *T* is the number of trees, and *C* is the number of parallel threads."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The derivation of the first criterion (threshold on number of samples) can be found in [FastTreeSHAP](https://arxiv.org/abs/2109.09847) paper. The derivation of the second criterion (memory constraint) is based on the implementation of parallel computing for FastTreeSHAP v2: Two versions of parallel computing have been implemented. Version 1 builds a parallel for-loop over all trees, which requires <img src=\"https://latex.codecogs.com/svg.latex?(MN+L2^D)\\cdot C\\cdot8Byte\"/> memory allocation (each thread has its own matrices to store both SHAP values and pre-computed values). Version 2 builds two consecutive parallel for-loops over all trees and over all samples respectively, which requires <img src=\"https://latex.codecogs.com/svg.latex?TL2^D\\cdot8Byte\"/> memory allocation (first parallel for-loop stores pre-computed values across all trees). In FastTreeSHAP package, based on empirical results, version 1 has a higher priority than version 2, i.e., version 1 will be selected if both version 1 and version 2 satisfy the memory constraint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Automatic algorithm selection in moderate models with varying number of samples to be explained"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In moderate models (i.e., memory constraint is not a big concern), `\"auto\"` selects `\"v2\"` when the number of samples to be explained exceeds a threshold as defined above, and selects `\"v1\"` otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_estimators = 200  # number of trees in random forest model\n",
    "max_depth = 8  # maximum depth of any trees in random forest model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train a random forest model\n",
    "rf_model = RandomForestClassifier(n_estimators = n_estimators, max_depth = max_depth, random_state = 0)\n",
    "rf_model.fit(train, label_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimated memory usage of FastTreeSHAP v2 shows that memory constraint is not a big concern\n",
    "# even for as many as 100,000 samples and parallel computing on all available cores\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model)\n",
    "memory_estimate_v2(shap_explainer, num_sample = 100000, num_feature = test.shape[1], n_jobs = -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When number of samples to be explained is 100, `\"auto\"` selects `\"v2\"` as the most appropriate TreeSHAP algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of samples to be explained\n",
    "num_sample = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v0 (i.e., original TreeSHAP)\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v0\")\n",
    "shap_values_v0 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v1\")\n",
    "shap_values_v1 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via FastTreeSHAP v2\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v2\")\n",
    "shap_values_v2 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via automatic TreeSHAP algorithm selection\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"auto\")\n",
    "shap_values_auto = shap_explainer(test.iloc[:num_sample]).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"auto\" selects \"v2\" as the most appropriate TreeSHAP algorithm when number of samples is 100\n",
    "print(\"Maximum difference of SHAP values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v0))))\n",
    "print(\"Maximum difference of SHAP values between auto and v1 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v1))))\n",
    "print(\"Maximum difference of SHAP values between auto and v2 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v2))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When number of samples to be explained is 50, `\"auto\"` selects `\"v1\"` as the most appropriate TreeSHAP algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of samples to be explained\n",
    "num_sample = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v0 (i.e., original TreeSHAP)\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v0\")\n",
    "shap_values_v0 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v1\")\n",
    "shap_values_v1 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via FastTreeSHAP v2\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v2\")\n",
    "shap_values_v2 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via automatic TreeSHAP algorithm selection\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"auto\")\n",
    "shap_values_auto = shap_explainer(test.iloc[:num_sample]).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"auto\" selects \"v1\" as the most appropriate TreeSHAP algorithm when number of samples is 50\n",
    "print(\"Maximum difference of SHAP values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v0))))\n",
    "print(\"Maximum difference of SHAP values between auto and v1 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v1))))\n",
    "print(\"Maximum difference of SHAP values between auto and v2 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v2))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Automatic algorithm selection in very large models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In very large models, `\"auto\"` selects `\"v1\"` instead of `\"v2\"` when the potential memory risk is detected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_estimators = 200  # number of trees in random forest model\n",
    "max_depth = 20  # maximum depth of any trees in random forest model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train a random forest model\n",
    "rf_model = RandomForestClassifier(n_estimators = n_estimators, max_depth = max_depth, random_state = 0)\n",
    "rf_model.fit(train, label_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimated memory usage of FastTreeSHAP v2 shows a potential memory risk\n",
    "# even for only 10 samples and no parallel computing\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model)\n",
    "memory_estimate_v2(shap_explainer, num_sample = 10, num_feature = test.shape[1], n_jobs = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of samples to be explained\n",
    "num_sample = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute SHAP values via FastTreeSHAP v0 (i.e., original TreeSHAP)\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v0\")\n",
    "shap_values_v0 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via FastTreeSHAP v1\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v1\")\n",
    "shap_values_v1 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via FastTreeSHAP v2\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"v2\")\n",
    "shap_values_v2 = shap_explainer(test.iloc[:num_sample]).values\n",
    "\n",
    "# compute SHAP values via automatic TreeSHAP algorithm selection\n",
    "shap_explainer = fasttreeshap.TreeExplainer(rf_model, algorithm = \"auto\")\n",
    "shap_values_auto = shap_explainer(test.iloc[:num_sample]).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"v2\" is automatically switched to \"v1\" as potential memory risk is detected\n",
    "print(\"Maximum difference of SHAP values between v2 and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v2 - shap_values_v0))))\n",
    "print(\"Maximum difference of SHAP values between v2 and v1 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_v2 - shap_values_v1))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"auto\" selects \"v1\" as the most appropriate TreeSHAP algorithm as potential memory risk is detected\n",
    "print(\"Maximum difference of SHAP values between auto and v0 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v0))))\n",
    "print(\"Maximum difference of SHAP values between auto and v1 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v1))))\n",
    "print(\"Maximum difference of SHAP values between auto and v2 is {:.2e}.\".format(\n",
    "    np.max(abs(shap_values_auto - shap_values_v2))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}