{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "israeli-mayor",
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.components.common"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef81d065-d1c5-4e82-9742-15b0917d295c",
   "metadata": {},
   "source": [
    "# Common components for models\n",
    "> Common functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "diagnostic-elimination",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import torch as t\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils import weight_norm\n",
    "from torch.autograd.function import Function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "powerful-luther",
   "metadata": {},
   "source": [
    "# Chomp1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "friendly-church",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Chomp1d(nn.Module):\n",
    "    \"\"\"\n",
    "    Receives x input of dim [N,C,T], and trims it so that only \n",
    "    'time available' information is used. Used for one dimensional \n",
    "    causal convolutions.\n",
    "    : param chomp_size: lenght of outsample values to skip.\n",
    "    \"\"\"\n",
    "    def __init__(self, chomp_size):\n",
    "        super(Chomp1d, self).__init__()\n",
    "        self.chomp_size = chomp_size\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x[:, :, :-self.chomp_size].contiguous()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "modified-salad",
   "metadata": {},
   "source": [
    "# CausalConv1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oriental-general",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "ACTIVATIONS = ['ReLU',\n",
    "               'Softplus',\n",
    "               'Tanh',\n",
    "               'SELU',\n",
    "               'LeakyReLU',\n",
    "               'PReLU',\n",
    "               'Sigmoid']\n",
    "\n",
    "class CausalConv1d(nn.Module):\n",
    "    \"\"\"\n",
    "    Receives x input of dim [N,C,T], computes a unidimensional \n",
    "    causal convolution.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    in_channels: int\n",
    "    out_channels: int\n",
    "    activation: str \n",
    "        https://discuss.pytorch.org/t/call-activation-function-from-string\n",
    "    padding: int\n",
    "    kernel_size: int\n",
    "    dilation: int\n",
    "    \n",
    "    Returns:\n",
    "    x: tesor\n",
    "        torch tensor of dim [N,C,T]\n",
    "        activation(conv1d(inputs, kernel) + bias)\n",
    "    \"\"\"\n",
    "    def __init__(self, in_channels, out_channels, kernel_size,\n",
    "                 padding, dilation, activation, stride:int=1, with_weight_norm:bool=False):\n",
    "        super(CausalConv1d, self).__init__()\n",
    "        assert activation in ACTIVATIONS, f'{activation} is not in {ACTIVATIONS}'\n",
    "        \n",
    "        self.conv       = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, \n",
    "                                    kernel_size=kernel_size, stride=stride, padding=padding,\n",
    "                                    dilation=dilation)\n",
    "        if with_weight_norm: self.conv = weight_norm(self.conv)\n",
    "        \n",
    "        self.chomp      = Chomp1d(padding)\n",
    "        self.activation = getattr(nn, activation)()\n",
    "        self.causalconv = nn.Sequential(self.conv, self.chomp, self.activation)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.causalconv(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ahead-dating",
   "metadata": {},
   "source": [
    "# TimeDistributed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "entitled-habitat",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class TimeDistributed2d(nn.Module):\n",
    "    \"\"\"\n",
    "    Receives x input of dim [N,C,T], reshapes it to [T,N,C]\n",
    "    Collapses input of dim [T,N,C] to [TxN,C] and applies module to C.\n",
    "    Finally reshapes it to [N,C_out,T].\n",
    "    Allows handling of variable sequence lengths and minibatch sizes.\n",
    "    : param module: Module to apply input to.\n",
    "    \"\"\"\n",
    "    def __init__(self, module):\n",
    "        super(TimeDistributed2d, self).__init__()\n",
    "        self.module = module\n",
    "\n",
    "    def forward(self, x):\n",
    "        N, C, T = x.size()\n",
    "        x = x.permute(2, 0, 1).contiguous()\n",
    "        x = x.view(T * N, -1)\n",
    "        x = self.module(x)\n",
    "        x = x.view(T, N, -1)\n",
    "        x = x.permute(1, 2, 0).contiguous()\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "french-scene",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class TimeDistributed3d(nn.Module):\n",
    "    \"\"\"\n",
    "    Receives x input of dim [N,L,C,T], reshapes it to [T,N,L,C]\n",
    "    Collapses input of dim [T,N,L,C] to [TxNxL,C] and applies module to C.\n",
    "    Finally reshapes it to [N,L,C_out,T].\n",
    "    Allows handling of variable sequence lengths and minibatch sizes.\n",
    "    : param module: Module to apply input to.\n",
    "    \"\"\"\n",
    "    def __init__(self, module):\n",
    "        super(TimeDistributed3d, self).__init__()\n",
    "        self.module = module\n",
    "\n",
    "    def forward(self, x):\n",
    "        N, L, C, T = x.size()\n",
    "        x = x.permute(3, 0, 1, 2).contiguous() #[N,L,C,T] --> [T,N,L,C]\n",
    "        x = x.view(T * N * L, -1)\n",
    "        x = self.module(x)\n",
    "        x = x.view(T, N, L, -1)\n",
    "        x = x.permute(1, 2, 3, 0).contiguous() #[T,N,L,C] --> [N,L,C,T]\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fossil-chess",
   "metadata": {},
   "source": [
    "# RepeatVector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "boxed-continent",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class RepeatVector(nn.Module):\n",
    "    \"\"\"\n",
    "    Receives x input of dim [N,C], and repeats the vector \n",
    "    to create tensor of shape [N, C, K]\n",
    "    : repeats: int, the number of repetitions for the vector.\n",
    "    \"\"\"\n",
    "    def __init__(self, repeats):\n",
    "        super(RepeatVector, self).__init__()\n",
    "        self.repeats = repeats\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.unsqueeze(-1).repeat(1, 1, self.repeats) # <------------ Mejorar?\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "premium-wednesday",
   "metadata": {},
   "source": [
    "# L1Regularizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "royal-mount",
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class L1Regularizer(nn.Module):\n",
    "    \"\"\"\n",
    "    Layer meant to apply elementwise L1 regularization to a dimension.\n",
    "    Receives x input of dim [N,C] and returns the input [N,C].\n",
    "    \"\"\"\n",
    "    def __init__(self, in_features, l1_lambda):\n",
    "        super(L1Regularizer, self).__init__()\n",
    "        self.l1_lambda = l1_lambda\n",
    "        self.weight = t.nn.Parameter(t.rand((in_features), dtype=t.float),\n",
    "                                     requires_grad=True)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # channelwise regularization, turns on or off channels\n",
    "        x = t.einsum('bp,p->bp', x, self.weight)\n",
    "        return x\n",
    "\n",
    "    def regularization(self):\n",
    "        return self.l1_lambda * t.norm(self.weight, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "stainless-tunnel",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X.shape (1000, 100)\n",
      "beta.shape (100,)\n",
      "Y.shape (1000, 1)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from sklearn import linear_model\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "X1  = np.random.normal(0, 1, (1000,1))\n",
    "X   = np.random.normal(0, 1, (1000, 99))\n",
    "X   = np.concatenate([X1, X], axis=1)\n",
    "eps = np.random.normal(0, 0.1, (1000))\n",
    "beta = np.array([1] + [0]*99)\n",
    "Y =  X @ beta.T + eps\n",
    "Y = np.expand_dims(Y, 1)\n",
    "print(\"X.shape\", X.shape)\n",
    "print(\"beta.shape\", beta.shape)\n",
    "print(\"Y.shape\", Y.shape)\n",
    "\n",
    "# model = linear_model.Lasso(alpha=0.1)\n",
    "# model.fit(X, Y)\n",
    "# print(\"model.coef_.shape\", model.coef_.shape)\n",
    "# model.coef_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ready-henry",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "import numpy as np\n",
    "import time\n",
    "from scipy.stats import hmean\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "czech-bulgarian",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _Model(nn.Module):  \n",
    "\n",
    "    def __init__(self, in_features, l1_lambda):\n",
    "        super(_Model, self).__init__()\n",
    "        self.l1 = L1Regularizer(in_features, l1_lambda)\n",
    "        self.linear_layer = nn.Linear(in_features=in_features, \n",
    "                                      out_features=1, \n",
    "                                      bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.l1(x.float())\n",
    "        y_hat = self.linear_layer(x)\n",
    "        return y_hat\n",
    "    \n",
    "class Data(Dataset):\n",
    "    \n",
    "    # Constructor\n",
    "    def __init__(self, Y, X):\n",
    "        self.X = X\n",
    "        self.Y = Y\n",
    "        self.len = Y.shape[0]\n",
    "\n",
    "    # Getter\n",
    "    def __getitem__(self, index):          \n",
    "        return self.X[index], self.Y[index]\n",
    "    \n",
    "    # Get Length\n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aquatic-mathematics",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_Model(\n",
      "  (l1): L1Regularizer()\n",
      "  (linear_layer): Linear(in_features=100, out_features=1, bias=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = _Model(in_features=X.shape[1], l1_lambda=0.07)\n",
    "dataloader = DataLoader(dataset=Data(X=X, Y=Y), batch_size=512)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "print(model)\n",
    "\n",
    "def train_model(model, epochs, print_progress=False):\n",
    "    start = time.time()\n",
    "    step = 0 \n",
    "    training_trajectory = {'epoch': [],\n",
    "                           'train_loss': []}\n",
    "    \n",
    "    criterion = t.nn.MSELoss()\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        for x, y in dataloader:\n",
    "            x, y = x.float(), y.float() # Type compatibility\n",
    "            \n",
    "            step += 1\n",
    "            y_hat = model(x)\n",
    "            \n",
    "            training_loss = criterion(y, y_hat) + model.l1.regularization()\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            training_loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "        if epoch % 100 == 0: \n",
    "            training_trajectory['epoch'].append(epoch)\n",
    "            train_loss = training_loss.detach().numpy()\n",
    "            training_trajectory['train_loss'].append(train_loss)\n",
    "\n",
    "            \n",
    "            display_str = f'epoch: {epoch} step: {step} time: {time.time()-start:03.3f} ** '\n",
    "            display_str += f'train_loss: {train_loss:.4f}'\n",
    "            print(display_str)\n",
    "            \n",
    "    return model, training_trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "least-might",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model, training_trajectory = train_model(model=model, epochs=2000)\n",
    "# plt.plot(training_trajectory['epoch'], training_trajectory['train_loss'])\n",
    "# plt.xlabel('Epochs')\n",
    "# plt.ylabel('MSE + L1 Loss')\n",
    "# plt.grid()\n",
    "# plt.show()\n",
    "# model.l1.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "creative-pregnancy",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
