{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "perfect-gossip",
   "metadata": {},
   "source": [
    "# Fitting And Analyzing Contextualized Models\n",
    "\n",
    "Let's walk through an example of Contextualized analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "convenient-karma",
   "metadata": {},
   "source": [
    "# Download and Prepare Data\n",
    "\n",
    "First, we will load the data into a standard pandas dataframe or a numpy array, and create a train / test split. There's only 1 step of preprocessing required: deciding the context variables.."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "empty-wisdom",
   "metadata": {},
   "source": [
    "````{margin}\n",
    "```{note}\n",
    "**Deciding context variables is an experiment-driven question.**\n",
    "\n",
    "Since we will typically interpret the contextualized models in terms of the predictor, it's often helpful to use interpretable variables as the predictors.\n",
    "\n",
    "In this example, we will use age, sex, and BMI as the contexts to look for context-specific predictors of diabetes progression.\n",
    "```\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "grand-angle",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.datasets import load_diabetes\n",
    "\n",
    "X, Y = load_diabetes(return_X_y=True, as_frame=True)\n",
    "Y = np.expand_dims(Y.values, axis=-1)\n",
    "C = X[['age', 'sex', 'bmi']]\n",
    "X.drop(['age', 'sex', 'bmi'], axis=1, inplace=True)\n",
    "\n",
    "seed = 1\n",
    "C_train, C_test, X_train, X_test, Y_train, Y_test = train_test_split(C, X, Y, test_size=0.20, random_state=seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "shaped-calcium",
   "metadata": {},
   "source": [
    "# Train a Contextualized Model\n",
    "\n",
    "Contextualized models follow an sklearn-like interface to make training easy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "550ebdef",
   "metadata": {},
   "source": [
    "````{margin}\n",
    "```{note}\n",
    "Common constructor keywords include:\n",
    "\n",
    "- n_bootstraps: The integer number of bootstrap runs to fit.\n",
    "- encoder_type: 'mlp' or 'ngam', which type of model to make as context encoder.\n",
    "- alpha: non-negative float, regularization strength.\n",
    "- mu_ratio: float in range (0.0, 1.0), governs how much the regularization applies to context-specific parameters or context-specific offsets.\n",
    "- l1_ratio: float in range (0.0, 1.0), governs how much the regularization penalizes l1 vs l2 parameter norms.\n",
    "\n",
    "Common fitting keywords include:\n",
    "\n",
    "- max_epochs: positive number, the maximum number of epochs to fit. Early stopping is turned on by default.\n",
    "- learning_rate: positive float, default is 1e-3.\n",
    "```\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "af41a81d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1789: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:616: UserWarning: Checkpoint directory /Users/blengerich/Dropbox/Professional/Research/Libraries/Contextualized/docs/lightning_logs/boot_0_checkpoints exists and is not empty.\n",
      "  rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
      "\n",
      "  | Name      | Type           | Params\n",
      "---------------------------------------------\n",
      "0 | metamodel | NaiveMetamodel | 958   \n",
      "---------------------------------------------\n",
      "958       Trainable params\n",
      "0         Non-trainable params\n",
      "958       Total params\n",
      "0.004     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3beea94906d5443ab14e6614c7993796",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:236: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n",
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:236: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d180ba6b8bd341f19d5ed86805af097d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00286b8ad8cb4ccfb3d88657d51eb86d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "240934ba3a234d62aad36ad5fca6aedd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7de31ad60f674b568295f34778465d6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6818c360dc72497aa1b3f785d9bac4da",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "afc275257e9c4b12849c38df0346a67e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2dfa1634f72f467bafa07933161735e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c5510019c47d43a5a18515a5f4a013ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f9c55eed586e4868bcb921fe6d15b8d0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ed5358433ce44febb4a3d023fc5db92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e6c164394244af0bc76b9cf6438f1d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d13646f522b411c8e5bf91a9dfcf723",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c82915bda5f445f8de6631d1b4626be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eaf70d000b45499c996224d5ab2863c6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "738986e17143413c993390a58b3c0426",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de7aa7fa56874c2196c1686aa1c79f52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=20` reached.\n",
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:616: UserWarning: Checkpoint directory /Users/blengerich/Dropbox/Professional/Research/Libraries/Contextualized/docs/lightning_logs/boot_1_checkpoints exists and is not empty.\n",
      "  rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
      "\n",
      "  | Name      | Type           | Params\n",
      "---------------------------------------------\n",
      "0 | metamodel | NaiveMetamodel | 958   \n",
      "---------------------------------------------\n",
      "958       Trainable params\n",
      "0         Non-trainable params\n",
      "958       Total params\n",
      "0.004     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a1c0a0d26f35489b82e0fbc472505a33",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7cb175d3c794516b5011a29b67764d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc9b27f7875b46b783c79ed9f0074cd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "08819671ba754dfeafb09eb11cf8806e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "774b5f0f08e84f4f9771484731d784f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da1b58d5cd7d48f283508fba9be979fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "722796c87b6f4a91a655cc26226f4d49",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b89985edded440bbc707d44d0e12115",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66c533e703d145c2b296566524763af2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b896c9ecac574c608e3ae1e3ddf2c561",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9162ca1d78ec4db6a85cad00734830b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95ecdcf91e324fc9b0d63c9f5d03d03e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e37d970d1da4aec9d3ebe73e03048ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a3d6afce7e04975bc3482258c98c2a0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "24960e82dd8f48cba9a9383574e26ece",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=20` reached.\n",
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:616: UserWarning: Checkpoint directory /Users/blengerich/Dropbox/Professional/Research/Libraries/Contextualized/docs/lightning_logs/boot_2_checkpoints exists and is not empty.\n",
      "  rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
      "\n",
      "  | Name      | Type           | Params\n",
      "---------------------------------------------\n",
      "0 | metamodel | NaiveMetamodel | 958   \n",
      "---------------------------------------------\n",
      "958       Trainable params\n",
      "0         Non-trainable params\n",
      "958       Total params\n",
      "0.004     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18668b5ff5ae47209ced5bfc5bbd8aa0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "19274a38bfff45d8b7bbfad2b74aac86",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c9b56d824e141429a57f98cd97447a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "988d56eb25ae4eceaf24a96a06588a6b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40f3a97352d14bc3a7c45aeecb4b7024",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02ade43686ea475b82aa80031fcf859f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39ff79b1d09e45dfb771ea1c5ca3f22a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18e59631db7b495584859e7d13183420",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a29de0c7dc3436f987a7d1c51ed994f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d48a60fbcb8d4c6b9f707c09c6f79229",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c15bafa5d85144fd97cd7edec24797c3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "facc69ce419b47c1b5919ca2a5951e74",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c12a1ab9971b4fa1ad6f6132ff1af097",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3034664e7fdc48f1a7c94f1c6a6e2da5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from contextualized.easy import ContextualizedRegressor\n",
    "\n",
    "model = ContextualizedRegressor(n_bootstraps=3)\n",
    "model.fit(C_train.values, X_train.values, Y_train,\n",
    "          encoder_type=\"mlp\", max_epochs=20,\n",
    "          learning_rate=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa62ba84",
   "metadata": {},
   "source": [
    "# Save/load the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f19c3c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextualized.utils import save, load\n",
    "\n",
    "save_path = './easy_demo_model.pt'\n",
    "save(model, path=save_path)\n",
    "model = load(save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "qualified-friendship",
   "metadata": {},
   "source": [
    "# Inspect the model predictions.\n",
    "\n",
    "We can use standard plotting tools to inspect the model predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3cd19b0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:236: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "688b8467f2cb48039761696e4a0ae681",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/loops/epoch/prediction_epoch_loop.py:174: UserWarning: Lightning couldn't infer the indices fetched for your dataloader.\n",
      "  warning_cache.warn(\"Lightning couldn't infer the indices fetched for your dataloader.\")\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b895fcbf8f04c19b34291c381b8b775",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ebb3a37ab2144e5a0d4aad42517ae04",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 52944 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.rcParams.update({'font.size': 18})\n",
    "\n",
    "plt.scatter(Y[:, 0], model.predict(C.values, X.values)[:, 0])\n",
    "plt.xlabel(\"True Value\")\n",
    "plt.ylabel(\"Predicted Value\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf6b63a2",
   "metadata": {},
   "source": [
    "# Check what the individual bootstrap models learned.\n",
    "\n",
    "Since we bootstrapping for robustness, we can also access individual bootstrap runs with the `individual_preds` keyword to get confidence intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "66e30786",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6d8906b0b694dceb3541d4af9851dcf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d33279ce37fe43f6ae883f6443e98788",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a56705b7c24c4ab598e9e4ef2550b769",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 120848 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model_preds = model.predict(C.values, X.values, individual_preds=True)\n",
    "for i, pred in enumerate(model_preds):\n",
    "    plt.scatter(Y[:, 0], pred[:, 0], label='Bootstrap {}'.format(i))\n",
    "plt.xlabel(\"True Value\")\n",
    "plt.ylabel(\"Predicted Value\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "756f3ef6",
   "metadata": {},
   "source": [
    "# Check what contextualized parameters the models learned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f7c525bf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ab820de3f174480bf889071a666478a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28b66869b2934b7d93b60ba57068a152",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89aef7b8e0824a53a764e7e18bf77094",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 119392 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 105928 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 110108 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "betas, mus = model.predict_params(C.values, individual_preds=False)\n",
    "\n",
    "# To get a sense of how the models are clusterd, let's embed the models is a 2-D space and visualize them.\n",
    "\n",
    "# Betas are shape n_samples, n_outputs, n_predictors\n",
    "# Let's squeeze out the n_outputs axis since we only have 1 output in this example.\n",
    "betas = np.squeeze(betas)\n",
    "\n",
    "from umap import UMAP\n",
    "um = UMAP(n_neighbors=5)\n",
    "model_reps = um.fit_transform(betas)\n",
    "\n",
    "from contextualized.analysis.embeddings import plot_embedding_for_all_covars\n",
    "plot_embedding_for_all_covars(model_reps, C, xlabel='Model UMAP 1', ylabel='Model UMAP 2')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f05011f1",
   "metadata": {},
   "source": [
    "## What are the homogeneous predictors?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "001d30ef",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:236: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 10 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e352ae22a58476eb986281874d30279",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ab8ee6b16e24f79be55cb32cd9516e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17d5ff54eb7042a697a22297d7b656ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54980 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9cb9b46069594425b84c2db3c09fbaf2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "558732b1dbce475789fb316b2b8bc96e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9fae55f3eae941a8ad42a127100c9abb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53744 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d9ed528dbea4f8e86899b0170720599",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6f58c417e22436bab8208b07a2910d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e7d959a613341529a1c75e2f36b39ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 63428 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2e0b838e6384d3ea2879b787963f7f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cc3798f59fa486b94a6bcd63541aa39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ebd65c27b6a458183e238b7407d767e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54148 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55716 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 46948 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 58264 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55112 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44628 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53640 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from contextualized.analysis.effects import (\n",
    "    plot_homogeneous_context_effects,\n",
    "    plot_homogeneous_predictor_effects,\n",
    "    plot_heterogeneous_predictor_effects,\n",
    ")\n",
    "\n",
    "plot_homogeneous_context_effects(\n",
    "    model, C, \n",
    "    classification=False, verbose=False, \n",
    "    ylabel=\"Diabetes Progression\")\n",
    "\n",
    "plot_homogeneous_predictor_effects(\n",
    "    model, C, X,\n",
    "    ylabel=\"Diabetes Progression\",\n",
    "    classification=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "460da0bf",
   "metadata": {},
   "source": [
    "## What are the heterogeneous predictors (effects change based on context)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "684f3805",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4b17379b04954c1599bc1755b781fec4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc38f6b36eea45d7816d54dacbc77f28",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ae61022609646708329b832cf8db413",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79cc084c1e7f422489cd62ba0fb25604",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ccd0d2fdbba41df87a8dc19ca8bedd0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f2eb216058e4228968924a0faeea12d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83c57c2ea4a042c0a4fd802a3209fe0b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d25132bb96784e21bfb872a2fdb61fa0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64f3f27d6496415aac93d8310894a5a0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 282it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55776 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 53048 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 39304 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54996 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 41508 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 51340 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 49716 characters]\n",
      "text/plain": [
       "<Figure size 800x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_heterogeneous_predictor_effects(model, C, X, min_effect_size=10,\n",
    "                  ylabel=\"Influence of\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fde22ee",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
