Bayesian_LSTM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "from tqdm import trange, tqdm\n",
        "\n",
        "from io import BytesIO\n",
        "from urllib.request import urlopen\n",
        "from zipfile import ZipFile\n",
        "\n",
        "from pandas import read_csv\n",
        "from scipy import stats\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from tqdm import tqdm\n",
        "\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "from torch.utils.data.sampler import RandomSampler"
      ],
      "metadata": {
        "id": "29cnjOrEGR8N"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "6c8--hKmErOz"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Prepare the data for the DeepAR model\n",
        "def prep_data(data, covariates, data_start, train = True):\n",
        "    # Get the number of time steps in the data\n",
        "    time_len = data.shape[0]\n",
        "    # Calculate the size of the input window based on the stride size\n",
        "    input_size = window_size-stride_size\n",
        "    # Compute the number of windows for each time series\n",
        "    windows_per_series = np.full((num_series), (time_len-input_size-target_window_size) // stride_size)\n",
        "    # Adjust window count for training data\n",
        "    if train: windows_per_series -= (data_start+stride_size-1) // stride_size\n",
        "    # Calculate the total number of windows across all series\n",
        "    total_windows = np.sum(windows_per_series)\n",
        "    # Initialize the input matrix for the model\n",
        "    x_input = np.zeros((total_windows, window_size, 1 + num_covariates), dtype='float32')\n",
        "    # Initialize the labels matrix for the model\n",
        "    label = np.zeros((total_windows, target_window_size, 1 + num_covariates), dtype='float32')\n",
        "    # Initialize the matrix for additional input features\n",
        "    v_input = np.zeros((total_windows, 2), dtype='float32')\n",
        "    # Initialize counter for the number of processed windows\n",
        "    count = 0\n",
        "    # Iterate over each series\n",
        "    for series in trange(num_series):\n",
        "        # Iterate over each window in the current series\n",
        "        for i in range(windows_per_series[series]):\n",
        "            # Calculate the start of the current window\n",
        "            window_start = stride_size*i+data_start[series] if train else stride_size*i\n",
        "            # Calculate the end of the current window\n",
        "            window_end = window_start+window_size\n",
        "            # Calculate the end of the target window\n",
        "            target_window_end = window_end+target_window_size\n",
        "            # Set the current window data in the input matrix\n",
        "            x_input[count, :, 0] = data[window_start:window_end, series]\n",
        "            # Set the covariates for the current window in the input matrix\n",
        "            x_input[count, :, 1:1+num_covariates] = covariates[window_start:window_end, :]\n",
        "            # Set the target window data in the label matrix\n",
        "            label[count, :, 0] = data[window_end:target_window_end, series]\n",
        "            # Set the covariates for the target window in the label matrix\n",
        "            label[count,:, 1:1+num_covariates] = covariates[window_end:target_window_end, :]\n",
        "            # Calculate the sum of non-zero values in the input window\n",
        "            nonzero_sum = (x_input[count, 1:input_size, 0]!=0).sum()\n",
        "            # Normalize the input data if there are non-zero values\n",
        "            if nonzero_sum == 0:\n",
        "                v_input[count, 0] = 0\n",
        "            else:\n",
        "                v_input[count, 0] = np.true_divide(x_input[count, :input_size, 0].sum(),nonzero_sum)+1\n",
        "                x_input[count, :, 0] = x_input[count, :, 0]/v_input[count, 0]\n",
        "                label[count, :, 0] = label[count, :, 0]/v_input[count, 0]\n",
        "            # Increment the count after processing a window\n",
        "            count += 1\n",
        "    # Return the prepared data matrices\n",
        "    return x_input, v_input, label"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate covariates based on time attributes\n",
        "def gen_covariates(times, num_covariates):\n",
        "    # Initialize an array for storing covariates\n",
        "    covariates = np.zeros((times.shape[0], num_covariates))\n",
        "    # Iterate over each time point\n",
        "    for i, input_time in enumerate(times):\n",
        "        # Extract weekday, hour, and month as covariates\n",
        "        covariates[i, 0] = input_time.weekday()\n",
        "        covariates[i, 1] = input_time.hour\n",
        "        covariates[i, 2] = input_time.month\n",
        "    # Normalize the covariates using z-score normalization\n",
        "    for i in range(num_covariates):\n",
        "        covariates[:,i] = stats.zscore(covariates[:,i])\n",
        "    # Return the covariates array\n",
        "    return covariates[:, :num_covariates]\n"
      ],
      "metadata": {
        "id": "f5uj0OVYF6Fn"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "name = 'LD2011_2014.txt'\n",
        "save_name = 'elect'\n",
        "window_size = 192\n",
        "stride_size = 24\n",
        "target_window_size = 24\n",
        "num_covariates = 3\n",
        "train_start = '2011-01-01 00:00:00'\n",
        "train_end = '2013-12-31 23:00:00'\n",
        "validation_start = '2014-01-01 23:00:00'\n",
        "validation_end = '2014-08-31 23:00:00'\n",
        "test_start = '2014-08-25 00:00:00'\n",
        "test_end = '2014-09-07 23:00:00'\n",
        "\n",
        "save_path = os.path.join('data', save_name)\n",
        "if not os.path.exists(save_path):\n",
        "    os.makedirs(save_path)\n",
        "csv_path = os.path.join(save_path, name)\n",
        "if not os.path.exists(csv_path):\n",
        "    zipurl = 'https://archive.ics.uci.edu/ml/machine-learning-databases/00321/LD2011_2014.txt.zip'\n",
        "    with urlopen(zipurl) as zipresp:\n",
        "        with ZipFile(BytesIO(zipresp.read())) as zfile:\n",
        "            zfile.extractall(save_path)\n",
        "\n",
        "data_frame = pd.read_csv(csv_path, sep=\";\", index_col=0, parse_dates=True, decimal=',')\n",
        "data_frame = data_frame.resample('1H',label = 'left',closed = 'right').sum()[train_start:test_end]\n",
        "data_frame.fillna(0, inplace=True)\n",
        "\n",
        "covariates = gen_covariates(data_frame[train_start:test_end].index, num_covariates)\n",
        "\n",
        "train_data = data_frame[train_start:train_end].values\n",
        "validation_data = data_frame[validation_start:validation_end].values\n",
        "test_data = data_frame[test_start:test_end].values\n",
        "\n",
        "data_start = (train_data!=0).argmax(axis=0) #find first nonzero value in each time series\n",
        "total_time = data_frame.shape[0] #32304\n",
        "num_series = data_frame.shape[1] #370\n",
        "\n",
        "X_train, v_train, y_train = prep_data(train_data, covariates, data_start)\n",
        "X_validation, v_validation, y_validation = prep_data(validation_data, covariates, data_start, train=False)\n",
        "X_test, v_test, y_test = prep_data(test_data, covariates, data_start, train=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jfgeSG9qF-gM",
        "outputId": "ce5c7237-ccd3-4bd4-a9e5-2690a2128647"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 370/370 [00:06<00:00, 54.16it/s]\n",
            "100%|██████████| 370/370 [00:01<00:00, 199.84it/s]\n",
            "100%|██████████| 370/370 [00:00<00:00, 4296.69it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Class VariationalDropout\n"
      ],
      "metadata": {
        "id": "HxMseMe4HsTH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Variational Dropout implementation based on the paper: https://arxiv.org/abs/1512.05287\n",
        "class VariationalDropout(nn.Module):\n",
        "    # Initialize VariationalDropout with a specified dropout rate\n",
        "    def __init__(self, dropout):\n",
        "        super().__init__()\n",
        "        self.dropout = dropout\n",
        "\n",
        "    # Forward pass for applying dropout\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        # Return input as is if not in training mode\n",
        "        if not self.training:\n",
        "            return x\n",
        "        # Determine the maximum batch size from input\n",
        "        max_batch_size = x.size(1)\n",
        "        # Generate a dropout mask with the same mask applied across the entire sequence\n",
        "        m = x.new_empty(1, max_batch_size, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)\n",
        "        # Apply the mask and scale the activations accordingly\n",
        "        x = x.masked_fill(m == 0, 0) / (1 - self.dropout)\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "I3kebHvUGOkX"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Class LSTM\n"
      ],
      "metadata": {
        "id": "S7qMaA9gH1iC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Custom LSTM class with variational dropout for input, weights, and output\n",
        "class LSTM(nn.LSTM):\n",
        "    # Initialize LSTM with specific dropout rates for input, weight, and output\n",
        "    def __init__(self, *args, dropouti=0., dropoutw=0., dropouto=0., unit_forget_bias=True, **kwargs):\n",
        "        super().__init__(*args, **kwargs)\n",
        "        self.unit_forget_bias = unit_forget_bias\n",
        "        self.dropoutw = dropoutw\n",
        "        self.input_drop = VariationalDropout(dropouti)\n",
        "        self.output_drop = VariationalDropout(dropouto)\n",
        "\n",
        "    # Function to apply dropout to hidden-to-hidden weights\n",
        "    def _drop_weights(self):\n",
        "        for name, param in self.named_parameters():\n",
        "            if \"weight_hh\" in name:\n",
        "                getattr(self, name).data = torch.nn.functional.dropout(param.data, p=self.dropoutw, training=self.training).contiguous()\n",
        "\n",
        "    # Forward pass with custom dropout for LSTM\n",
        "    def forward(self, input):\n",
        "        # Apply dropout to weights\n",
        "        self._drop_weights()\n",
        "        # Apply variational dropout to input\n",
        "        input = self.input_drop(input)\n",
        "        # Standard LSTM forward pass\n",
        "        seq, state = super().forward(input)\n",
        "        # Apply variational dropout to output\n",
        "        return self.output_drop(seq), state\n"
      ],
      "metadata": {
        "id": "X2pBo4ZtHxFC"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        " ### Class VDEncoder\n"
      ],
      "metadata": {
        "id": "xhL4cjhvH7yU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Variational Dropout Encoder module\n",
        "class VDEncoder(nn.Module):\n",
        "    # Initialize VDEncoder with specific input and output features and dropout probability\n",
        "    def __init__(self, in_features, out_features, p):\n",
        "        super(VDEncoder, self).__init__()\n",
        "        # Define LSTM layers with variational dropout for the encoder\n",
        "        self.model = nn.ModuleDict({\n",
        "            'lstm1': LSTM(in_features, 32, dropouto=p),\n",
        "            'lstm2': LSTM(32, 8, dropouto=p),\n",
        "            'lstm3': LSTM(8, out_features, dropouto=p)\n",
        "        })\n",
        "\n",
        "    # Forward pass through the VDEncoder\n",
        "    def forward(self, x):\n",
        "        # Sequentially pass input through each LSTM layer\n",
        "        out, _ = self.model['lstm1'](x)\n",
        "        out, _ = self.model['lstm2'](out)\n",
        "        out, _ = self.model['lstm3'](out)\n",
        "        return out\n"
      ],
      "metadata": {
        "id": "lu_DaNeeIARc"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Class VDDecoder\n"
      ],
      "metadata": {
        "id": "MBEd7KLbIDyJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Variational Dropout Decoder module\n",
        "class VDDecoder(nn.Module):\n",
        "    # Initialize VDDecoder with dropout probability\n",
        "    def __init__(self, p):\n",
        "        super(VDDecoder, self).__init__()\n",
        "        # Define LSTM layers with variational dropout for the decoder\n",
        "        self.model = nn.ModuleDict({\n",
        "            'lstm1': LSTM(1, 2, dropouto=p),\n",
        "            'lstm2': LSTM(2, 2, dropouto=p),\n",
        "            'lstm3': LSTM(2, 1, dropouto=p)\n",
        "        })\n",
        "\n",
        "    # Forward pass through the VDDecoder\n",
        "    def forward(self, x):\n",
        "        # Sequentially pass input through each LSTM layer\n",
        "        out, _ = self.model['lstm1'](x)\n",
        "        out, _ = self.model['lstm2'](out)\n",
        "        out, _ = self.model['lstm3'](out)\n",
        "        return out\n"
      ],
      "metadata": {
        "id": "eflg8b8XIHcJ"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Class VDEncoderDecoder\n"
      ],
      "metadata": {
        "id": "RxxJpAMbIJ8_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the Variational Dropout Encoder-Decoder model class.\n",
        "class VDEncoderDecoder(nn.Module):\n",
        "    # Initialize the class with specific input features, output steps, and dropout probability.\n",
        "    def __init__(self, in_features, output_steps, p):\n",
        "        super(VDEncoderDecoder, self).__init__()\n",
        "        self.enc_in_features = in_features  # Store the number of input features for the encoder.\n",
        "        self.output_steps = output_steps  # Store the number of output steps (f in the paper).\n",
        "        self.enc_out_features = 1  # Define the number of output features for the encoder.\n",
        "        self.traffic_col = 4  # Set the index of the traffic column (specific to use case).\n",
        "\n",
        "        # Create a module dictionary to store encoder, decoder, and fully connected layers.\n",
        "        self.model = nn.ModuleDict({\n",
        "            'encoder': VDEncoder(self.enc_in_features, self.enc_out_features, p),  # Initialize the encoder.\n",
        "            'decoder': VDDecoder(p),  # Initialize the decoder.\n",
        "            'fc1': nn.Linear(window_size, 32),  # First fully connected layer to transform encoded output.\n",
        "            'fc2': nn.Linear(32, self.output_steps)  # Second fully connected layer to generate final output.\n",
        "        })\n",
        "\n",
        "    # Define the forward pass of the model.\n",
        "    def forward(self, x):\n",
        "        # Pass the input through the encoder.\n",
        "        out = self.model['encoder'](x)\n",
        "        # Pass the encoder's output through the decoder.\n",
        "        out = self.model['decoder'](out)\n",
        "        # Flatten and pass the decoder's output through the first fully connected layer.\n",
        "        out = self.model['fc1'](out.squeeze().view(64,-1))\n",
        "        # Pass the output of the first fully connected layer through the second to get final predictions.\n",
        "        out = self.model['fc2'](out)\n",
        "\n",
        "        return out  # Return the final output.\n"
      ],
      "metadata": {
        "id": "0qmypVuTIMMw"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Class TrainDataset\n"
      ],
      "metadata": {
        "id": "wOYvWL8WJZhZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a dataset class for training.\n",
        "class TrainDataset(Dataset):\n",
        "    # Initialize the dataset with data and labels.\n",
        "    def __init__(self, data, label):\n",
        "        self.data = data  # Store the input data.\n",
        "        self.label = label  # Store the labels corresponding to the data.\n",
        "        self.train_len = self.data.shape[0]  # Calculate the length of the dataset.\n",
        "\n",
        "    # Return the length of the dataset.\n",
        "    def __len__(self):\n",
        "        return self.train_len\n",
        "\n",
        "    # Return the data and label for a given index.\n",
        "    def __getitem__(self, index):\n",
        "        # Return the main feature, additional covariates, and corresponding label.\n",
        "        return (self.data[index,:,0],  # Extract the main feature from the data.\n",
        "                self.data[index,:,1:1+num_covariates],  # Extract covariates from the data.\n",
        "                self.label[index,:,0],  # Extract the main label.\n",
        "                self.label[index,:,1:1+num_covariates])  # Extract covariates from the label.\n"
      ],
      "metadata": {
        "id": "xYeS9XnKJdr1"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Class ValidationAndTestDataset"
      ],
      "metadata": {
        "id": "hv41aihEJhh0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a dataset class for validation and testing.\n",
        "class ValidationAndTestDataset(Dataset):\n",
        "    # Initialize the dataset with data, scaling factors (v), and labels.\n",
        "    def __init__(self, data, v, label):\n",
        "        self.data = data  # Store the input data.\n",
        "        self.v = v  # Store the scaling factors for the data.\n",
        "        self.label = label  # Store the labels corresponding to the data.\n",
        "        self.test_len = self.data.shape[0]  # Calculate the length of the dataset.\n",
        "\n",
        "    # Return the length of the dataset.\n",
        "    def __len__(self):\n",
        "        return self.test_len\n",
        "\n",
        "    # Return the data, scaling factor, and label for a given index.\n",
        "    def __getitem__(self, index):\n",
        "        # Return the main feature, additional covariates, scaling factor, and corresponding label.\n",
        "        return (self.data[index,:,0],  # Extract the main feature from the data.\n",
        "                self.data[index,:,1:1+num_covariates],  # Extract covariates from the data.\n",
        "                self.v[index],  # Extract the scaling factor for the given index.\n",
        "                self.label[index,:,0],  # Extract the main label.\n",
        "                self.label[index,:,1:1+num_covariates])  # Extract covariates from the label.\n"
      ],
      "metadata": {
        "id": "obqrKMPzJlXc"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_batch_size = 64\n",
        "\n",
        "train_set = TrainDataset(data=X_train, label=y_train)\n",
        "validation_set = ValidationAndTestDataset(data=X_validation, v=v_validation, label=y_validation)\n",
        "test_set = ValidationAndTestDataset(data=X_test, v=v_test, label=y_test)\n",
        "\n",
        "train_loader = DataLoader(train_set, batch_size=train_batch_size, drop_last=True)\n",
        "validation_loader = DataLoader(validation_set, batch_size=train_batch_size, sampler=RandomSampler(test_set))\n",
        "test_loader = DataLoader(test_set, batch_size=train_batch_size, sampler=RandomSampler(test_set), drop_last=True)"
      ],
      "metadata": {
        "id": "3KnSRenHJu9r"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a training function for the encoder-decoder model.\n",
        "def train_encdec(model, device=torch.device('cuda'), num_epochs=2, learning_rate=1e-3):\n",
        "    # Get the length of the training dataset.\n",
        "    train_len = len(train_loader)\n",
        "    # Initialize the optimizer with the model's parameters and the learning rate.\n",
        "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "    # Initialize a list to store loss values for each epoch.\n",
        "    loss_summary = []\n",
        "    # Set the loss function to mean squared error.\n",
        "    loss_fn = F.mse_loss\n",
        "\n",
        "    # Loop through the specified number of epochs.\n",
        "    for epoch in range(num_epochs):\n",
        "        # Set the model to training mode.\n",
        "        model.train()\n",
        "        # Initialize variables to track the sum of loss and total number of samples.\n",
        "        epoch_loss_sum = 0.0\n",
        "        total_sample = 0\n",
        "\n",
        "        # Initialize a progress bar for the training loader.\n",
        "        pbar = tqdm(train_loader)\n",
        "        # Loop through each batch in the training dataset.\n",
        "        for (train_batch, current_covs_batch, labels_batch, next_covs_batch) in pbar:\n",
        "            # Determine batch size, sequence length, and horizon size.\n",
        "            batch_size, seq_len, horizon_size = train_batch.shape[0], train_batch.shape[1], labels_batch.shape[0]\n",
        "            # Update the total number of samples processed.\n",
        "            total_sample += batch_size * seq_len * horizon_size\n",
        "            # Reset the gradients of the model's parameters.\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Adjust the dimensions of the training batch and move it to the specified device.\n",
        "            train_batch = train_batch.unsqueeze(2).permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "            # Pass the training batch through the model to get the output.\n",
        "            out = model(train_batch)\n",
        "            # Calculate the loss between the model's output and the actual labels.\n",
        "            loss = loss_fn(out.float(), labels_batch.squeeze().to(device).float())\n",
        "\n",
        "            # Update the progress bar with the current loss value.\n",
        "            pbar.set_description(f\"Loss:{loss.item()}\")\n",
        "            # Perform backpropagation to calculate gradients.\n",
        "            loss.backward()\n",
        "            # Update the model's parameters based on the gradients.\n",
        "            optimizer.step()\n",
        "\n",
        "        # Store the loss of the current epoch.\n",
        "        loss_summary.append(loss.cpu().detach())\n",
        "    # Return the summary of losses after all epochs.\n",
        "    return loss_summary\n"
      ],
      "metadata": {
        "id": "IK-wHGKQJwuo"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encdec_model = VDEncoderDecoder(in_features=1, output_steps=target_window_size, p=0.25).cuda()\n",
        "train_encdec(encdec_model, num_epochs = 2, device=torch.device('cuda'))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g9PEQ3hyKW4Q",
        "outputId": "15394f56-267e-4f21-c6b4-40843ba7e4f9"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Loss:0.07597240805625916: 100%|██████████| 5004/5004 [01:23<00:00, 60.12it/s]\n",
            "Loss:0.08595262467861176: 100%|██████████| 5004/5004 [01:18<00:00, 63.81it/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[tensor(0.0760), tensor(0.0860)]"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### class PredictionNetwork"
      ],
      "metadata": {
        "id": "MX-l_lcDLCiy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a prediction network class that inherits from nn.Module.\n",
        "class PredictionNetwork(nn.Module):\n",
        "    # Initialize the network with an encoder-decoder model and a dropout probability.\n",
        "    def __init__(self, encoder_decoder, p=0.25):\n",
        "        # Initialize the base class.\n",
        "        super(PredictionNetwork, self).__init__()\n",
        "        # Retrieve the encoder part of the passed encoder-decoder model and set it to evaluation mode.\n",
        "        self.encoder = encoder_decoder.model['encoder'].eval()\n",
        "        # Define a sequential model with linear layers, dropout, and ReLU activation functions.\n",
        "        self.model = nn.Sequential(\n",
        "            # First linear layer takes input size equal to window size times the number of features (1 + num_covariates).\n",
        "            nn.Linear((1 + num_covariates)*window_size, 128),\n",
        "            # Apply dropout with the specified probability.\n",
        "            nn.Dropout(p),\n",
        "            # Apply ReLU activation function.\n",
        "            nn.ReLU(),\n",
        "            # Apply dropout again.\n",
        "            nn.Dropout(p),\n",
        "            # Second linear layer to reduce dimension from 128 to 64.\n",
        "            nn.Linear(128, 64),\n",
        "            # Apply ReLU activation function.\n",
        "            nn.ReLU(),\n",
        "            # Apply dropout once more.\n",
        "            nn.Dropout(p),\n",
        "            # Final linear layer to produce output of size equal to target window size.\n",
        "            nn.Linear(64, target_window_size)\n",
        "        )\n",
        "\n",
        "    # Define the forward pass of the model.\n",
        "    def forward(self, x_input, cov_input):\n",
        "        # Pass the input through the encoder to get extracted features.\n",
        "        extracted = self.encoder(x_input)\n",
        "        # Concatenate the extracted features with the covariate input.\n",
        "        x_concat = torch.cat([extracted, cov_input], dim=-1)\n",
        "        # Flatten the concatenated features and pass them through the sequential model.\n",
        "        # Squeeze the output to remove dimensions of size 1.\n",
        "        out = self.model(x_concat.view(train_batch_size, -1)).squeeze()\n",
        "        # Return the final output.\n",
        "        return out\n"
      ],
      "metadata": {
        "id": "HBpetBjqK9o5"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a function to train the prediction network.\n",
        "def train_prediction_network(model, device=torch.device('cuda'), num_epochs = 2, learning_rate = 1e-3):\n",
        "    # Calculate the total number of batches in the train loader.\n",
        "    train_len = len(train_loader)\n",
        "    # Set up the optimizer for the model.\n",
        "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "    # Initialize a list to keep track of the loss after each epoch.\n",
        "    loss_summary = []\n",
        "    # Define the loss function as mean squared error (MSE).\n",
        "    loss_fn = F.mse_loss\n",
        "\n",
        "    # Iterate over each epoch.\n",
        "    for epoch in range(num_epochs):\n",
        "        # Set the model to training mode.\n",
        "        model.train()\n",
        "        # Initialize variables to track loss and sample count.\n",
        "        epoch_loss_sum = 0.0\n",
        "        total_sample = 0\n",
        "\n",
        "        # Create a progress bar for iterating through the data loader.\n",
        "        pbar = tqdm(train_loader)\n",
        "        # Iterate over each batch in the data loader.\n",
        "        for (train_batch, current_covs_batch, labels_batch, next_covs_batch) in pbar:\n",
        "            # Determine the size of the batch, sequence length, and horizon size.\n",
        "            batch_size, seq_len, horizon_size = train_batch.shape[0], train_batch.shape[1], labels_batch.shape[0]\n",
        "            # Calculate the total number of samples processed.\n",
        "            total_sample += batch_size * seq_len * horizon_size\n",
        "            # Reset gradients in the optimizer.\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Prepare the training batch and covariates batch for the model, adjusting dimensions and moving to the device.\n",
        "            train_batch = train_batch.unsqueeze(2).permute(1,0,2).to(torch.float32).to(device)\n",
        "            current_covs_batch = current_covs_batch.permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "            # Pass the data through the model.\n",
        "            out = model(train_batch, current_covs_batch)\n",
        "            # Calculate the loss between model output and labels.\n",
        "            loss = loss_fn(out.float(), labels_batch.squeeze().to(device).float())\n",
        "\n",
        "            # Update the progress bar with the current loss.\n",
        "            pbar.set_description(f\"Loss:{loss.item()}\")\n",
        "            # Backpropagate the loss.\n",
        "            loss.backward()\n",
        "            # Update model parameters.\n",
        "            optimizer.step()\n",
        "\n",
        "        # Store the loss of the last batch in the summary.\n",
        "        loss_summary.append(loss.cpu().detach())\n",
        "    # Return the loss summary and the optimizer for further use.\n",
        "    return loss_summary, optimizer\n",
        "\n",
        "# Define a function to evaluate the prediction network.\n",
        "def evaluate_prediction_network(model, optimizer, device=torch.device('cuda')):\n",
        "    # Set the criterion as mean squared error (MSE) loss for evaluation.\n",
        "    criterion = nn.MSELoss()\n",
        "    # Initialize a list to store root mean squared error (RMSE) results for each batch.\n",
        "    rmse_results = []\n",
        "\n",
        "    # Disable gradient calculations for evaluation.\n",
        "    with torch.no_grad():\n",
        "        # Set the model to evaluation mode.\n",
        "        model.eval()\n",
        "        # Initialize an array to store loss per epoch.\n",
        "        loss_epoch = np.zeros(len(train_loader))\n",
        "\n",
        "        # Create a progress bar for iterating through the test data loader.\n",
        "        pbar = tqdm(test_loader)\n",
        "        # Iterate over each batch in the test loader.\n",
        "        for (ts_data_batch, current_covs_batch, v_batch, labels_batch, next_covs_batch) in pbar:\n",
        "            # Reset gradients in the optimizer.\n",
        "            optimizer.zero_grad()\n",
        "            # Prepare the test data batch and covariates batch for the model, adjusting dimensions and moving to the device.\n",
        "            ts_data_batch = ts_data_batch.unsqueeze(2).permute(1,0,2).to(torch.float32).to(device)\n",
        "            current_covs_batch = current_covs_batch.permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "            # Pass the data through the model.\n",
        "            out = model(ts_data_batch, current_covs_batch)\n",
        "            # Calculate RMSE for the batch and add it to the results.\n",
        "            rmse_results.append(torch.sqrt(criterion(out.detach().cpu(), labels_batch)).item())\n",
        "\n",
        "    # Calculate the average RMSE across all test batches.\n",
        "    test_rmse = np.mean(rmse_results)\n",
        "    # Return the average test RMSE.\n",
        "    return test_rmse\n"
      ],
      "metadata": {
        "id": "uWbATmisLLLH"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a function to train the prediction network.\n",
        "def train_prediction_network(model, device=torch.device('cuda'), num_epochs = 2, learning_rate = 1e-3):\n",
        "    # Calculate the total number of batches in the train loader.\n",
        "    train_len = len(train_loader)\n",
        "    # Set up the optimizer for the model.\n",
        "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "    # Initialize a list to keep track of the loss after each epoch.\n",
        "    loss_summary = []\n",
        "    # Define the loss function as mean squared error (MSE).\n",
        "    loss_fn = F.mse_loss\n",
        "\n",
        "    # Iterate over each epoch.\n",
        "    for epoch in range(num_epochs):\n",
        "        # Set the model to training mode.\n",
        "        model.train()\n",
        "        # Initialize variables to track loss and sample count.\n",
        "        epoch_loss_sum = 0.0\n",
        "        total_sample = 0\n",
        "\n",
        "        # Create a progress bar for iterating through the data loader.\n",
        "        pbar = tqdm(train_loader)\n",
        "        # Iterate over each batch in the data loader.\n",
        "        for (train_batch, current_covs_batch, labels_batch, next_covs_batch) in pbar:\n",
        "            # Determine the size of the batch, sequence length, and horizon size.\n",
        "            batch_size, seq_len, horizon_size = train_batch.shape[0], train_batch.shape[1], labels_batch.shape[0]\n",
        "            # Calculate the total number of samples processed.\n",
        "            total_sample += batch_size * seq_len * horizon_size\n",
        "            # Reset gradients in the optimizer.\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Prepare the training batch and covariates batch for the model, adjusting dimensions and moving to the device.\n",
        "            train_batch = train_batch.unsqueeze(2).permute(1,0,2).to(torch.float32).to(device)\n",
        "            current_covs_batch = current_covs_batch.permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "            # Pass the data through the model.\n",
        "            out = model(train_batch, current_covs_batch)\n",
        "            # Calculate the loss between model output and labels.\n",
        "            loss = loss_fn(out.float(), labels_batch.squeeze().to(device).float())\n",
        "\n",
        "            # Update the progress bar with the current loss.\n",
        "            pbar.set_description(f\"Loss:{loss.item()}\")\n",
        "            # Backpropagate the loss.\n",
        "            loss.backward()\n",
        "            # Update model parameters.\n",
        "            optimizer.step()\n",
        "\n",
        "        # Store the loss of the last batch in the summary.\n",
        "        loss_summary.append(loss.cpu().detach())\n",
        "    # Return the loss summary and the optimizer for further use.\n",
        "    return loss_summary, optimizer\n",
        "\n",
        "# Define a function to evaluate the prediction network.\n",
        "def evaluate_prediction_network(model, optimizer, device=torch.device('cuda')):\n",
        "    # Set the criterion as mean squared error (MSE) loss for evaluation.\n",
        "    criterion = nn.MSELoss()\n",
        "    # Initialize a list to store root mean squared error (RMSE) results for each batch.\n",
        "    rmse_results = []\n",
        "\n",
        "    # Disable gradient calculations for evaluation.\n",
        "    with torch.no_grad():\n",
        "        # Set the model to evaluation mode.\n",
        "        model.eval()\n",
        "        # Initialize an array to store loss per epoch.\n",
        "        loss_epoch = np.zeros(len(train_loader))\n",
        "\n",
        "        # Create a progress bar for iterating through the test data loader.\n",
        "        pbar = tqdm(test_loader)\n",
        "        # Iterate over each batch in the test loader.\n",
        "        for (ts_data_batch, current_covs_batch, v_batch, labels_batch, next_covs_batch) in pbar:\n",
        "            # Reset gradients in the optimizer.\n",
        "            optimizer.zero_grad()\n",
        "            # Prepare the test data batch and covariates batch for the model, adjusting dimensions and moving to the device.\n",
        "            ts_data_batch = ts_data_batch.unsqueeze(2).permute(1,0,2).to(torch.float32).to(device)\n",
        "            current_covs_batch = current_covs_batch.permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "            # Pass the data through the model.\n",
        "            out = model(ts_data_batch, current_covs_batch)\n",
        "            # Calculate RMSE for the batch and add it to the results.\n",
        "            rmse_results.append(torch.sqrt(criterion(out.detach().cpu(), labels_batch)).item())\n",
        "\n",
        "    # Calculate the average RMSE across all test batches.\n",
        "    test_rmse = np.mean(rmse_results)\n",
        "    # Return the average test RMSE.\n",
        "    return test_rmse\n"
      ],
      "metadata": {
        "id": "AKyVisNLL7ul"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def accuracy_RMSE(mu: torch.Tensor, labels: torch.Tensor, relative = False):\n",
        "    zero_index = (labels != 0)\n",
        "    diff = torch.sum(torch.mul((mu[zero_index] - labels[zero_index]), (mu[zero_index] - labels[zero_index]))).item()\n",
        "    if relative is False:\n",
        "        return [diff, torch.sum(zero_index).item(), torch.sum(zero_index).item()]\n",
        "    else:\n",
        "        summation = torch.sum(torch.abs(labels[zero_index])).item()\n",
        "        if summation == 0:\n",
        "            logger.error('summation denominator error! ')\n",
        "        return [diff, summation, torch.sum(zero_index).item()]\n",
        "\n",
        "def update_metrics(raw_metrics, sample_mu, labels, relative=False):\n",
        "    # TODO: use samples to calcualte rou50, rou90 metrics\n",
        "    raw_metrics['RMSE'] = raw_metrics['RMSE'] + accuracy_RMSE(sample_mu, labels, relative=relative)\n",
        "    return raw_metrics\n",
        "\n",
        "def final_metrics(raw_metrics):\n",
        "    summary_metric = {}\n",
        "    summary_metric['RMSE'] = np.sqrt(raw_metrics['RMSE'][0] / raw_metrics['RMSE'][2]) / (\n",
        "                raw_metrics['RMSE'][1] / raw_metrics['RMSE'][2])\n",
        "    return summary_metric\n",
        "\n",
        "def dropout_on(m):\n",
        "    if type(m) in [torch.nn.Dropout, LSTM]:\n",
        "        m.train()\n",
        "\n",
        "def dropout_off(m):\n",
        "    if type(m) in [torch.nn.Dropout, LSTM]:\n",
        "        m.eval()\n",
        "\n",
        "def mc_dropout(model, B, device):\n",
        "    model = model.apply(dropout_on)\n",
        "\n",
        "    pbar = range(B)\n",
        "    pbar = tqdm(pbar)\n",
        "\n",
        "    y_hats = []\n",
        "    for b in pbar:\n",
        "        for (x, cov, v, y, ncov) in test_loader:\n",
        "            x,cov,y = x.to(device), cov.to(device), y.to(device)\n",
        "            break\n",
        "        x = x.unsqueeze(2).permute(1, 0, 2).to(torch.float32).to(device)\n",
        "        cov = cov.permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "        y_hat_b = model(x, cov).float()\n",
        "        y_hats.append(y_hat_b.cpu().detach().numpy())\n",
        "\n",
        "    ymc_hats = np.mean(y_hats, axis=0)\n",
        "    eta_1s   = np.mean((ymc_hats[:,0] - np.stack(y_hats)[:,:,0])**2, axis=0)\n",
        "    return ymc_hats, eta_1s\n",
        "\n",
        "\n",
        "def inference(model, B=100, device=torch.device('cuda')):\n",
        "    # mc dropout\n",
        "    ymc_hats, eta_1s = mc_dropout(model, B, device)\n",
        "\n",
        "    # inherent noise\n",
        "    model.apply(dropout_off)\n",
        "    for (x, cov, v, y, ncov) in validation_loader:\n",
        "        x,cov,y = x.to(device), cov.to(device), y.to(device)\n",
        "        break\n",
        "    x = x.unsqueeze(2).permute(1, 0, 2).to(torch.float32).to(device)\n",
        "    cov = cov.permute(1,0,2).to(torch.float32).to(device)\n",
        "    y_hat_b = model(x, cov)\n",
        "\n",
        "    eta_2sq = np.mean(y_hat_b.cpu().detach().numpy()[:,0])\n",
        "    # total noise\n",
        "    etas = np.sqrt(eta_1s + eta_2sq)\n",
        "\n",
        "    for (x, cov, v, y, ncov) in test_loader:\n",
        "        break\n",
        "    normalized_label = y.T/v[:,0]\n",
        "    diff = torch.nansum(torch.mul((normalized_label.T - torch.tensor(ymc_hats)), (normalized_label.T - torch.tensor(ymc_hats)))).item()\n",
        "    test_rmse = np.sqrt(diff/len(test_set))\n",
        "\n",
        "    return test_rmse"
      ],
      "metadata": {
        "id": "4GsLkZC_McaA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a function to calculate the root mean squared error (RMSE).\n",
        "def accuracy_RMSE(mu: torch.Tensor, labels: torch.Tensor, relative = False):\n",
        "    # Identify non-zero elements in labels.\n",
        "    zero_index = (labels != 0)\n",
        "    # Calculate the squared difference between predictions and labels for non-zero elements.\n",
        "    diff = torch.sum(torch.mul((mu[zero_index] - labels[zero_index]), (mu[zero_index] - labels[zero_index]))).item()\n",
        "    # Return different metrics based on whether relative error is required.\n",
        "    if relative is False:\n",
        "        return [diff, torch.sum(zero_index).item(), torch.sum(zero_index).item()]\n",
        "    else:\n",
        "        # Calculate the sum of absolute values of labels for non-zero elements.\n",
        "        summation = torch.sum(torch.abs(labels[zero_index])).item()\n",
        "        # Log an error if the summation is zero.\n",
        "        if summation == 0:\n",
        "            logger.error('summation denominator error! ')\n",
        "        return [diff, summation, torch.sum(zero_index).item()]\n",
        "\n",
        "# Define a function to update the raw metrics with new data.\n",
        "def update_metrics(raw_metrics, sample_mu, labels, relative=False):\n",
        "    # Update the RMSE metric in the raw metrics dictionary.\n",
        "    raw_metrics['RMSE'] = raw_metrics['RMSE'] + accuracy_RMSE(sample_mu, labels, relative=relative)\n",
        "    return raw_metrics\n",
        "\n",
        "# Define a function to calculate final metrics from raw metrics.\n",
        "def final_metrics(raw_metrics):\n",
        "    summary_metric = {}\n",
        "    # Calculate the final RMSE value.\n",
        "    summary_metric['RMSE'] = np.sqrt(raw_metrics['RMSE'][0] / raw_metrics['RMSE'][2]) / (\n",
        "                raw_metrics['RMSE'][1] / raw_metrics['RMSE'][2])\n",
        "    return summary_metric\n",
        "\n",
        "# Define a function to enable dropout in a model.\n",
        "def dropout_on(m):\n",
        "    if type(m) in [torch.nn.Dropout, LSTM]:\n",
        "        m.train()\n",
        "\n",
        "# Define a function to disable dropout in a model.\n",
        "def dropout_off(m):\n",
        "    if type(m) in [torch.nn.Dropout, LSTM]:\n",
        "        m.eval()\n",
        "\n",
        "# Define a function to perform Monte Carlo dropout.\n",
        "def mc_dropout(model, B, device):\n",
        "    # Enable dropout in the model.\n",
        "    model = model.apply(dropout_on)\n",
        "\n",
        "    # Initialize a progress bar for Monte Carlo iterations.\n",
        "    pbar = range(B)\n",
        "    pbar = tqdm(pbar)\n",
        "\n",
        "    y_hats = []\n",
        "    # Iterate over each Monte Carlo sample.\n",
        "    for b in pbar:\n",
        "        # Retrieve a batch of data from the test loader.\n",
        "        for (x, cov, v, y, ncov) in test_loader:\n",
        "            x,cov,y = x.to(device), cov.to(device), y.to(device)\n",
        "            break\n",
        "        # Prepare the data for the model.\n",
        "        x = x.unsqueeze(2).permute(1, 0, 2).to(torch.float32).to(device)\n",
        "        cov = cov.permute(1,0,2).to(torch.float32).to(device)\n",
        "\n",
        "        # Make predictions with the model.\n",
        "        y_hat_b = model(x, cov).float()\n",
        "        # Store the predictions.\n",
        "        y_hats.append(y_hat_b.cpu().detach().numpy())\n",
        "\n",
        "    # Calculate the mean prediction over all Monte Carlo samples.\n",
        "    ymc_hats = np.mean(y_hats, axis=0)\n",
        "    # Calculate the variance of the first component across samples.\n",
        "    eta_1s   = np.mean((ymc_hats[:,0] - np.stack(y_hats)[:,:,0])**2, axis=0)\n",
        "    return ymc_hats, eta_1s\n",
        "\n",
        "\n",
        "def inference(model, B=100, device=torch.device('cuda')):\n",
        "    # mc dropout\n",
        "    ymc_hats, eta_1s = mc_dropout(model, B, device)\n",
        "\n",
        "    # inherent noise\n",
        "    model.apply(dropout_off)\n",
        "    for (x, cov, v, y, ncov) in validation_loader:\n",
        "        x,cov,y = x.to(device), cov.to(device), y.to(device)\n",
        "        break\n",
        "    x = x.unsqueeze(2).permute(1, 0, 2).to(torch.float32).to(device)\n",
        "    cov = cov.permute(1,0,2).to(torch.float32).to(device)\n",
        "    y_hat_b = model(x, cov)\n",
        "\n",
        "    eta_2sq = np.mean(y_hat_b.cpu().detach().numpy()[:,0])\n",
        "    # total noise\n",
        "    etas = np.sqrt(eta_1s + eta_2sq)\n",
        "\n",
        "    for (x, cov, v, y, ncov) in test_loader:\n",
        "        break\n",
        "    normalized_label = y.T/v[:,0]\n",
        "    diff = torch.nansum(torch.mul((normalized_label.T - torch.tensor(ymc_hats)), (normalized_label.T - torch.tensor(ymc_hats)))).item()\n",
        "    test_rmse = np.sqrt(diff/len(test_set))\n",
        "\n",
        "    return test_rmse"
      ],
      "metadata": {
        "id": "CAj25HHvNJw6"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prednet_model = PredictionNetwork(encdec_model).cuda()\n",
        "loss_summary, optimizer = train_prediction_network(prednet_model, num_epochs = 2, device=torch.device('cuda'))\n",
        "evaluate_prediction_network(prednet_model, optimizer)\n",
        "test_rmse = inference(prednet_model, B=500, device=torch.device('cuda'))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QK_Jis4-MRCY",
        "outputId": "918a40d1-5445-496d-f42a-feeaa4076676"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Loss:0.07591520249843597: 100%|██████████| 5004/5004 [00:58<00:00, 85.86it/s]\n",
            "Loss:0.1028592586517334: 100%|██████████| 5004/5004 [00:58<00:00, 86.15it/s]\n",
            "100%|██████████| 34/34 [00:00<00:00, 227.80it/s]\n",
            "100%|██████████| 500/500 [00:02<00:00, 223.80it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(test_rmse)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "l_9gyMhPRaYF",
        "outputId": "c5dede8d-9e8d-4cde-d4b3-19867c258bf7"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.8222752754229762\n"
          ]
        }
      ]
    }
  ]
}