{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb81cda8",
   "metadata": {},
   "source": [
    "# Death to Cluster Models; Long Live Context Encoders\n",
    "\n",
    "Modeling allows us to estimate useful parameters from data while ignoring noise, and then interpret these parameters as general mechanisms in complex systems. Traditional approaches to parameter estimation break down when true parameters vary continuously from sample to sample, which is often the case in biological systems. Our framework remedies this by estimating N models for N samples simultaneously, sharing information between each estimation task. We do this by learning to estimate parameters as a function of sample context, combining modern deep learning with classic statistical inference.\n",
    "\n",
    "\n",
    "Outline of this demo:\n",
    "- Linear Models: population vs. cluster vs. contextualized models, compared on prediction mean squared error (MSE) and parameter MSE\n",
    "- Graphical Models: population vs. cluster vs. contextualized models, compared on prediction MSE and parameter MSE\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fef98aa",
   "metadata": {},
   "source": [
    "## Section 1: Regression Models\n",
    "\n",
    "### 1.1: Problem Overview\n",
    "Say we have a dataset of patients with a disease, where we've measured each patient's response to treatments A and B under an uninformed (randomized) dosing of each treatment, as well as their risk factors (diet, exercise). We know that diet and exercise have some effect on response to treatment, but it isn't a simple or obvious relationship. Formally, we'll define this as a linear model with context-varying parameters:\n",
    "\n",
    "```\n",
    "Context:            C0, C1 ~ Uniform[-1, 1]\n",
    "Treatment dose:     X0, X1 ~ Uniform[-1, 1]\n",
    "Treatment effect:   W0 = C0 + C1^2, W1 = C0^2 + C1\n",
    "Treatment response: Y = (X0 W0) + (X1 W1)\n",
    "```\n",
    "\n",
    "We want to estimate the treatment effect $\\hat{W}$ based on observations $C, X, Y$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c283856b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams.update({'font.size': 14})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f82b434f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate our data and do some plotting\n",
    "n_z = 75\n",
    "Z_space = np.linspace(-1, 1, n_z)\n",
    "Z = np.zeros((n_z**2, 2))\n",
    "for i, z_i in enumerate(Z_space):\n",
    "    for j, z_j in enumerate(Z_space):\n",
    "        Z[i * n_z + j] = [z_i, z_j]\n",
    "z_samples = 1\n",
    "Z_labels = np.repeat(Z, z_samples, axis=0)\n",
    "C = Z_labels.copy()\n",
    "C = C[::-1]\n",
    "\n",
    "W = np.zeros_like(Z_labels)\n",
    "W[:,0] = C[:,0] + C[:,1]**2 \n",
    "W[:,1] = C[:,0]**2 + C[:,1] \n",
    "W = np.repeat(W, z_samples, axis=0)\n",
    "# W += np.random.normal(0, .1, W.shape)\n",
    "X = np.random.uniform(-1, 1, W.shape)\n",
    "Y = (X * W).sum(axis=1)[:, np.newaxis]\n",
    "# Y += np.random.normal(0, .1, Y.shape)\n",
    "\n",
    "def get_gaps(Z):\n",
    "    dist = lambda a, b: np.sqrt(np.sum((a - b)**2)) \n",
    "    gap_centers = np.array([\n",
    "        [0, 0],\n",
    "        [.5, .5],\n",
    "        [.5, -.5],\n",
    "        [-.5, .5],\n",
    "        [-.5, -.5],\n",
    "    ])\n",
    "    gap_radius = .1\n",
    "    gap_idx = np.ones(len(Z))\n",
    "    for i, z in enumerate(Z):\n",
    "        for gap_center in gap_centers:\n",
    "            if dist(z, gap_center) < gap_radius:\n",
    "                gap_idx[i] = 0\n",
    "    return gap_idx\n",
    "    \n",
    "    \n",
    "gap_idx = get_gaps(Z_labels)\n",
    "train_idx, test_idx = gap_idx == 0, gap_idx == 1\n",
    "split = lambda arr: (arr[train_idx], arr[test_idx])\n",
    "Z_train, Z_test = split(Z_labels)\n",
    "C_train_pre, C_test_pre = split(C)\n",
    "C_mean, C_std = C_train_pre.mean(axis=0), C_train_pre.std(axis=0)\n",
    "C_train = (C_train_pre - C_mean) / C_std\n",
    "C_test = (C_test_pre - C_mean) / C_std\n",
    "W_train, W_test = split(W)\n",
    "X_train, X_test = split(X)\n",
    "Y_train, Y_test = split(Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fbf9e0fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 552076 characters]\n",
      "text/plain": [
       "<Figure size 1500x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Preliminary Exploration and Plotting\n",
    "def cbar(arr):\n",
    "    c = (arr - arr.mean()) / arr.std()\n",
    "    return c\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15, 5))\n",
    "axs[0]\n",
    "\n",
    "# Plot the context and treatment effect spaces\n",
    "axs[0].scatter(C[:,0], C[:,1], c=C[:,0]+C[:,1], cmap='coolwarm')\n",
    "axs[0].set_xlabel('Eats Healthy')\n",
    "axs[0].set_xticks([-1, 0, 1])\n",
    "axs[0].set_ylabel('Exercises')\n",
    "axs[0].set_yticks([-1, 0, 1])\n",
    "axs[0].set_title('Context: Risk Factors')\n",
    "\n",
    "axs[1].scatter(W[:,0], W[:,1], c=C[:,0]+C[:,1], cmap='coolwarm')\n",
    "axs[1].set_xlabel('Effect of Treatment A per dose')\n",
    "axs[1].set_xticks([-1, 0, 1, 2])\n",
    "axs[1].set_ylabel('Effect of Treatment B per dose')\n",
    "axs[1].set_yticks([-1, 0, 1, 2])\n",
    "axs[1].set_title('True Parameters: Treatment Effect')\n",
    "\n",
    "cplot = axs[2].scatter(X[:,0], X[:,1], c=cbar(Y), s=6., cmap='RdYlGn')\n",
    "axs[2].set_xlabel('Dose of Treatment A')\n",
    "axs[2].set_xticks([-1, 0, 1])\n",
    "axs[2].set_ylabel('Dose of Treatment B')\n",
    "axs[2].set_yticks([-1, 0, 1])\n",
    "axs[2].set_title('Patient Outcomes')\n",
    "plt.colorbar(cplot, ax=axs[2], ticks=[], label='Treatment Response')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b996f57",
   "metadata": {},
   "source": [
    "### 1.2: Initial Thoughts\n",
    "\n",
    "Notice that the patient outcomes seem noisy, but there is no noise in the sample generation. \n",
    "The variance of patient responses is due to meaningful variation in each patient's model parameters.\n",
    "We never observe true parameters in real life, so it's easy to attribute variation to meaningless noise.\n",
    "Here, we won't have that luxury: we'll compare all of our estimates against the true parameters to see how different estimators can generalize over the whole space.\n",
    "\n",
    "### 1.3 Realistic Constraints\n",
    "\n",
    "Let's say our dataset only contains highly localized populations of samples but we need to generalize to a much broader population. \n",
    "Below, we visualize the observed and unobserved populations, and pick two patients of interest to see how different methods generalize near observed populations and far away from observed populations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "78aeca74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 130652 characters]\n",
      "text/plain": [
       "<Figure size 1500x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 83312 characters]\n",
      "text/plain": [
       "<Figure size 1300x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the observed and unobserved samples\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15, 5))\n",
    "\n",
    "axs[0].scatter(C_test_pre[:,0], C_test_pre[:,1], c='lightsteelblue', label='unobserved patients')\n",
    "axs[0].scatter(C_train_pre[:,0], C_train_pre[:,1], c='cornflowerblue', label='observed patients')\n",
    "axs[0].set_xlabel('Eats Healthy')\n",
    "axs[0].set_xticks([-1, 0, 1])\n",
    "axs[0].set_ylabel('Exercises')\n",
    "axs[0].set_yticks([-1, 0, 1])\n",
    "axs[0].set_title('Context: Risk Factors')\n",
    "axs[0].legend()\n",
    "\n",
    "axs[1].scatter(W_test[:,0], W_test[:,1], c='lightsteelblue', label='unobserved patients')\n",
    "axs[1].scatter(W_train[:,0], W_train[:,1], c='cornflowerblue', label='observed patients')\n",
    "axs[1].set_xlabel('Effect of Treatment A per dose')\n",
    "axs[1].set_xticks([-1, 0, 1, 2])\n",
    "axs[1].set_ylabel('Effect of Treatment B per dose')\n",
    "axs[1].set_yticks([-1, 0, 1, 2])\n",
    "axs[1].set_title('True Parameters: Treatment Effect')\n",
    "\n",
    "cplot = axs[2].scatter(X_train[:,0], X_train[:,1], c=cbar(Y_train), s=5., cmap='RdYlGn')\n",
    "axs[2].set_xlabel('Dose of Treatment A')\n",
    "axs[2].set_xticks([-1, 0, 1])\n",
    "axs[2].set_ylabel('Dose of Treatment B')\n",
    "axs[2].set_yticks([-1, 0, 1])\n",
    "axs[2].set_title('Patient Outcomes (Observed)')\n",
    "plt.colorbar(cplot, ax=axs[2], ticks=[], label='Treatment Response')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# Identify patients of interest\n",
    "C_patients_pre = np.array([\n",
    "    [.6, -.6],\n",
    "    [-1, 0]\n",
    "])\n",
    "W_patients = np.zeros_like(C_patients_pre)\n",
    "W_patients[:,0] = C_patients_pre[:,0] + C_patients_pre[:,1]**2 \n",
    "W_patients[:,1] = C_patients_pre[:,0]**2 + C_patients_pre[:,1] \n",
    "\n",
    "C_patients = (C_patients_pre - C_mean) / C_std\n",
    "X_patients = np.zeros_like(W_patients)\n",
    "Y_patients = np.zeros((len(X_patients), 1))\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(13, 5))\n",
    "axs[0].scatter(C_test_pre[:,0], C_test_pre[:,1], c='lightsteelblue', label='unobserved')\n",
    "axs[0].scatter(C_train_pre[:,0], C_train_pre[:,1], c='cornflowerblue', label='observed')\n",
    "axs[0].scatter(C_patients_pre[:,0], C_patients_pre[:,1], c='darkorange', label='patients of interest')\n",
    "axs[0].set_xlabel('Eats Healthy')\n",
    "axs[0].set_xticks([-1, 0, 1])\n",
    "axs[0].set_ylabel('Exercises')\n",
    "axs[0].set_yticks([-1, 0, 1])\n",
    "axs[0].set_title('Context: Risk Factors')\n",
    "\n",
    "axs[1].scatter(W_test[:,0], W_test[:,1], c='lightsteelblue', label='unobserved patients')\n",
    "axs[1].scatter(W_train[:,0], W_train[:,1], c='cornflowerblue', label='observed patients')\n",
    "axs[1].scatter(W_patients[:,0], W_patients[:,1], c='darkorange', label='patients of interest')\n",
    "axs[1].set_xlabel('Effect of Treatment A per dose')\n",
    "axs[1].set_xticks([-1, 0, 1, 2])\n",
    "axs[1].set_ylabel('Effect of Treatment B per dose')\n",
    "axs[1].set_yticks([-1, 0, 1, 2])\n",
    "axs[1].set_title('True Parameters: Treatment Effect')\n",
    "axs[1].legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')\n",
    "plt.tight_layout()\n",
    "# plt.savefig('figures/patient_trueeffects.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68facabd",
   "metadata": {},
   "source": [
    "### 1.4: Understanding the Problem\n",
    "\n",
    "It is important to generalize to the unobserved areas because because these gaps often represent:\n",
    "\n",
    "- Undersampled / underrepresented patient populations\n",
    "- Undocumented / undiscovered treatments\n",
    "- Rare cell types / diseases / risk factors\n",
    "\n",
    "Let's pre-define our MSE evaluation metric and visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "86eb02e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mse = lambda true, pred: ((true - pred)**2).mean()\n",
    "\n",
    "def plot_all(w_hat, patient_w_hat, model_name):\n",
    "    fig, axs = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(16, 6))\n",
    "    axs[0].scatter(W_test[:,0], W_test[:,1], c='lightsteelblue')\n",
    "    axs[0].scatter(W_train[:,0], W_train[:,1], c='cornflowerblue')\n",
    "    axs[0].scatter(w_hat[:,0], w_hat[:,1], c='crimson', label='predicted effects (all unobserved)')\n",
    "    axs[0].set_xticks([-1, 0, 1])\n",
    "    axs[0].set_yticks([-1, 0, 1])\n",
    "\n",
    "\n",
    "    axs[1].scatter(W_test[:,0], W_test[:,1], c='lightsteelblue', label='true effects (unobserved patients)')\n",
    "    axs[1].scatter(W_train[:,0], W_train[:,1], c='cornflowerblue', label='true effects (observed patients)')\n",
    "    axs[1].scatter([], [], c='crimson', label='predicted effects (unobserved patients)')  # dummy plot for legend\n",
    "    axs[1].scatter(W_patients[:,0], W_patients[:,1], c='darkorange', label='true effects (patients of interest)')\n",
    "    axs[1].scatter(patient_w_hat[:,0], patient_w_hat[:,1], c='gold', label='predicted effects (patients of interest)')\n",
    "    plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')\n",
    "    fig.add_subplot(111, frameon=False)  # subplot axis hack\n",
    "    plt.tick_params(labelcolor='none', which='both', top=False, bottom=False, left=False, right=False)\n",
    "    plt.xlabel('Effect of treatment A per dose')\n",
    "    axs[1].set_xticks([-1, 0, 1, 2])\n",
    "    plt.ylabel('Effect of treatment B per dose')\n",
    "    axs[1].set_yticks([-1, 0, 1, 2])\n",
    "    plt.title('Model Parameters: Treatment Effects')\n",
    "    plt.tight_layout()\n",
    "#     plt.savefig(f'figures/{model_name}_effects.png', dpi=300)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd5fea75",
   "metadata": {},
   "source": [
    "### 1.5: Approaching the Problem\n",
    "\n",
    "We'll compare three approaches to modeling treatment response in this patient population:\n",
    "- Population Regression: Disregard context, learn a single model for all patients\n",
    "- Context-clustered Regression: We know context has some effect on model parameters, so cluster patients according to context and learn a model for each cluster\n",
    "- Contextualized Regression: Learn the function from context to model parameters\n",
    "\n",
    "For each of these models, we'll evaluate how well they recover the following in terms of MSE:\n",
    "- True patient outcomes / treatment response\n",
    "- True model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9f881267",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define our models\n",
    "from sklearn.cluster import KMeans\n",
    "from contextualized.regression import ContextualizedRegression, RegressionTrainer\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "\n",
    "class NaiveRegression:\n",
    "    def __init__(self):\n",
    "        pass\n",
    "    \n",
    "    def fit(self, X, Y):\n",
    "        self.w = np.linalg.inv(X.T @ X) @ X.T @ Y\n",
    "        return self\n",
    "        \n",
    "    def predict_w(self, X):\n",
    "        return np.tile(self.w.T, (len(X), 1))\n",
    "    \n",
    "    def predict_y(self, X):\n",
    "        w_hat = self.predict_w(X)\n",
    "        y_hat = (X * w_hat).sum(axis=1)[:, np.newaxis]\n",
    "        return y_hat\n",
    "\n",
    "\n",
    "class ClusterRegression:\n",
    "    def __init__(self, K):\n",
    "        self.K = K\n",
    "        self.kmeans = KMeans(n_clusters=K)\n",
    "        self.models = {k: NaiveRegression() for k in range(K)}\n",
    "    \n",
    "    def fit(self, C, X, Y):\n",
    "        self.wdim = X.shape[-1]\n",
    "        self.kmeans.fit(C)\n",
    "        for k in range(self.K):\n",
    "            k_idx = self.kmeans.labels_ == k\n",
    "            X_k, Y_k = X[k_idx], Y[k_idx]\n",
    "            self.models[k].fit(X_k, Y_k)\n",
    "        return self\n",
    "            \n",
    "    def predict_l(self, C):\n",
    "        return self.kmeans.predict(C)\n",
    "    \n",
    "    def predict_w(self, C, X):\n",
    "        labels = self.predict_l(C)\n",
    "        w_hat = np.zeros_like(X)\n",
    "        for label in np.unique(labels):\n",
    "            l_idx = labels == label\n",
    "            X_l = X[l_idx]\n",
    "            w_hat[l_idx] = self.models[label].predict_w(X_l)\n",
    "        return w_hat\n",
    "    \n",
    "    def predict_y(self, C, X):\n",
    "        labels = self.predict_l(C)\n",
    "        y_hat = np.zeros((len(X), 1))\n",
    "        for label in np.unique(labels):\n",
    "            l_idx = labels == label\n",
    "            X_l = X[l_idx]\n",
    "            y_hat[l_idx] = self.models[label].predict_y(X_l)\n",
    "        return y_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ac2bf42",
   "metadata": {},
   "source": [
    "### 1.6: Population Regression\n",
    "\n",
    "In our initial pass we learn a single model from data, which turns out to be the average of all our true models.\n",
    "This has terrible coverage of our model space, and completely disregards context. \n",
    "\n",
    "The predicted effects show how this model only learns a single general trend and ignores meaningful variation among patients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "41881436",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 117296 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Treatment-response error: \t0.3237248877344565\n",
      "Effect-size error: \t\t0.4739677672379643\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 89788 characters]\n",
      "text/plain": [
       "<Figure size 505x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Population regression\n",
    "naive_model = NaiveRegression().fit(X_train, Y_train)\n",
    "\n",
    "w_hat = naive_model.predict_w(X_test)\n",
    "patient_w_hat = naive_model.predict_w(X_patients)\n",
    "y_hat = naive_model.predict_y(X_test)\n",
    "\n",
    "plot_all(w_hat, patient_w_hat, 'population')\n",
    "y_mse = mse(Y_test, y_hat)\n",
    "w_mse = mse(W_test, w_hat)\n",
    "print(f\"Treatment-response error: \\t{y_mse}\")\n",
    "print(f\"Effect-size error: \\t\\t{w_mse}\")\n",
    "\n",
    "plt.figure(figsize=(5.05,4))\n",
    "plt.scatter(X_train[:,0], X_train[:,1], c=naive_model.predict_y(X_train), cmap='RdYlGn')\n",
    "plt.xlabel('Dose of treatment A')\n",
    "plt.ylabel('Dose of treatment B')\n",
    "plt.colorbar(label='Change in Health')\n",
    "plt.title('Treatment Response (Estimated)')\n",
    "plt.tight_layout()\n",
    "# plt.savefig('figures/population_responses.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03bdc8dd",
   "metadata": {},
   "source": [
    "### 1.7: Context-clustered Regression\n",
    "\n",
    "Context-clustering does much better on the unobserved patients than the population model, but we're severely constrained by the breadth of the data we've observed.\n",
    "For the patient of interest near our observed data, this will work just fine.\n",
    "However, with higher dimensional data the dataset coverage will always be sparse and new patients rarely match an observed population exactly.\n",
    "\n",
    "Overall, we get very poor coverage of the overall parameter space, and our our outlier patient will have problems with her treatment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "68e04693",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 119572 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Treatment-response error: \t0.1008637600937118\n",
      "Effect-size error: \t\t0.1493301834696637\n"
     ]
    }
   ],
   "source": [
    "# Cluster regression\n",
    "cluster_model = ClusterRegression(15).fit(C_train, X_train, Y_train)\n",
    "w_hat = cluster_model.predict_w(C_test, X_test)\n",
    "y_hat = cluster_model.predict_y(C_test, X_test)\n",
    "patient_w_hat = cluster_model.predict_w(C_patients, X_patients)\n",
    "\n",
    "plot_all(w_hat, patient_w_hat, 'cluster_neat')\n",
    "\n",
    "y_mse = mse(Y_test, y_hat)\n",
    "w_mse = mse(W_test, w_hat)\n",
    "print(f\"Treatment-response error: \\t{y_mse}\")\n",
    "print(f\"Effect-size error: \\t\\t{w_mse}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f32a921",
   "metadata": {},
   "source": [
    "### 1.7: Context-clustered Regression with Noise\n",
    "\n",
    "When we add noise parameters to the context, the cluster approach degrades. \n",
    "Because we're not estimating a function, it's difficult to do feature selection beyond 1st or 2nd order (1 or 2-feature) effects. \n",
    "Here, the coverage we get is due to noisy context and new samples will tend to align with the wrong clusters.\n",
    "\n",
    "Clustering in this realistic case marginally improves on the population model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7f6499ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 119096 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response error: 0.1459280291885149\n",
      "Effect-size error: 0.2139779446208418\n"
     ]
    }
   ],
   "source": [
    "# Regression from noise-context clustering\n",
    "noise = np.random.normal(0, 1, (len(C), 2))\n",
    "noise_train, noise_test = split(noise)\n",
    "noise_mean, noise_std = noise_train.mean(axis=0), noise_train.std(axis=0)\n",
    "noise_train = (noise_train - noise_mean) / noise_std\n",
    "noise_test = (noise_test - noise_mean) / noise_std\n",
    "C_noise_train = np.concatenate((C_train, noise_train), axis=-1)\n",
    "C_noise_test = np.concatenate((C_test, noise_test), axis=-1)\n",
    "C_noise_patients = np.concatenate((C_patients, noise_test[:len(C_patients)]), axis=-1)\n",
    "\n",
    "cluster_model = ClusterRegression(10).fit(C_noise_train, X_train, Y_train)\n",
    "w_hat = cluster_model.predict_w(C_noise_test, X_test)\n",
    "y_hat = cluster_model.predict_y(C_noise_test, X_test)\n",
    "patient_w_hat = cluster_model.predict_w(C_noise_patients, X_patients)\n",
    "\n",
    "plot_all(w_hat, patient_w_hat, 'cluster_noisy')\n",
    "\n",
    "y_mse = mse(Y_test, y_hat)\n",
    "w_mse = mse(W_test, w_hat)\n",
    "print(f\"Response error: {y_mse}\")\n",
    "print(f\"Effect-size error: {w_mse}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cda7ade",
   "metadata": {},
   "source": [
    "### 1.8: Contextualized Regression\n",
    "\n",
    "Now let's learn the function from context to model parameters with our approach.\n",
    "We get much better generalization and coverage of the parameter space by learning how context affects our models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "79a1fd26",
   "metadata": {},
   "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/trainer/configuration_validator.py:107: PossibleUserWarning: You defined a `validation_step` but have no `val_dataloader`. Skipping val loop.\n",
      "  rank_zero_warn(\n",
      "\n",
      "  | Name      | Type             | Params\n",
      "-----------------------------------------------\n",
      "0 | metamodel | SubtypeMetamodel | 1.0 K \n",
      "-----------------------------------------------\n",
      "1.0 K     Trainable params\n",
      "0         Non-trainable params\n",
      "1.0 K     Total params\n",
      "0.004     Total estimated model params size (MB)\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": "247f418a3ce6478897ddd1b57d54fb91",
       "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",
      "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",
      "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",
      "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",
      "`Trainer.fit` stopped: `max_epochs=50` reached.\n",
      "/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": "bdb2c437afdc4934ba3b4197deb1df61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 221it [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": "ebd2a8d1b03640c9ac33c9c7f4cf3868",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 221it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 124364 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response error: 0.03500504772071586\n",
      "Effect-size error: 0.05217640476705207\n"
     ]
    }
   ],
   "source": [
    "# Regression from context encoding\n",
    "encoder_kwargs = {'width': 25, 'layers': 2}\n",
    "model = ContextualizedRegression(C.shape[-1], X.shape[-1], Y.shape[-1], encoder_kwargs=encoder_kwargs)\n",
    "train_dataloader = model.dataloader(C_train, X_train, Y_train, batch_size=1)\n",
    "test_dataloader = model.dataloader(C_test, X_test, Y_test, batch_size=100)\n",
    "patient_dataloader = model.dataloader(C_patients, X_patients, Y_patients, batch_size=1)\n",
    "\n",
    "trainer = RegressionTrainer(max_epochs=50, auto_lr_find=True)\n",
    "trainer.fit(model, train_dataloader)\n",
    "\n",
    "w_hat, _ = trainer.predict_params(model, test_dataloader)\n",
    "w_hat = w_hat.squeeze()\n",
    "y_hat = (X_test * w_hat).sum(axis=1)[:,np.newaxis]\n",
    "\n",
    "patient_w_hat, _ = trainer.predict_params(model, patient_dataloader)\n",
    "patient_w_hat = patient_w_hat.squeeze()\n",
    "\n",
    "plot_all(w_hat, patient_w_hat, 'encoder')\n",
    "\n",
    "y_mse = mse(Y_test, y_hat)\n",
    "w_mse = mse(W_test, w_hat)\n",
    "print(f\"Response error: {y_mse}\")\n",
    "print(f\"Effect-size error: {w_mse}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c417c392",
   "metadata": {},
   "source": [
    "### 1.9: Contextualized Regression on Noise\n",
    "\n",
    "For a quick sanity check, let's test what happens when we feed our context encoder __only__ noise. \n",
    "We see that for a completely uninformative context, we can do no worse than the population model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "636c631c",
   "metadata": {},
   "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",
      "\n",
      "  | Name      | Type             | Params\n",
      "-----------------------------------------------\n",
      "0 | metamodel | SubtypeMetamodel | 3.2 K \n",
      "-----------------------------------------------\n",
      "3.2 K     Trainable params\n",
      "0         Non-trainable params\n",
      "3.2 K     Total params\n",
      "0.013     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d5d9e38f27f4c81a633f876b9bc33dd",
       "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": [
      "`Trainer.fit` stopped: `max_epochs=2` reached.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ccfb8daae4f458b9427208a8f9bd799",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 221it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2484704f27b447ba22e5913fc6da93f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 221it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 118148 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response error: 0.38570995551708015\n",
      "Effect-size error: 0.569284028913694\n"
     ]
    }
   ],
   "source": [
    "# Regression from noise encoding\n",
    "encoder_kwargs = {'width': 50, 'layers': 2}\n",
    "model = ContextualizedRegression(noise.shape[-1], X.shape[-1], Y.shape[-1], encoder_kwargs=encoder_kwargs)\n",
    "train_dataloader = model.dataloader(noise_train, X_train, Y_train, batch_size=1)\n",
    "test_dataloader = model.dataloader(noise_test, X_test, Y_test, batch_size=100)\n",
    "patient_dataloader = model.dataloader(noise_test[:len(C_patients)], X_patients, Y_patients, batch_size=1)\n",
    "trainer = RegressionTrainer(max_epochs=2, auto_lr_find=True)\n",
    "trainer.fit(model, train_dataloader)\n",
    "\n",
    "w_hat, _ = trainer.predict_params(model, test_dataloader)\n",
    "w_hat = w_hat.squeeze()\n",
    "y_hat = (X_test * w_hat).sum(axis=1)[:,np.newaxis]\n",
    "\n",
    "patient_w_hat, _ = trainer.predict_params(model, patient_dataloader)\n",
    "patient_w_hat = patient_w_hat.squeeze()\n",
    "\n",
    "plot_all(w_hat, patient_w_hat, 'encoder_noise')\n",
    "\n",
    "y_mse = mse(Y_test, y_hat)\n",
    "w_mse = mse(W_test, w_hat)\n",
    "print(f\"Response error: {y_mse}\")\n",
    "print(f\"Effect-size error: {w_mse}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b698aaf",
   "metadata": {},
   "source": [
    "### 1.9: Contextualized Regression with Some Noise\n",
    "\n",
    "Now, we can use the context encoder to do automatic feature selection and knowledge distillation.\n",
    "Because only some features are useful, the context encoder learns to use these over the noise features and produces a near-identical mapping as the contextualized regression with no noise in 1.7.\n",
    "\n",
    "With interpretable deep learning frameworks like generalized additive models, we can discover the exact effect of each context feature on the model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b1164cdf",
   "metadata": {},
   "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",
      "\n",
      "  | Name      | Type             | Params\n",
      "-----------------------------------------------\n",
      "0 | metamodel | SubtypeMetamodel | 3.3 K \n",
      "-----------------------------------------------\n",
      "3.3 K     Trainable params\n",
      "0         Non-trainable params\n",
      "3.3 K     Total params\n",
      "0.013     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50463ba372fb449789a7ff04e15ea151",
       "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",
      "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",
      "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",
      "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",
      "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": {
      "image/png": "[encoded data, 139360 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response error: 0.04201372922004467\n",
      "Effect-size error: 0.062325649347889804\n"
     ]
    }
   ],
   "source": [
    "# Regression from noisy context encoding\n",
    "encoder_kwargs = {'width': 50, 'layers': 2}\n",
    "model = ContextualizedRegression(C_noise_train.shape[-1], X.shape[-1], Y.shape[-1], encoder_kwargs=encoder_kwargs)\n",
    "train_dataloader = model.dataloader(C_noise_train, X_train, Y_train, batch_size=1)\n",
    "test_dataloader = model.dataloader(C_noise_test, X_test, Y_test, batch_size=100)\n",
    "patient_dataloader = model.dataloader(C_noise_patients, X_patients, Y_patients, batch_size=1)\n",
    "\n",
    "trainer = RegressionTrainer(max_epochs=50, auto_lr_find=True)\n",
    "trainer.fit(model, train_dataloader)\n",
    "\n",
    "w_hat, _ = trainer.predict_params(model, test_dataloader)\n",
    "w_hat = w_hat.squeeze()\n",
    "y_hat = (X_test * w_hat).sum(axis=1)[:,np.newaxis]\n",
    "\n",
    "patient_w_hat, _ = trainer.predict_params(model, patient_dataloader)\n",
    "patient_w_hat = patient_w_hat.squeeze()\n",
    "\n",
    "plot_all(w_hat, patient_w_hat, 'encoder_somenoise')\n",
    "\n",
    "y_mse = mse(Y_test, y_hat)\n",
    "w_mse = mse(W_test, w_hat)\n",
    "print(f\"Response error: {y_mse}\")\n",
    "print(f\"Effect-size error: {w_mse}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26bdceb7",
   "metadata": {},
   "source": [
    "### 1.10: Throw a deep learner at it!\n",
    "\n",
    "Given that our approach requires the use of a deep learner, it is natural to wonder whether the \"secret sauce\" here is our formulation or the deep learner itself.\n",
    "Specificially, we believe our approach succeeds because it enforces an explicit representation of model parameters in the deep learner, but we can also see what happens when this explicit representation is removed.\n",
    "For linear models, we can use the differentiability of neural networks to discover implicit representations of feature derivatives (i.e. linear model parameters)!\n",
    "\n",
    "By training a standard multi-layer perceptron (neural network) to predict outcomes from context and dose, we can then use standard backpropagation to find the gradient of the outcome with respect to the dose features and infer the implicit parameters.\n",
    "\n",
    "Below, you'll see that the implicit model representation is highly disorganized compared to our approach.\n",
    "We believe this might be because our approach reduces the hypothesis space of potential models, requiring a much smaller sample complexity to learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9760532e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 157564 characters]\n",
      "text/plain": [
       "<Figure size 1600x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response error: 0.04875406250357628\n",
      "Effect-size error: 0.10195370763540268\n"
     ]
    }
   ],
   "source": [
    "# gradient models from MLP(CX) -> Y\n",
    "from contextualized.modules import MLP\n",
    "\n",
    "def get_grad(model, X):\n",
    "    X.requires_grad = True\n",
    "    for i in range(len(X)):\n",
    "        yhat = model(X)\n",
    "        yhat[i].backward()\n",
    "    w_hat = X.grad.clone()\n",
    "    X.requires_grad = False\n",
    "    return w_hat\n",
    "\n",
    "CX = torch.tensor(np.concatenate((C, X), axis=-1), dtype=torch.float32)\n",
    "CX_patients = torch.tensor(np.concatenate((C_patients, X_patients), axis=-1), dtype=torch.float32)\n",
    "Y_train = torch.tensor(Y_train, dtype=torch.float32)\n",
    "Y_test = torch.tensor(Y_test, dtype=torch.float32)\n",
    "Y_patients = torch.tensor(Y_patients, dtype=torch.float32)\n",
    "\n",
    "CX_train, CX_test = split(CX)\n",
    "mlp = MLP(CX.shape[-1], Y.shape[-1], 50, 2)\n",
    "opt = torch.optim.Adam(mlp.parameters(), lr=1e-3)\n",
    "for _ in range(1000):\n",
    "    loss = mse(mlp(CX_train), Y_train)\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "train_w_hat_train = get_grad(mlp, CX_train)[:,-X.shape[-1]:]\n",
    "test_w_hat = get_grad(mlp, CX_test)[:,-X.shape[-1]:]\n",
    "patient_w_hat = get_grad(mlp, CX_patients)[:,-X_patients.shape[-1]:]\n",
    "plot_all(test_w_hat, patient_w_hat, 'cx_mlp')\n",
    "\n",
    "y_mse = mse(Y_test, mlp(CX_test))\n",
    "w_mse = mse(torch.tensor(W_test, dtype=torch.float32), test_w_hat)\n",
    "print(f\"Response error: {y_mse}\")\n",
    "print(f\"Effect-size error: {w_mse}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "938b181a",
   "metadata": {},
   "source": [
    "## Section 2: Graphical Models\n",
    "\n",
    "### 2.1: Problem Setup: Bayesian Networks\n",
    "We take a similar approach here, simulating true Bayesian networks according to context and evaluating population, cluster, and contextualized network models in terms of how well they represent the samples they're assigned to (model MSE) and how close the estimated parameters are to the true parameters (parameter MSE). \n",
    "\n",
    "We use [NO-TEARS](https://arxiv.org/abs/1803.01422) to learn Bayesian networks with standard optimization libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "84e7a356",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'contextualized.dags.torch_notmad'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [14], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpytorch_lightning\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Trainer\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mcontextualized\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdags\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtorch_notmad\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcallbacks\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mcontextualized\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdags\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtorch_notmad\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtorch_notmad\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m NOTMAD_model\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mcontextualized\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdags\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtorch_notmad\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdatamodules\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m CX_DataModule\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'contextualized.dags.torch_notmad'"
     ]
    }
   ],
   "source": [
    "from pytorch_lightning import Trainer\n",
    "from contextualized.dags.torch_notmad.callbacks import *\n",
    "from contextualized.dags.torch_notmad.torch_notmad import NOTMAD_model\n",
    "from contextualized.dags.torch_notmad.datamodules import CX_DataModule\n",
    "from contextualized.dags.torch_notmad.baselines import BayesianNetwork, GroupedNetworks, project_all, dag_pred, mse_loss\n",
    "from contextualized.dags.notmad_helpers.simulation import simulate_linear_sem\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "dag_pred_torch = dag_pred\n",
    "dag_pred = lambda a, b: dag_pred_torch(torch.Tensor(a), torch.Tensor(b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8574ae3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextualized.dags.torch_notmad import baselines\n",
    "import importlib\n",
    "importlib.reload(baselines)\n",
    "BayesianNetwork = baselines.BayesianNetwork\n",
    "project_all = baselines.project_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0c7a58",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 1000\n",
    "C = np.linspace(1, 2, n).reshape((n, 1))\n",
    "# C = np.ones((n, 1))\n",
    "blank = np.zeros_like(C)\n",
    "W_00 = blank\n",
    "W_01 = C-2\n",
    "# W_01 = np.ones_like(C)\n",
    "W_02 = blank\n",
    "W_03 = blank\n",
    "W_10 = blank\n",
    "W_11 = blank\n",
    "W_12 = blank\n",
    "W_13 = blank\n",
    "W_20 = blank\n",
    "W_21 = C**2\n",
    "W_22 = blank\n",
    "W_23 = blank\n",
    "W_30 = blank\n",
    "W_31 = C**3\n",
    "W_32 = C\n",
    "W_33 = blank\n",
    "W = np.array([\n",
    "    [W_00, W_01, W_02, W_03],\n",
    "    [W_10, W_11, W_12, W_13],\n",
    "    [W_20, W_21, W_22, W_23],\n",
    "    [W_30, W_31, W_32, W_33],\n",
    "]).squeeze()\n",
    "W = np.transpose(W, (2, 0, 1))\n",
    "X = np.array([simulate_linear_sem(w, 1, 'gauss')[0] for w in W])\n",
    "\n",
    "idx = np.logical_and(C > 1.7, C < 1.9).squeeze()\n",
    "test_idx = np.argwhere(idx).squeeze()\n",
    "train_idx = np.argwhere(~idx).squeeze()\n",
    "split = lambda arr: (arr[train_idx], arr[test_idx])\n",
    "W_train, W_test = split(W)\n",
    "C_train, C_test = split(C)\n",
    "X_train, X_test = split(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc51f9ca",
   "metadata": {},
   "source": [
    "### 2.2: Population Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d73a21cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Population Model\n",
    "notears = BayesianNetwork().fit(X_train)\n",
    "\n",
    "naive_preds_train = project_all(notears.predict(len(C_train)))\n",
    "naive_preds = project_all(notears.predict(len(C_test)))\n",
    "print(f\"train L2:  {mse(naive_preds_train, W_train)}\")\n",
    "print(f\"test L2:   {mse(naive_preds, W_test)}\")\n",
    "print(f\"train MSE: {mse(dag_pred(X_train, naive_preds_train), X_train)}\")\n",
    "print(f\"test MSE:  {mse(dag_pred(X_test, naive_preds), X_test)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c45266",
   "metadata": {},
   "source": [
    "### 2.3: Context-clustered Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd550abc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cluster Model\n",
    "kmeans = KMeans(n_clusters=10).fit(C_train)\n",
    "train_cluster_labels = kmeans.predict(C_train)\n",
    "test_cluster_labels = kmeans.predict(C_test)\n",
    "clustered_model = GroupedNetworks(BayesianNetwork).fit(X_train, train_cluster_labels)\n",
    "\n",
    "clustered_preds_train = project_all(clustered_model.predict(train_cluster_labels))\n",
    "clustered_preds = project_all(clustered_model.predict(test_cluster_labels))\n",
    "print(f\"train L2: {mse(clustered_preds_train, W_train)}\")\n",
    "print(f\"test L2:  {mse(clustered_preds, W_test)}\")\n",
    "print(f\"train mse: {mse(dag_pred(X_train, clustered_preds_train), X_train)}\")\n",
    "print(f\"test mse:  {mse(dag_pred(X_test, clustered_preds), X_test)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c076b8b",
   "metadata": {},
   "source": [
    "### 2.4: Contextualized Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3edb5e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "C_train_sansval, C_val, X_train_sansval, X_val = train_test_split(C_train, X_train, test_size=0.2)\n",
    "train_datamodule = CX_DataModule(C_train_sansval, X_train_sansval, batch_size=10)\n",
    "val_datamodule = CX_DataModule(C_val, X_val, batch_size=100)\n",
    "test_datamodule = CX_DataModule(C_test, X_test, batch_size=100)\n",
    "contextualized_model = NOTMAD_model(\n",
    "    train_datamodule,\n",
    "    n_archetypes=10,\n",
    "    encoder_type='mlp',\n",
    "    encoder_kwargs={\"width\": 100, \"layers\": 3},\n",
    "    sample_specific_loss_params={\"l1\": 1e-3, \"alpha\": 1e-1, \"rho\": 1e-2},\n",
    "    archetype_loss_params={\"l1\": 0.0, \"alpha\": 1e-1, \"rho\": 1e-2},\n",
    "    learning_rate=1e-3,\n",
    ")\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    monitor=\"val_loss\", \n",
    "    dirpath=f'notmad_checkpoints',\n",
    "    filename='{epoch}-{val_loss:.2f}'\n",
    ")\n",
    "es_callback = EarlyStopping(\n",
    "    monitor=\"val_loss\", \n",
    "    min_delta=0.01, \n",
    "    patience=5, \n",
    "    verbose=True, \n",
    "    mode=\"min\"\n",
    ")\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "trainer = Trainer(max_epochs=100, auto_lr_find=True, accelerator=accelerator, devices=1, callbacks=[checkpoint_callback, es_callback])\n",
    "trainer.fit(contextualized_model, train_datamodule)\n",
    "\n",
    "# Get best checkpoint by val_loss, test and save\n",
    "best_checkpoint = torch.load(checkpoint_callback.best_model_path)\n",
    "contextualized_model.load_state_dict(best_checkpoint['state_dict'])\n",
    "\n",
    "notmad_preds_train = project_all(contextualized_model(torch.Tensor(C_train)).detach().numpy())\n",
    "notmad_preds = project_all(contextualized_model(torch.Tensor(C_test)).detach().numpy())\n",
    "print(f\"train L2: {mse(notmad_preds_train, W_train)}\")\n",
    "print(f\"test L2:  {mse(notmad_preds, W_test)}\")\n",
    "print(f\"train mse: {mse(dag_pred(X_train, notmad_preds_train), X_train)}\")\n",
    "print(f\"test mse:  {mse(dag_pred(X_test, notmad_preds), X_test)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c703d04a",
   "metadata": {},
   "source": [
    "### 2.4 Visualizing Networks\n",
    "Use the code below to visualize true and predicted networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92b0da22",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2)\n",
    "i = 0\n",
    "axs[0].imshow(W_train[i], cmap='coolwarm')\n",
    "axs[0].set_title('True')\n",
    "axs[1].imshow(notmad_preds_train[i], cmap='coolwarm')\n",
    "axs[1].set_title('Predicted')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7b3dcc2",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
