{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb800276",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c03deed6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([100, 4]), torch.Size([100, 2]), torch.Size([100, 3]))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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)\n",
    "C.shape, X.shape, Y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "a86d2b2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20\n",
    "w = torch.nn.functional.softmax(torch.rand((n, 3)), dim=1)\n",
    "u = torch.nn.functional.softmax(torch.rand((n, 4)), dim=1)\n",
    "v = torch.nn.functional.softmax(torch.rand((n, 5)), dim=1)\n",
    "A = torch.rand((7, 5, 4, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "cfc6ef45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 1, 7, 5, 4, 3])\n",
      "torch.Size([20, 1, 1, 1, 3, 1])\n",
      "\n",
      "torch.Size([20, 1, 7, 5, 4])\n",
      "torch.Size([20, 1, 1, 4, 1])\n",
      "\n",
      "torch.Size([20, 1, 7, 5])\n",
      "torch.Size([20, 1, 5, 1])\n",
      "\n",
      "torch.Size([20, 1, 7])\n"
     ]
    }
   ],
   "source": [
    "b = A.unsqueeze(0).expand(v.shape[0], -1, -1, -1, -1).unsqueeze(1)\n",
    "print(b.shape)\n",
    "w_batch = w.unsqueeze(-1).unsqueeze(1).unsqueeze(1).unsqueeze(1)\n",
    "print(w_batch.shape)\n",
    "print()\n",
    "b = torch.matmul(b, w_batch).squeeze(-1)\n",
    "print(b.shape)\n",
    "u_batch = u.unsqueeze(-1).unsqueeze(1).unsqueeze(1)\n",
    "print(u_batch.shape)\n",
    "print()\n",
    "b = torch.matmul(b, u_batch).squeeze(-1)\n",
    "print(b.shape)\n",
    "v_batch = v.unsqueeze(-1).unsqueeze(1)\n",
    "print(v_batch.shape)\n",
    "print()\n",
    "b = torch.matmul(b, v_batch).squeeze(-1)\n",
    "print(b.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "4a36ad29",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20\n",
    "w = torch.nn.functional.softmax(torch.rand((n, 5)), dim=1)\n",
    "u = torch.nn.functional.softmax(torch.rand((n, 4)), dim=1)\n",
    "v = torch.nn.functional.softmax(torch.rand((n, 3)), dim=1)\n",
    "batch_weights = [w, u, v]\n",
    "\n",
    "archetypes = torch.rand((7, 49, 5, 4, 3))\n",
    "# first dim is the batch size\n",
    "# last d dims are the catcols\n",
    "# middle n-d-1 dims are the archetypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "a99c4259",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 7, 49, 5, 4, 3])\n",
      "torch.Size([20, 1, 1, 1, 3, 1])\n",
      "\n",
      "torch.Size([20, 7, 49, 5, 4])\n",
      "torch.Size([20, 1, 1, 4, 1])\n",
      "\n",
      "torch.Size([20, 7, 49, 5])\n",
      "torch.Size([20, 1, 5, 1])\n",
      "\n",
      "torch.Size([20, 7, 49])\n"
     ]
    }
   ],
   "source": [
    "batch_size = batch_weights[0].shape[0]\n",
    "expand_dims = [batch_size] + [-1 for _ in range(len(archetypes.shape))]\n",
    "batch_archetypes = archetypes.unsqueeze(0).expand(expand_dims)\n",
    "d = len(batch_weights)\n",
    "for batch_w in batch_weights[::-1]:\n",
    "    batch_w = batch_w.unsqueeze(-1)\n",
    "    d = len(batch_archetypes.shape) - len(batch_w.shape)\n",
    "    for _ in range(d):\n",
    "        batch_w = batch_w.unsqueeze(1)\n",
    "    print(batch_archetypes.shape)\n",
    "    print(batch_w.shape)\n",
    "    batch_archetypes = torch.matmul(batch_archetypes, batch_w).squeeze(-1)\n",
    "    print()\n",
    "print(batch_archetypes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "f4637a26",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "class CatCol(nn.Module):\n",
    "    def __init__(self, in_dims, out_shape):\n",
    "        super(CatCol, self).__init__()\n",
    "        init_mat = torch.rand(list(out_shape) + list(in_dims)) * 2e-2 - 1e-2\n",
    "        self.archetypes = nn.parameter.Parameter(init_mat, requires_grad=True)\n",
    "\n",
    "    def forward(self, *batch_weights):\n",
    "        batch_size = batch_weights[0].shape[0]\n",
    "        expand_dims = [batch_size] + [-1 for _ in range(len(self.archetypes.shape))]\n",
    "        batch_archetypes = self.archetypes.unsqueeze(0).expand(expand_dims)\n",
    "        for batch_w in batch_weights[::-1]:\n",
    "            batch_w = batch_w.unsqueeze(-1)\n",
    "            d = len(batch_archetypes.shape) - len(batch_w.shape)\n",
    "            for _ in range(d):\n",
    "                batch_w = batch_w.unsqueeze(1)\n",
    "            batch_archetypes = torch.matmul(batch_archetypes, batch_w).squeeze(-1)\n",
    "        return batch_archetypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "a31b0145",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 7, 8, 5, 4, 3])\n",
      "torch.Size([20, 1, 1, 1, 3, 1])\n",
      "\n",
      "torch.Size([20, 7, 8, 5, 4])\n",
      "torch.Size([20, 1, 1, 4, 1])\n",
      "\n",
      "torch.Size([20, 7, 8, 5])\n",
      "torch.Size([20, 1, 5, 1])\n",
      "\n",
      "torch.Size([20, 7, 8])\n"
     ]
    }
   ],
   "source": [
    "w = torch.nn.functional.softmax(torch.ones((n, 5)), dim=1)\n",
    "u = torch.nn.functional.softmax(torch.ones((n, 4)), dim=1)\n",
    "v = torch.nn.functional.softmax(torch.ones((n, 3)), dim=1)\n",
    "catcol = CatCol((5, 4, 3), (7, 8))\n",
    "batch_archetypes = catcol(w, u, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "04c3c8e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_archetypes[0] == batch_archetypes[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "28ce5842",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NGAM(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, width=25, layers=2, activation=nn.ReLU):\n",
    "        super(NGAM, self).__init__()\n",
    "        self.intput_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        hidden_layers = lambda: [layer for _ in range(0, layers - 2) for layer in (nn.Linear(width, width), activation())]\n",
    "        nam_layers = lambda: [nn.Linear(1, width), activation()] + hidden_layers() + [nn.Linear(width, output_dim)]\n",
    "        self.nams = [nn.Sequential(*nam_layers()) for _ in range(input_dim)]\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        ret = torch.zeros((batch_size, self.output_dim))\n",
    "        for i, nam in enumerate(self.nams):\n",
    "            ret += nam(x[:, i].unsqueeze(-1))\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7133bfb5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([100, 4]), torch.Size([100, 2]), torch.Size([100, 3]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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)\n",
    "C.shape, X.shape, Y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "286f739b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\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.nn.parameter import Parameter\n",
    "\n",
    "\n",
    "def MSE(beta, mu, x, y):\n",
    "    y_hat = (beta * x).sum(axis=1).unsqueeze(-1) + mu\n",
    "    residual = y_hat - y\n",
    "    return residual.pow(2).mean()\n",
    "\n",
    "\n",
    "class SoftSelect(nn.Module):\n",
    "    \"\"\"\n",
    "    Parameter sharing for multiple context encoders:\n",
    "    Batched computation for mapping many subtypes onto d-dimensional archetypes\n",
    "    \"\"\"\n",
    "    def __init__(self, in_dims, out_shape):\n",
    "        super(SoftSelect, self).__init__()\n",
    "        init_mat = torch.rand(list(out_shape) + list(in_dims)) * 2e-2 - 1e-2\n",
    "        self.archetypes = nn.parameter.Parameter(init_mat, requires_grad=True)\n",
    "\n",
    "    def forward(self, *batch_weights):\n",
    "        batch_size = batch_weights[0].shape[0]\n",
    "        expand_dims = [batch_size] + [-1 for _ in range(len(self.archetypes.shape))]\n",
    "        batch_archetypes = self.archetypes.unsqueeze(0).expand(expand_dims)\n",
    "        for batch_w in batch_weights[::-1]:\n",
    "            batch_w = batch_w.unsqueeze(-1)\n",
    "            d = len(batch_archetypes.shape) - len(batch_w.shape)\n",
    "            for _ in range(d):\n",
    "                batch_w = batch_w.unsqueeze(1)\n",
    "            batch_archetypes = torch.matmul(batch_archetypes, batch_w).squeeze(-1)\n",
    "        return batch_archetypes\n",
    "\n",
    "\n",
    "class NGAM(nn.Module):\n",
    "    \"\"\"\n",
    "    Generalized additive model implemented with neural networks for feature-specific functions\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dim, output_dim, width=25, layers=2, activation=nn.ReLU):\n",
    "        super(NGAM, self).__init__()\n",
    "        self.intput_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        hidden_layers = lambda: [layer for _ in range(0, layers - 2) for layer in (nn.Linear(width, width), activation())]\n",
    "        nam_layers = lambda: [nn.Linear(1, width), activation()] + hidden_layers() + [nn.Linear(width, output_dim)]\n",
    "        self.nams = [nn.Sequential(*nam_layers()) for _ in range(input_dim)]\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        ret = torch.zeros((batch_size, self.output_dim))\n",
    "        for i, nam in enumerate(self.nams):\n",
    "            ret += nam(x[:, i].unsqueeze(-1))\n",
    "        return ret\n",
    "\n",
    "    \n",
    "class ContextualizedRegressionModule(nn.Module):\n",
    "    \"\"\"\n",
    "    Estimates the weights and offset in a context-specific and task-specific regression of X onto Y\n",
    "    \"\"\"\n",
    "    def __init__(self, context_dim, task_dim, beta_dim=1, context_archetypes=10, task_archetypes=10, nam_width=25, nam_layers=2, activation=nn.ReLU):\n",
    "        super(ContextualizedRegressionModule, self).__init__()\n",
    "        self.context_dim = context_dim\n",
    "        self.task_dim = task_dim\n",
    "        self.beta_dim = beta_dim\n",
    "\n",
    "        self.context_encoder = NGAM(context_dim, context_archetypes, width=nam_width, layers=nam_layers)\n",
    "        self.task_encoder = NGAM(task_dim, task_archetypes, width=nam_width, layers=nam_layers)\n",
    "        self.softselect = SoftSelect((context_archetypes, task_archetypes), (beta_dim + 1, ))\n",
    "\n",
    "    def forward(self, c, t):\n",
    "        Z_context = self.context_encoder(c)\n",
    "        Z_task = self.task_encoder(t)\n",
    "        W = self.softselect(F.softmax(Z_context, dim=1), F.softmax(Z_task, dim=1))\n",
    "        beta = W[:, :-1]\n",
    "        mu = W[:, -1:]\n",
    "        return beta, mu\n",
    "    \n",
    "\n",
    "class UnivariateDataset:\n",
    "    def __init__(self, C, X, Y, batch_size=1, dtype=torch.float, device=torch.device('cpu')):\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, device=device)\n",
    "        self.X = torch.tensor(X, dtype=dtype, device=device)\n",
    "        self.Y = torch.tensor(Y, dtype=dtype, device=device)\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.task_dim = self.x_dim + self.y_dim\n",
    "        self.batch_size = batch_size\n",
    "        self.dtype = dtype\n",
    "        self.device = device\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",
    "    def sample(self):\n",
    "        t = torch.zeros(self.task_dim)\n",
    "        t[self.x_i] = 1\n",
    "        t[self.x_dim + self.y_i] = 1\n",
    "        ret = (\n",
    "            self.C[self.n_i].unsqueeze(0),\n",
    "            t.unsqueeze(0),\n",
    "            self.X[self.n_i, self.x_i:self.x_i+1].unsqueeze(0),\n",
    "            self.Y[self.n_i, self.y_i:self.y_i+1].unsqueeze(0),\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 __next__(self):\n",
    "        \"\"\"\n",
    "        Returns a batch_size paired sample (c, t, x, y)\n",
    "        If there are fewer than batch_size samples remaining, returns n - batch_size samples\n",
    "        c: (batch_size, c_dim)\n",
    "        t: (batch_size, task_dim * 2)  [x_task y_task]\n",
    "        x: (batch_size, 1)\n",
    "        y: (batch_size, 1)\n",
    "        \"\"\"\n",
    "        if self.n_i >= self.n:\n",
    "            self.n_i = 0\n",
    "            raise StopIteration\n",
    "        C_batch, T_batch, X_batch, Y_batch = self.sample()\n",
    "        while len(C_batch) < self.batch_size and self.n_i < self.n:\n",
    "            C_s, T_s, X_s, Y_s = self.sample()\n",
    "            C_batch = torch.cat((C_batch, C_s))\n",
    "            T_batch = torch.cat((T_batch, T_s))\n",
    "            X_batch = torch.cat((X_batch, X_s))\n",
    "            Y_batch = torch.cat((Y_batch, Y_s))\n",
    "        return C_batch, T_batch, X_batch, Y_batch\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.n * self.x_dim * self.y_dim\n",
    "    \n",
    "\n",
    "class ContextualizedUnivariateRegression:\n",
    "    def __init__(self, context_dim, x_dim, y_dim, context_archetypes=10, task_archetypes=10, \n",
    "                l1=0, encoder_width=25, encoder_layers=2, activation=nn.ReLU, bootstraps=None, \n",
    "                device=torch.device('cpu')):\n",
    "        self.context_dim = context_dim\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.taskpair_dim = x_dim + y_dim\n",
    "        module_params = {\n",
    "            'context_dim': context_dim,\n",
    "            'task_dim': self.taskpair_dim,\n",
    "            'beta_dim': 1,\n",
    "            'context_archetypes': context_archetypes,\n",
    "            'task_archetypes': task_archetypes,\n",
    "            'nam_width': encoder_width,\n",
    "            'nam_layers': encoder_layers,\n",
    "            'activation': activation,\n",
    "        }\n",
    "        if bootstraps is None:\n",
    "            self.model = ContextualizedRegressionModule(**module_params)\n",
    "            self.models = [self.model]\n",
    "        else:\n",
    "            self.model = None\n",
    "            self.models = [ContextualizedRegressionModule(**module_params) for _ in range(bootstraps)]\n",
    "        self.to(device)\n",
    "\n",
    "    def to(self, device):\n",
    "        self.device = device\n",
    "        for model in self.models:\n",
    "            model.to(self.device)\n",
    "        return self\n",
    "\n",
    "    def _loss(self, outputs, X, Y):\n",
    "        beta, mu = outputs\n",
    "        mse = MSE(beta, mu, X, Y)\n",
    "        return mse\n",
    "\n",
    "    def _fit(self, model, C, X, Y, epochs, batch_size, optimizer=torch.optim.Adam, lr=1e-3, \n",
    "             validation_set=None, es_patience=np.inf, es_epoch=0, silent=False):\n",
    "        model.train()\n",
    "        device = next(model.parameters()).device\n",
    "        opt = optimizer(model.parameters(), lr=lr)\n",
    "        dataset = UnivariateDataset(C, X, Y, batch_size=batch_size, dtype=torch.float, device=device)\n",
    "        val_dataset = None\n",
    "        if validation_set is not None:\n",
    "            val_dataset = UnivariateDataset(*validation_set, batch_size=batch_size, dtype=torch.float, device=device)\n",
    "        progress_bar = tqdm(range(epochs), disable=silent)\n",
    "        min_loss = np.inf  # for early stopping\n",
    "        es_count = 0\n",
    "        for epoch in progress_bar:\n",
    "            for batch_i, (c, t, x, y) in enumerate(dataset):\n",
    "                loss = None\n",
    "                # Train\n",
    "                outputs = model(c, t)\n",
    "                loss = self._loss(outputs, x, y)\n",
    "                opt.zero_grad()\n",
    "                loss.backward()\n",
    "                opt.step()\n",
    "                # Update UI and check validation set for early stopping\n",
    "                train_desc = f'[Train MSE: {loss.item():.4f}] [Sample: {batch_size * batch_i}/{len(dataset)}] Epoch'\n",
    "                if val_dataset is not None:\n",
    "                    val_loss_batches = []\n",
    "                    for c_val, t_val, x_val, y_val in val_dataset:\n",
    "                        val_outputs = model(c_val, t_val)\n",
    "                        val_loss_batch = self._loss(val_outputs, x_val, y_val).item()\n",
    "                        val_loss_batches.append(val_loss_batch)\n",
    "                    val_loss = np.mean(val_loss_batches)\n",
    "                    train_desc = f\"[Val MSE: {val_loss:.4f}] \" + train_desc\n",
    "                    if epoch >= es_epoch:\n",
    "                        if val_loss < min_loss:\n",
    "                            min_loss = val_loss\n",
    "                            es_count = 0\n",
    "                        else:\n",
    "                            es_count += 1\n",
    "                progress_bar.set_description(train_desc)\n",
    "                if es_count > es_patience:\n",
    "                    model.eval()\n",
    "                    return\n",
    "        model.eval()\n",
    "\n",
    "    def fit(self, C, X, Y, epochs, batch_size, optimizer=torch.optim.Adam, lr=1e-3, validation_set=None, es_patience=np.inf, es_epoch=0, silent=False):\n",
    "        fit_params = {\n",
    "            'C': C, 'X': X, 'Y': Y, 'epochs': epochs, 'batch_size': batch_size,\n",
    "            'optimizer': optimizer, 'lr': lr,\n",
    "            'validation_set': validation_set, 'es_epoch': es_epoch, 'es_patience': es_patience,\n",
    "            'silent': silent,\n",
    "        }\n",
    "        if self.model:\n",
    "            fit_params['model'] = self.model\n",
    "            self._fit(**fit_params)\n",
    "        else:\n",
    "            for model in self.models:\n",
    "                boot_idx = np.random.choice(np.arange(len(X)), size=len(X), replace=True)\n",
    "                fit_params.update({\n",
    "                    'model': model, 'C': C[boot_idx], 'X': X[boot_idx], 'Y': Y[boot_idx],\n",
    "                })\n",
    "                self._fit(**fit_params)\n",
    "\n",
    "    def _predict_coefs(self, model, C):\n",
    "        \"\"\"\n",
    "        Predict a (1, c_dim, x_dim, y_dim) matrix of regression coefficients and offsets for each context\n",
    "        beta[i,j] and mu[i,j] solve the regression problem y_j = beta[i,j] * x_i + mu[i,j]\n",
    "        returns a numpy matrix\n",
    "        \"\"\"\n",
    "        n = C.shape[0]\n",
    "        betas = torch.zeros((n, self.x_dim, self.y_dim))\n",
    "        mus = torch.zeros((n, self.x_dim, self.y_dim))\n",
    "        for i in range(n):  # Predict per-sample to avoid OOM\n",
    "            C_i = torch.Tensor(C[i])\n",
    "            for t_x in range(self.x_dim):\n",
    "                for t_y in range(self.y_dim):\n",
    "                    task = torch.zeros(self.taskpair_dim)\n",
    "                    task[t_x] = 1\n",
    "                    task[self.x_dim + t_y] = 1\n",
    "                    beta, mu = model(C_i.unsqueeze(0), task.unsqueeze(0))\n",
    "                    beta, mu = beta.cpu().detach(), mu.cpu().detach()\n",
    "                    betas[i, t_x, t_y] = beta.squeeze()\n",
    "                    mus[i, t_x, t_y] = mu.squeeze()\n",
    "        return betas.unsqueeze(0).numpy(), mus.unsqueeze(0).numpy()\n",
    "\n",
    "    def predict_coefs(self, C, all_bootstraps=False):\n",
    "        \"\"\"\n",
    "        Predict an (x_dim, y_dim) matrix of regression coefficients for each context\n",
    "        Returns a numpy matrix (1 or bootstraps, n, x_dim, y_dim)\n",
    "        \"\"\"\n",
    "        betas, mus = self._predict_coefs(self.models[0], C)\n",
    "        for model in self.models[1:]:\n",
    "            betas_i, mus_i = self._predict_coefs(model, C)\n",
    "            betas = np.concatenate((betas, betas_i), axis=0)\n",
    "            mus = np.concatenate((mus, mus_i), axis=0)\n",
    "        if all_bootstraps:\n",
    "            return betas, mus\n",
    "        return betas.mean(axis=0), mus.mean(axis=0)\n",
    "\n",
    "    def predict_correlation(self, C, all_bootstraps=False):\n",
    "        \"\"\"\n",
    "        Requires x_dim == y_dim\n",
    "        Predict an (x_dim, y_dim) matrix of squared Pearson's correlation coefficients for each context\n",
    "        Returns a numpy matrix (n, x_dim, y_dim, 1 or bootstraps)\n",
    "        \"\"\"\n",
    "        betas, mus = self.predict_coefs(C, all_bootstraps=True)\n",
    "        betas_T = np.transpose(betas, (0, 1, 3, 2))\n",
    "        rho = betas * betas_T\n",
    "        if all_bootstraps:\n",
    "            return rho\n",
    "        return rho.mean(axis=0)\n",
    "\n",
    "    def get_mse(self, C, X, Y, all_bootstraps=False):\n",
    "        \"\"\"\n",
    "        Returns the MSE of the model on a dataset\n",
    "        Returns a numpy array (1 or bootstraps, )\n",
    "        \"\"\"\n",
    "        dataset = UnivariateDataset(C, X, Y, batch_size=1, device=self.device)\n",
    "        mses = np.zeros(len(self.models))\n",
    "        for i, model in enumerate(self.models):\n",
    "            for c, t, x, y in dataset:\n",
    "                betas, mus = model(c, t)\n",
    "                mse = MSE(betas, mus, x, y).item()\n",
    "                mses[i] += 1 / len(dataset) * mse\n",
    "        if not all_bootstraps:\n",
    "            return mses.mean()\n",
    "        return mses\n",
    "\n",
    "\n",
    "class MultivariateDataset:\n",
    "    def __init__(self, C, X, Y, batch_size=1, dtype=torch.float, device=torch.device('cpu')):\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, device=device)\n",
    "        self.X = torch.tensor(X, dtype=dtype, device=device)\n",
    "        self.Y = torch.tensor(Y, dtype=dtype, device=device)\n",
    "        self.n_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.task_dim = self.y_dim\n",
    "        self.batch_size = batch_size\n",
    "        self.dtype = dtype\n",
    "        self.device = device\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.n_i = 0\n",
    "        self.y_i = 0\n",
    "        return self\n",
    "        \n",
    "    def sample(self):\n",
    "        t = torch.zeros(self.task_dim)\n",
    "        t[self.y_i] = 1\n",
    "        ret = (\n",
    "            self.C[self.n_i].unsqueeze(0),\n",
    "            t.unsqueeze(0),\n",
    "            self.X[self.n_i].unsqueeze(0),\n",
    "            self.Y[self.n_i, self.y_i:self.y_i+1].unsqueeze(0),\n",
    "        )\n",
    "        self.y_i += 1\n",
    "        if self.y_i >= self.y_dim:\n",
    "            self.n_i += 1\n",
    "            self.y_i = 0\n",
    "        return ret\n",
    "    \n",
    "    def __next__(self):\n",
    "        \"\"\"\n",
    "        Returns a batch_size sample (c, t, x, y)\n",
    "        If there are fewer than batch_size samples remaining, returns n - batch_size samples\n",
    "        c: (batch_size, c_dim)\n",
    "        t: (batch_size, task_dim * 2)  [x_task y_task]\n",
    "        x: (batch_size, x_dim)\n",
    "        y: (batch_size, 1)\n",
    "        \"\"\"\n",
    "        if self.n_i >= self.n:\n",
    "            self.n_i = 0\n",
    "            raise StopIteration\n",
    "        C_batch, T_batch, X_batch, Y_batch = self.sample()\n",
    "        while len(C_batch) < self.batch_size and self.n_i < self.n:\n",
    "            C_s, T_s, X_s, Y_s = self.sample()\n",
    "            C_batch = torch.cat((C_batch, C_s))\n",
    "            T_batch = torch.cat((T_batch, T_s))\n",
    "            X_batch = torch.cat((X_batch, X_s))\n",
    "            Y_batch = torch.cat((Y_batch, Y_s))\n",
    "        return C_batch, T_batch, X_batch, Y_batch\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.n * self.y_dim\n",
    "\n",
    "\n",
    "class ContextualizedRegression:\n",
    "    def __init__(self, context_dim, x_dim, y_dim, context_archetypes=10, task_archetypes=10, \n",
    "                l1=0, encoder_width=25, encoder_layers=2, activation=nn.ReLU, bootstraps=None, \n",
    "                device=torch.device('cpu')):\n",
    "        self.context_dim = context_dim\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.task_dim = y_dim\n",
    "        module_params = {\n",
    "            'context_dim': context_dim,\n",
    "            'task_dim': y_dim,\n",
    "            'beta_dim': x_dim,\n",
    "            'context_archetypes': context_archetypes,\n",
    "            'task_archetypes': task_archetypes,\n",
    "            'nam_width': encoder_width,\n",
    "            'nam_layers': encoder_layers,\n",
    "            'activation': activation,\n",
    "        }\n",
    "        if bootstraps is None:\n",
    "            self.model = ContextualizedRegressionModule(**module_params)\n",
    "            self.models = [self.model]\n",
    "        else:\n",
    "            self.model = None\n",
    "            self.models = [ContextualizedRegressionModule(**module_params) for _ in range(bootstraps)]\n",
    "        self.to(device)\n",
    "\n",
    "    def to(self, device):\n",
    "        self.device = device\n",
    "        for model in self.models:\n",
    "            model.to(self.device)\n",
    "        return self\n",
    "\n",
    "    def _loss(self, outputs, X, Y):\n",
    "        beta, mu = outputs\n",
    "        mse = MSE(beta, mu, X, Y)\n",
    "        return mse\n",
    "\n",
    "    def _fit(self, model, C, X, Y, epochs, batch_size, optimizer=torch.optim.Adam, lr=1e-3, \n",
    "             validation_set=None, es_patience=np.inf, es_epoch=0, silent=False):\n",
    "        model.train()\n",
    "        device = next(model.parameters()).device\n",
    "        opt = optimizer(model.parameters(), lr=lr)\n",
    "        dataset = MultivariateDataset(C, X, Y, batch_size=batch_size, dtype=torch.float, device=device)\n",
    "        val_dataset = None\n",
    "        if validation_set is not None:\n",
    "            val_dataset = MultivariateDataset(*validation_set, batch_size=batch_size, dtype=torch.float, device=device)\n",
    "        progress_bar = tqdm(range(epochs), disable=silent)\n",
    "        min_loss = np.inf  # for early stopping\n",
    "        es_count = 0\n",
    "        for epoch in progress_bar:\n",
    "            for batch_i, (c, t, x, y) in enumerate(dataset):\n",
    "                loss = None\n",
    "                # Train\n",
    "                outputs = model(c, t)\n",
    "                loss = self._loss(outputs, x, y)\n",
    "                opt.zero_grad()\n",
    "                loss.backward()\n",
    "                opt.step()\n",
    "                # Update UI and check validation set for early stopping\n",
    "                train_desc = f'[Train MSE: {loss.item():.4f}] [Sample: {batch_size * batch_i}/{len(dataset)}] Epoch'\n",
    "                if val_dataset is not None:\n",
    "                    val_loss_batches = []\n",
    "                    for c_val, t_val, x_val, y_val in val_dataset:\n",
    "                        val_outputs = model(c_val, t_val)\n",
    "                        val_loss_batch = self._loss(val_outputs, x_val, y_val).item()\n",
    "                        val_loss_batches.append(val_loss_batch)\n",
    "                    val_loss = np.mean(val_loss_batches)\n",
    "                    train_desc = f\"[Val MSE: {val_loss:.4f}] \" + train_desc\n",
    "                    if epoch >= es_epoch:\n",
    "                        if val_loss < min_loss:\n",
    "                            min_loss = val_loss\n",
    "                            es_count = 0\n",
    "                        else:\n",
    "                            es_count += 1\n",
    "                progress_bar.set_description(train_desc)\n",
    "                if es_count > es_patience:\n",
    "                    model.eval()\n",
    "                    return\n",
    "        model.eval()\n",
    "\n",
    "    def fit(self, C, X, Y, epochs, batch_size, optimizer=torch.optim.Adam, lr=1e-3, validation_set=None, es_patience=np.inf, es_epoch=0, silent=False):\n",
    "        fit_params = {\n",
    "            'C': C, 'X': X, 'Y': Y, 'epochs': epochs, 'batch_size': batch_size,\n",
    "            'optimizer': optimizer, 'lr': lr,\n",
    "            'validation_set': validation_set, 'es_epoch': es_epoch, 'es_patience': es_patience,\n",
    "            'silent': silent,\n",
    "        }\n",
    "        if self.model:\n",
    "            fit_params['model'] = self.model\n",
    "            self._fit(**fit_params)\n",
    "        else:\n",
    "            for model in self.models:\n",
    "                boot_idx = np.random.choice(np.arange(len(X)), size=len(X), replace=True)\n",
    "                fit_params.update({\n",
    "                    'model': model, 'C': C[boot_idx], 'X': X[boot_idx], 'Y': Y[boot_idx],\n",
    "                })\n",
    "                self._fit(**fit_params)\n",
    "\n",
    "    def _predict_coefs(self, model, C):\n",
    "        \"\"\"\n",
    "        Predict a (1, c_dim, x_dim, y_dim) matrix of regression coefficients and offsets for each context\n",
    "        beta[i,j] and mu[i,j] solve the regression problem y_j = beta[i,j] * x_i + mu[i,j]\n",
    "        returns a numpy matrix\n",
    "        \"\"\"\n",
    "        n = C.shape[0]\n",
    "        betas = torch.zeros((n, self.x_dim, self.y_dim))\n",
    "        mus = torch.zeros((n, self.y_dim))\n",
    "        for i in range(n):  # Predict per-sample to avoid OOM\n",
    "            C_i = torch.Tensor(C[i])\n",
    "            for t_y in range(self.y_dim):\n",
    "                task = torch.zeros(self.y_dim)\n",
    "                task[t_y] = 1\n",
    "                beta, mu = model(C_i.unsqueeze(0), task.unsqueeze(0))\n",
    "                beta, mu = beta.cpu().detach(), mu.cpu().detach()\n",
    "                betas[i, :, t_y] = beta.squeeze()\n",
    "                mus[i, t_y] = mu.squeeze()\n",
    "        return betas.unsqueeze(0).numpy(), mus.unsqueeze(0).numpy()\n",
    "\n",
    "    def predict_coefs(self, C, all_bootstraps=False):\n",
    "        \"\"\"\n",
    "        Predict an (x_dim, y_dim) matrix of regression coefficients for each context\n",
    "        Returns a numpy matrix (1 or bootstraps, n, x_dim, y_dim)\n",
    "        \"\"\"\n",
    "        betas, mus = self._predict_coefs(self.models[0], C)\n",
    "        for model in self.models[1:]:\n",
    "            betas_i, mus_i = self._predict_coefs(model, C)\n",
    "            betas = np.concatenate((betas, betas_i), axis=0)\n",
    "            mus = np.concatenate((mus, mus_i), axis=0)\n",
    "        if all_bootstraps:\n",
    "            return betas, mus\n",
    "        return betas.mean(axis=0), mus.mean(axis=0)\n",
    "    \n",
    "    def predict_y(self, C, X, all_bootstraps=False):\n",
    "        betas, mus = self.predict_coefs(C, all_bootstraps=True)\n",
    "        y_hat = np.zeros((len(self.models), len(C), self.y_dim))\n",
    "        for i in range(len(self.models)):\n",
    "            for y_i in range(self.y_dim):\n",
    "                beta_y = betas[i, :, :, y_i]\n",
    "                mu_y = mus[i, :, y_i]\n",
    "                y_hat[i, :, y_i] = (beta_y * X).sum(axis=1) + mu_y\n",
    "        if all_bootstraps:\n",
    "            return y_hat\n",
    "        return y_hat.mean(axis=0)\n",
    "\n",
    "    def get_mse(self, C, X, Y, all_bootstraps=False):\n",
    "        \"\"\"\n",
    "        Returns the MSE of the model on a dataset\n",
    "        Returns a numpy array (1 or bootstraps, )\n",
    "        \"\"\"\n",
    "        dataset = MultivariateDataset(C, X, Y, batch_size=1, device=self.device)\n",
    "        mses = np.zeros(len(self.models))\n",
    "        for i, model in enumerate(self.models):\n",
    "            for c, t, x, y in dataset:\n",
    "                betas, mus = model(c, t)\n",
    "                mse = MSE(betas, mus, x, y).item()\n",
    "                mses[i] += 1 / len(dataset) * mse\n",
    "        if not all_bootstraps:\n",
    "            return mses.mean()\n",
    "        return mses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "cf3bea97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4 2 3\n"
     ]
    }
   ],
   "source": [
    "# from dataset import to_pairwise\n",
    "# C_pairwise, T_pairwise, X_pairwise, Y_pairwise = to_pairwise(C, X, Y)\n",
    "# uds = MultivariateDataset(C, X, Y)\n",
    "# count = 0\n",
    "# for dp in uds:\n",
    "#     print(dp)\n",
    "#     count += 1\n",
    "#     if count > 10:\n",
    "#         break\n",
    "# print(next(uds))\n",
    "# model = ContextualizedRegressionModule(C_pairwise.shape[-1], T_pairwise.shape[-1])\n",
    "# dataset = UnivariateDataset(C, X, Y, batch_size=1)\n",
    "# model = ContextualizedRegressionModule(dataset.c_dim, dataset.task_dim,) # beta_dim=dataset.x_dim)\n",
    "# opt = torch.optim.Adam(model.parameters(), lr=1e-2)\n",
    "# for epoch in range(1):\n",
    "#     for c, t, x, y in dataset:\n",
    "#         beta, mu = model(c, t)\n",
    "#         loss = MSE(beta, mu, x, y)  # this loop w unsqueeze is a hacky fix for the todo above\n",
    "#         opt.zero_grad()\n",
    "#         loss.backward()\n",
    "#         opt.step()\n",
    "#     loss_sum = 0\n",
    "#     for c, t, x, y in dataset:\n",
    "#         beta, mu = model(c, t)\n",
    "#         loss = MSE(beta, mu, x, y).detach().item()  # this loop w unsqueeze is a hacky fix for the todo above\n",
    "#         loss_sum += loss\n",
    "#     print(loss_sum)\n",
    "    \n",
    "# model = ContextualizedRegressionModule(C_pairwise.shape[-1], T_pairwise.shape[-1])\n",
    "# torch.stack((C[0], C[0], C[80])).shape\n",
    "print(C.shape[-1], X.shape[-1], Y.shape[-1])\n",
    "model = ContextualizedRegression(C.shape[-1], X.shape[-1], Y.shape[-1], bootstraps=2)\n",
    "# model.fit(C, X, Y, epochs=10, batch_size=1, lr=1e-2)\n",
    "# TODO: Write univariate predict fns, add multivariate contextualized regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "6093a126",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(100, 2, 3) (100, 3) (100, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/_6/468mk1cx3cb080z2r5gq0k_r0000gn/T/ipykernel_1526/205546479.py:325: 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, device=device)\n",
      "/var/folders/_6/468mk1cx3cb080z2r5gq0k_r0000gn/T/ipykernel_1526/205546479.py:326: 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, device=device)\n",
      "/var/folders/_6/468mk1cx3cb080z2r5gq0k_r0000gn/T/ipykernel_1526/205546479.py:327: 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, device=device)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.1312595792254387"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "betas, mus = model.predict_coefs(C)\n",
    "y_hat = model.predict_y(C.numpy(), X.numpy())\n",
    "print(betas.shape, mus.shape, y_hat.shape)\n",
    "model.get_mse(C, X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "8bde9c19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.10504741966724396\n",
      "0.08994530886411667\n",
      "0.07826512306928635\n",
      "0.06986884772777557\n",
      "0.0639653280377388\n",
      "0.05981753394007683\n",
      "0.056846413761377335\n",
      "0.05463786423206329\n",
      "0.05291229113936424\n",
      "0.05148674547672272\n",
      "0.05024436488747597\n",
      "0.04911275953054428\n",
      "0.04804883152246475\n",
      "0.047028183937072754\n",
      "0.046037837862968445\n",
      "0.04507139325141907\n",
      "0.04412606358528137\n",
      "0.04320095479488373\n",
      "0.04229608178138733\n",
      "0.041411858052015305\n",
      "0.04054883494973183\n",
      "0.03970754146575928\n",
      "0.03888847678899765\n",
      "0.03809204697608948\n",
      "0.037318550050258636\n",
      "0.03656821325421333\n",
      "0.035841163247823715\n",
      "0.03513743355870247\n",
      "0.034456998109817505\n",
      "0.03379973769187927\n",
      "0.03316546976566315\n",
      "0.03255394101142883\n",
      "0.031964849680662155\n",
      "0.03139783442020416\n",
      "0.030852477997541428\n",
      "0.03032832406461239\n",
      "0.029824890196323395\n",
      "0.02934165485203266\n",
      "0.028878064826130867\n",
      "0.028433555737137794\n",
      "0.028007542714476585\n",
      "0.027599439024925232\n",
      "0.02720862440764904\n",
      "0.02683452144265175\n",
      "0.026476513594388962\n",
      "0.02613401971757412\n",
      "0.025806443765759468\n",
      "0.02549322135746479\n",
      "0.025193793699145317\n",
      "0.02490762062370777\n"
     ]
    }
   ],
   "source": [
    "opt = torch.optim.Adam(model.parameters(), lr=1e-2)\n",
    "for epoch in range(5000):\n",
    "    beta, mu = model(C_pairwise, T_pairwise)\n",
    "    loss = MSE(beta.squeeze(), mu, X_pairwise, Y_pairwise)  # this loop w unsqueeze is a hacky fix for the todo above\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    if epoch % 100 == 0:\n",
    "        print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "id": "0fe80aa2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([400, 1]) torch.Size([400]) torch.Size([400]) torch.Size([400])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([400])"
      ]
     },
     "execution_count": 175,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "beta, mu = model(C_pairwise, T_pairwise)\n",
    "print(beta.shape, mu.shape, X_pairwise.shape, Y_pairwise.shape)\n",
    "Y_hat = beta.squeeze() * X_pairwise + mu\n",
    "Y_hat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "3b6eb0a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 0.2773, -0.1024,  0.1854, -0.2136,  0.2652, -0.1464,  0.2533, -0.3177,\n",
       "          0.2887,  0.0444], grad_fn=<SliceBackward0>),\n",
       " tensor([ 0.1800, -0.1988,  0.1800, -0.1988,  0.4678, -0.3332,  0.4678, -0.3332,\n",
       "          0.3164, -0.0254]))"
      ]
     },
     "execution_count": 176,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y_hat[:10], Y_pairwise[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2588f887",
   "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
}
