{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "825101a0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "2bee376a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from abc import ABC, abstractmethod\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import IterableDataset, DataLoader, random_split\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from contextualized.modules import NGAM, MLP, SoftSelect, Explainer\n",
    "\n",
    "\n",
    "ENCODERS = {\n",
    "    'mlp': MLP,\n",
    "    'ngam': NGAM,\n",
    "}\n",
    "MODELS = ['multivariate', 'univariate']\n",
    "METAMODELS = ['simple', 'subtype', 'multitask', 'tasksplit']\n",
    "LINK_FNS = [\n",
    "    lambda x: x,\n",
    "    lambda x: F.softmax(x, dim=1)\n",
    "]\n",
    "\n",
    "\n",
    "class CRTrainer(pl.Trainer):\n",
    "    def predict_coefs(self, model, dataloader):\n",
    "        preds = super().predict(model, dataloader)\n",
    "        return model._coef_preds(preds, dataloader)\n",
    "    \n",
    "    def predict_y(self, model, dataloader):\n",
    "        preds = super().predict(model, dataloader)\n",
    "        return model._y_preds(preds, dataloader)\n",
    "\n",
    "\n",
    "class Dataset:\n",
    "    def __init__(self, C, X, Y, dtype=torch.float):\n",
    "        \"\"\"\n",
    "        C: (n x c_dim)\n",
    "        X: (n x x_dim)\n",
    "        Y: (n x y_dim)\n",
    "        \"\"\"\n",
    "        self.C = torch.tensor(C, dtype=dtype)\n",
    "        self.X = torch.tensor(X, dtype=dtype)\n",
    "        self.Y = torch.tensor(Y, dtype=dtype)\n",
    "        self.n_i = 0\n",
    "        self.x_i = 0\n",
    "        self.y_i = 0\n",
    "        self.n = C.shape[0]\n",
    "        self.c_dim = C.shape[-1]\n",
    "        self.x_dim = X.shape[-1]\n",
    "        self.y_dim = Y.shape[-1]\n",
    "        self.dtype = dtype\n",
    "\n",
    "    def __iter__(self):\n",
    "        self.n_i = 0\n",
    "        self.x_i = 0\n",
    "        self.y_i = 0\n",
    "        return self\n",
    "    \n",
    "    @abstractmethod\n",
    "    def __next__(self):\n",
    "        pass\n",
    "    \n",
    "    @abstractmethod\n",
    "    def __len__(self):\n",
    "        pass\n",
    "\n",
    "\n",
    "class MultivariateDataset(Dataset):\n",
    "    def __next__(self):\n",
    "        pass\n",
    "        \n",
    "    def __len(self):\n",
    "        return self.n\n",
    "    \n",
    "class UnivariateDataset(Dataset):\n",
    "    def __next__(self):\n",
    "        pass\n",
    "        \n",
    "    def __len(self):\n",
    "        return self.n\n",
    "    \n",
    "class MultitaskMultivariateDataset(Dataset):\n",
    "    def __next__(self):\n",
    "        pass\n",
    "        \n",
    "    def __len(self):\n",
    "        return self.n * self.y_dim\n",
    "    \n",
    "class MultitaskUnivariateDataset(Dataset):\n",
    "    def __next__(self):\n",
    "        if self.n_i >= self.n:\n",
    "            self.n_i = 0\n",
    "            raise StopIteration\n",
    "        t = torch.zeros(self.x_dim + self.y_dim)\n",
    "        t[self.x_i] = 1\n",
    "        t[self.x_dim + self.y_i] = 1\n",
    "        ret = (\n",
    "            self.C[self.n_i],\n",
    "            t,\n",
    "            self.X[self.n_i, self.x_i:self.x_i+1],\n",
    "            self.Y[self.n_i, self.y_i:self.y_i+1],\n",
    "            self.n_i,\n",
    "            self.x_i,\n",
    "            self.y_i,\n",
    "        )\n",
    "        self.y_i += 1\n",
    "        if self.y_i >= self.y_dim:\n",
    "            self.x_i += 1\n",
    "            self.y_i = 0\n",
    "        if self.x_i >= self.x_dim:\n",
    "            self.n_i += 1\n",
    "            self.x_i = 0\n",
    "        return ret\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.n * self.x_dim * self.y_dim    \n",
    "\n",
    "class DataIterable(IterableDataset):\n",
    "    def __init__(self, dataset):\n",
    "        self.dataset = dataset\n",
    "    \n",
    "    def __iter__(self):\n",
    "        return iter(self.dataset)\n",
    "\n",
    "\n",
    "# class ContextualizedRegressionAbstract(pl.LightningModule):\n",
    "def MSE(beta, mu, x, y, link_fn=lambda x: x):\n",
    "    y_hat = link_fn((beta * x).sum(axis=1).unsqueeze(-1) + mu)\n",
    "    residual = y_hat - y\n",
    "    return residual.pow(2).mean()\n",
    "\n",
    "\n",
    "class NaiveContextualizedRegression(pl.LightningModule):\n",
    "    def __init__(self, context_dim, x_dim, y_dim, encoder_type='mlp', univariate=False,\n",
    "                 num_archetypes=10, encoder_width=25, encoder_layers=2):\n",
    "        super().__init__()\n",
    "        self.context_dim = context_dim\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.link_fn = link_fn\n",
    "\n",
    "        encoder = ENCODERS[encoder_type]\n",
    "        mu_dim = x_dim if univariate else 1\n",
    "        out_dim = (x_dim + mu_dim) * y_dim\n",
    "        self.context_encoder = encoder(context_dim, out_dim, encoder_width, encoder_layers)\n",
    "\n",
    "    def forward(self, C):\n",
    "        W = self.context_encoder(C)\n",
    "        W = torch.reshape(out, (out.shape[0], self.y_dim, x_dim + mu_dim))\n",
    "        beta = W[:, :, :self.x_dim]\n",
    "        mu = W[:, :, self.x_dim:]\n",
    "        return beta, mu\n",
    "\n",
    "\n",
    "class SubtypeContextualizedRegression(pl.LightningModule):\n",
    "    def __init__(self, context_dim, x_dim, y_dim, encoder_type='mlp', univariate=False,\n",
    "                 num_archetypes=10, encoder_width=25, encoder_layers=2, link_fn=lambda x: x):\n",
    "        super().__init__()\n",
    "        self.context_dim = context_dim\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.link_fn = link_fn\n",
    "\n",
    "        encoder = ENCODERS[encoder_type]\n",
    "        mu_dim = x_dim if univariate else 1\n",
    "        self.context_encoder = encoder(context_dim, num_archetypes, encoder_width, encoder_layers)\n",
    "        self.explainer = Explainer(num_archetypes, (self.y_dim, x_dim + mu_dim))\n",
    "\n",
    "    def forward(self, C):\n",
    "        Z_pre = self.context_encoder(C)\n",
    "        Z = self.link_fn(Z_pre)\n",
    "        W = self.explainer(Z)\n",
    "        beta = W[:, :, :self.x_dim]\n",
    "        mu = W[:, :, self.x_dim:]\n",
    "        return beta, mu\n",
    "\n",
    "\n",
    "class MultitaskContextualizedRegression(pl.LightningModule):\n",
    "    def __init__(self, context_dim, x_dim, y_dim, encoder_type='mlp', univariate=False,\n",
    "                 num_archetypes=10, encoder_width=25, encoder_layers=2, link_fn=lambda x: x):\n",
    "        super().__init__()\n",
    "        self.context_dim = context_dim\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.link_fn = link_fn\n",
    "\n",
    "        encoder = ENCODERS[encoder_type]\n",
    "        beta_dim = 1 if univariate else x_dim\n",
    "        task_dim = y_dim + x_dim if univariate else y_dim\n",
    "        self.context_encoder = encoder(context_dim + task_dim, num_archetypes, encoder_width, encoder_layers)\n",
    "        self.explainer = Explainer(num_archetypes, (beta_dim + 1, ))\n",
    "\n",
    "    def forward(self, C, T):\n",
    "        CT = torch.cat((C, T), 1)\n",
    "        Z_pre = self.context_encoder(CT)\n",
    "        Z = self.link_fn(Z_pre)\n",
    "        W = self.explainer(Z)\n",
    "        beta = W[:, :-1]\n",
    "        mu = W[:, -1:]\n",
    "        return beta, mu\n",
    "\n",
    "    \n",
    "class TasksplitContextualizedRegression(pl.LightningModule):\n",
    "    def __init__(self, context_dim, x_dim, y_dim, encoder_type='mlp', univariate=False,\n",
    "                 num_archetypes=10, encoder_width=25, encoder_layers=2, link_fn=lambda x: x):\n",
    "        super().__init__()\n",
    "        self.context_dim = context_dim\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.link_fn = link_fn\n",
    "        self.build()\n",
    "\n",
    "    def build(self):\n",
    "        encoder = ENCODERS[encoder_type]\n",
    "        beta_dim = 1 if univariate else x_dim\n",
    "        task_dim = y_dim + x_dim if univariate else y_dim\n",
    "        self.context_encoder = encoder(context_dim, num_archetypes, encoder_width, encoder_layers)\n",
    "        self.task_encoder = encoder(task_dim, num_archetypes, encoder_width, encoder_layers)\n",
    "        self.explainer = SoftSelect((num_archetypes, num_archetypes), (beta_dim + 1, ))\n",
    "\n",
    "    def forward(self, C, T):\n",
    "        Z_c_pre = self.context_encoder(C) \n",
    "        Z_t_pre = self.task_encoder(T)\n",
    "        Z_c = self.link_fn(Z_c_pre)\n",
    "        Z_t = self.link_fn(Z_t_pre)\n",
    "        W = self.explainer(Z_c, Z_t)\n",
    "        beta = W[:, :-1]\n",
    "        mu = W[:, -1:]\n",
    "        return beta, mu\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)\n",
    "        return optimizer\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        C, T, X, Y, _, _, _ = batch\n",
    "        beta_hat, mu_hat = self(C, T)\n",
    "        loss = MSE(beta_hat, mu_hat, X, Y)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        C, T, X, Y, _, _, _ = batch\n",
    "        beta_hat, mu_hat = self(C, T)\n",
    "        loss = MSE(beta_hat, mu_hat, X, Y)\n",
    "        return loss\n",
    "    \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        C, T, X, Y, _, _, _ = batch\n",
    "        beta_hat, mu_hat = self(C, T)\n",
    "        loss = MSE(beta_hat, mu_hat, X, Y)\n",
    "        self.log_dict({'test_mse': loss})\n",
    "    \n",
    "    def predict_step(self, batch, batch_idx):\n",
    "        C, T, X, Y, _, _, _ = batch\n",
    "        beta_hat, mu_hat = self(C, T)\n",
    "        return beta_hat, mu_hat\n",
    "    \n",
    "    def _coef_preds(self, preds, dataloader):\n",
    "        ds = dataloader.dataset.dataset\n",
    "        betas = np.zeros((ds.n, \n",
    "                          ds.x_dim, \n",
    "                          ds.y_dim))\n",
    "        mus = betas.copy()\n",
    "        for (beta_hats, mu_hats), data in zip(preds, dataloader):\n",
    "            _, _, _, _, n_idx, x_idx, y_idx = data\n",
    "            for beta_hat, mu_hat, n_i, x_i, y_i in zip(beta_hats, mu_hats, n_idx, x_idx, y_idx):\n",
    "                betas[n_i, x_i, y_i] = beta_hat\n",
    "                mus[n_i, x_i, y_i] = mu_hat\n",
    "        return betas, mus\n",
    "    \n",
    "    def _y_preds(self, preds, dataloader):\n",
    "        ds = dataloader.dataset.dataset\n",
    "        ys = np.zeros((ds.n, \n",
    "                       ds.x_dim, \n",
    "                       ds.y_dim))\n",
    "        for (beta_hats, mu_hats), data in zip(preds, dataloader):\n",
    "            _, _, X, Y, n_idx, x_idx, y_idx = data\n",
    "            for beta_hat, mu_hat, x, y, n_i, x_i, y_i in zip(beta_hats, mu_hats, X, Y, n_idx, x_idx, y_idx):\n",
    "                ys[n_i, x_i, y_i] = beta_hat * y + mu_hat\n",
    "        return ys\n",
    "    \n",
    "    def dataloader(self, C, X, Y, batch_size=32):\n",
    "        return DataLoader(dataset=DataIterable(MultitaskUnivariateDataset(C, X, Y)), batch_size=batch_size)\n",
    "    \n",
    "class ContextualizedUnivariateRegression:\n",
    "    def __init__(self, context_dim, x_dim, y_dim, metamodel='tasksplit'):\n",
    "        if metamodel == 'tasksplit':\n",
    "            self.model = TasksplitContextualizedRegression(context_dim, x_dim, y_dim, univariate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "ca6a9a83",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "c_dim = 4\n",
    "x_dim = 2\n",
    "y_dim = 3\n",
    "C = torch.rand((n, c_dim)) - .5 \n",
    "W_1 = C.sum(axis=1).unsqueeze(-1) ** 2\n",
    "W_2 = - C.sum(axis=1).unsqueeze(-1)\n",
    "b_1 = C[:, 0].unsqueeze(-1)\n",
    "b_2 = C[:, 1].unsqueeze(-1)\n",
    "W_full = torch.cat((W_1, W_2), axis=1)\n",
    "b_full = b_1 + b_2\n",
    "X = torch.rand((n, x_dim)) - .5\n",
    "Y_1 = X[:, 0].unsqueeze(-1) * W_1 + b_1\n",
    "Y_2 = X[:, 1].unsqueeze(-1) * W_2 + b_2\n",
    "Y_3 = X.sum(axis=1).unsqueeze(-1)\n",
    "Y = torch.cat((Y_1, Y_2, Y_3), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "4682d470",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/_6/468mk1cx3cb080z2r5gq0k_r0000gn/T/ipykernel_38844/2009399854.py:20: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.C = torch.tensor(C, dtype=dtype)\n",
      "/var/folders/_6/468mk1cx3cb080z2r5gq0k_r0000gn/T/ipykernel_38844/2009399854.py:21: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.X = torch.tensor(X, dtype=dtype)\n",
      "/var/folders/_6/468mk1cx3cb080z2r5gq0k_r0000gn/T/ipykernel_38844/2009399854.py:22: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.Y = torch.tensor(Y, dtype=dtype)\n",
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "\n",
      "  | Name            | Type       | Params\n",
      "-----------------------------------------------\n",
      "0 | context_encoder | MLP        | 385   \n",
      "1 | task_encoder    | MLP        | 410   \n",
      "2 | explainer       | SoftSelect | 200   \n",
      "-----------------------------------------------\n",
      "995       Trainable params\n",
      "0         Non-trainable params\n",
      "995       Total params\n",
      "0.004     Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0: : 20it [00:00, 122.00it/s, loss=0.116, v_num=36]\n",
      "Predicting: 38it [00:00, 752.14it/s]\n",
      "(100, 2, 3)\n",
      "Predicting: 38it [00:00, 672.61it/s]\n",
      "(100, 2, 3)\n",
      "Testing: 0it [00:00, ?it/s]0.10918955504894257\n",
      "0.13371214270591736\n",
      "0.09991994500160217\n",
      "0.12591280043125153\n",
      "0.10867510735988617\n",
      "0.18073561787605286\n",
      "0.1461264044046402\n",
      "0.05753888934850693\n",
      "0.10406116396188736\n",
      "0.14782333374023438\n",
      "0.07600175589323044\n",
      "0.10068490356206894\n",
      "0.04651330038905144\n",
      "0.12310661375522614\n",
      "0.13345009088516235\n",
      "0.11001557111740112\n",
      "0.10752174258232117\n",
      "0.0946996882557869\n",
      "0.15984398126602173\n",
      "--------------------------------------------------------------------------------\n",
      "DATALOADER:0 TEST RESULTS\n",
      "{'test_loss': 0.1133638247847557}\n",
      "--------------------------------------------------------------------------------\n",
      "Testing: 20it [00:00, 400.90it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.1133638247847557}]"
      ]
     },
     "execution_count": 151,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reg = TasksplitContextualizedRegression(c_dim, x_dim, y_dim, univariate=True)\n",
    "dataloader = reg.dataloader(C, X, Y, batch_size=32)\n",
    "trainer = CRTrainer(max_epochs=1)\n",
    "trainer.fit(reg, dataloader)\n",
    "beta_preds, mu_preds = trainer.predict_coefs(reg, dataloader)\n",
    "y_preds = trainer.predict_y(reg, dataloader)\n",
    "trainer.test(reg, dataloader)\n",
    "# reg.reshape_preds(preds, dataloader)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "5e462338",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name            | Type       | Params\n",
      "-----------------------------------------------\n",
      "0 | context_encoder | MLP        | 385   \n",
      "1 | task_encoder    | MLP        | 410   \n",
      "2 | explainer       | SoftSelect | 200   \n",
      "-----------------------------------------------\n",
      "995       Trainable params\n",
      "0         Non-trainable params\n",
      "995       Total params\n",
      "0.004     Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: : 20it [00:00, 176.42it/s, loss=0.0792, v_num=15]\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18a266b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:correlator]",
   "language": "python",
   "name": "conda-env-correlator-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
