earthquake-prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JqdSSxEJLtds"
   },
   "source": [
    "# *KAGGLE CHALLENGE: LANL Earthquake Prediction*\n",
    "\n",
    "Un projet de Matthieu Dagommer, Paul Boulgakoff, Godefroy Bichon, Germain L'Hostis\n",
    "\n",
    "Versions utilisées:\n",
    "\n",
    "Python: 3.10.4\n",
    "Torch: 1.11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt \n",
    "import time\n",
    "\n",
    "import torch as th\n",
    "th.cuda.empty_cache()\n",
    "\n",
    "import torch.autograd as autograd\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "\n",
    "from torchsummary import summary \n",
    "\n",
    "import gzip\n",
    "import pickle\n",
    "\n",
    "import pandas as pd\n",
    "import scipy\n",
    "import scipy.stats as stats\n",
    "from scipy.stats import kurtosis, skew\n",
    "import csv\n",
    "import os\n",
    "import pickle\n",
    "import gc\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler, MinMaxScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "GeForce GTX 1050 Ti with Max-Q Design\n"
     ]
    }
   ],
   "source": [
    "### Ensuring CPU is connected\n",
    "\n",
    "print(th.cuda.is_available())\n",
    "print(th.cuda.get_device_name())\n",
    "device = th.device('cuda' if th.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# *General Hyperparameters*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setting General Hyperparameters \n",
    "\n",
    "batch_size = 100 # number of patches per batch\n",
    "valid_rate = 0.1 # fraction of data dedicated to validation\n",
    "overlap_rate = 0.2 # overlap\n",
    "\n",
    "# Parameters\n",
    "\n",
    "nrows = 200_000_000\n",
    "seed = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "th.manual_seed(seed)\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# *Data Preparation*\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Loading Data\n",
    "\n",
    "rootpath = os.getcwd() + \"/\"\n",
    "\n",
    "train_data = pd.read_csv(rootpath + \"train.csv\", usecols = ['acoustic_data', 'time_to_failure'], \\\n",
    "                         dtype = {'acoustic_data': np.int16, 'time_to_failure': np.float64}, nrows = nrows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = th.squeeze(th.tensor(train_data['acoustic_data'].values, dtype = th.int16))\n",
    "Y = th.squeeze(th.tensor(train_data['time_to_failure'].values, dtype = th.float64))\n",
    "\n",
    "del train_data\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "tvNnBnxCsLNS"
   },
   "outputs": [],
   "source": [
    "### Retrieve Test Data\n",
    "\n",
    "path = rootpath + \"test/\" + \"seg_00030f.csv\"\n",
    "test = pd.read_csv(path)\n",
    "patch_size = test.shape[0]\n",
    "\n",
    "#path = rootpath + \"sample_submission.csv\"\n",
    "#sample_submission = pd.read_csv(path)\n",
    "\n",
    "#test_ttf_counts = sample_submission.time_to_failure.value_counts()\n",
    "\n",
    "#X_test = []\n",
    "\n",
    "#for filename in os.listdir(rootpath + \"test\"):\n",
    "#    temp_df = pd.read_csv(rootpath + \"test/\" + filename)\n",
    "#    X_test.append(temp_df)\n",
    "\n",
    "#Xtest = th.zeros((len(X_test),time_window))\n",
    "#for i in range(len(X_test)):\n",
    "#  Xtest[i,:] = th.tensor(X_test[i][\"acoustic_data\"], dtype = th.float32)\n",
    "\n",
    "#Ytest = th.tensor(sample_submission[\"time_to_failure\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "nn_5lrssRSpJ",
    "outputId": "e0eb114e-3c57-458a-bec2-53709bbd85d8"
   },
   "outputs": [],
   "source": [
    "### Function to create patched sequences with some overlap out of the training data\n",
    "\n",
    "# Overlap is expressed as a fraction of patch size\n",
    "def patching(patch_size, X, Y, overlap_rate):\n",
    "    \n",
    "    overlap = int(patch_size*overlap_rate)\n",
    "    L = X.shape[0] # total length of the training acoustic signal\n",
    "    n_patch = int(np.floor((L-patch_size)/(patch_size-overlap)+1))\n",
    "    ids_no_seism = []\n",
    "\n",
    "    X_patch = th.zeros(n_patch,patch_size)\n",
    "    Y_patch = th.zeros(n_patch)\n",
    "    \n",
    "    for i in range (n_patch):\n",
    "        X_patch[i,:] = X[i*(patch_size-overlap):(i+1)*patch_size-i*overlap] \n",
    "        Y_patch[i] = Y[(i+1)*patch_size - i*overlap] \n",
    "    \n",
    "        # Removing patches with no seism\n",
    "        if th.min(Y[i*(patch_size-overlap):(i+1)*patch_size - i*overlap]) > 0.001:\n",
    "            ids_no_seism.append(i)\n",
    "    \n",
    "    X_patch = X_patch[ids_no_seism]\n",
    "    Y_patch = Y_patch[ids_no_seism]\n",
    "    \n",
    "    return(n_patch, X_patch, Y_patch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Normalizing Data\n",
    "\n",
    "ss = StandardScaler()\n",
    "mm = MinMaxScaler()\n",
    "    \n",
    "X = th.squeeze(th.from_numpy(ss.fit_transform(np.array(X).reshape(-1, 1))))\n",
    "Y = th.squeeze(th.from_numpy(mm.fit_transform(np.array(Y).reshape(-1, 1))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1656, 150000])\n",
      "There are  1656 time series available for training and validation after patching with overlap. \n",
      "\n"
     ]
    }
   ],
   "source": [
    "### Distribution of data in patches\n",
    "\n",
    "_, X_patch, Y_patch = patching(patch_size, X, Y, overlap_rate = overlap_rate)\n",
    "\n",
    "del X; del Y\n",
    "gc.collect()\n",
    "\n",
    "print(X_patch.shape)\n",
    "print(\"There are \", X_patch.shape[0], \"time series available for training and validation after patching with overlap. \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Initializating Data Sets\n",
    "\n",
    "\n",
    "# Shuffling data\n",
    "idx = np.arange(X_patch.shape[0])\n",
    "shuffled_idx = np.random.shuffle(idx)\n",
    "X_patch = th.squeeze(X_patch[shuffled_idx,:])\n",
    "Y_patch = th.squeeze(Y_patch[shuffled_idx])\n",
    "\n",
    "N_samples = X_patch.shape[0]\n",
    "N_valid = int(valid_rate*N_samples) # number of validation patches\n",
    "\n",
    "X_valid = X_patch[:N_valid,:].cpu()\n",
    "Y_valid = Y_patch[:N_valid].cpu()\n",
    "\n",
    "\n",
    "# Inputs of nn.Conv1d must have the following shape: (N, C_in, *),\n",
    "#where N: number of samples (batch size), C_in: number of input channels, *: can be any dimension (150_000 for our time series)\n",
    "\n",
    "# Add channel dimension\n",
    "X_valid = th.unsqueeze(X_valid, 1)\n",
    "X_train = X_patch[N_valid:,:].cpu()\n",
    "Y_train = Y_patch[N_valid:].cpu()\n",
    "N_train = X_train.shape[0]\n",
    "\n",
    "n_batch = int(X_train.shape[0] / batch_size) # Number of batch per epoch\n",
    "\n",
    "X_train = X_train.reshape([N_train, 1, patch_size])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# *Function for Results plotting*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Plotting Graphs\n",
    "\n",
    "def plot_and_save_results(train_losses, valid_losses, best_mvd, best_mtd, min_tl, min_vl, mm, model_name):\n",
    "\n",
    "    best_epoch = np.fromiter(valid_losses, dtype=float).argmin()\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 2, figsize = (20,10))\n",
    "    plt.rcParams['font.size'] = '20'\n",
    "    ax[0].set(title = \"Losses: Training and Validation\") \n",
    "    ax[0].set_xlabel(\"epochs\", fontsize = 20)\n",
    "    ax[0].set_ylabel(\"MSE\", fontsize = 20)\n",
    "    ax[0].plot(train_losses,\"r\", label = \"Training\", linewidth = 3)\n",
    "    ax[0].plot(valid_losses, \"b\", label = \"Validation\", linewidth = 3)\n",
    "    ax[0].legend(loc = \"upper right\", fontsize = 18)\n",
    "    ax[0].axvline(x=int(best_epoch), color = 'black', linestyle =\"--\", linewidth = 3)\n",
    "    ax[0].annotate(\"Best epoch: {}\\nMSE_train: {:.3f}\\nMSE_valid: {:.3f}\".format(int(best_epoch), min_tl, min_vl), \\\n",
    "                   xy = (0.5,0.5), xycoords = 'axes fraction')\n",
    "    \n",
    "    best_mvd_plot = mm.inverse_transform(best_mvd.reshape(-1, 1))\n",
    "    best_mtd_plot = mm.inverse_transform(best_mtd.reshape(-1, 1))\n",
    "\n",
    "    N_valid = best_mvd_plot.shape[0]\n",
    "    \n",
    "    mean = float(np.mean(best_mvd_plot))\n",
    "    std_dev = float(np.std(best_mvd_plot))\n",
    "    kurt = float(kurtosis(best_mvd_plot))\n",
    "    skewn = float(skew(best_mvd_plot))\n",
    "    q1 = float(np.quantile(best_mvd_plot, 0.25))\n",
    "    median = float(np.quantile(best_mvd_plot, 0.5))\n",
    "    q3 = float(np.quantile(best_mvd_plot, 0.75))\n",
    "    mae = np.absolute(best_mvd_plot).sum() / N_valid\n",
    "    mse = np.square(best_mvd_plot).sum() / N_valid\n",
    "    \n",
    "    text = \"mean: {:.3f}\\nstd: {:.3f}\\nkurt: {:.3f}\\nskew: {:.3f}\\nq1: {:.3f}\\nmed: {:.3f}\\nq3: {:.3f}\\niqr: {:.3f}\\nmae: {:.3f}\\nmse: {:.3f}\\n\".format(mean, std_dev, kurt, skewn, q1, median, q3, q3-q1, mae, mse)\n",
    "    \n",
    "    ax[1].hist(best_mvd_plot, alpha = 0.3, label = \"validation set\", bins = 100, density = True, range = (-16, 16))\n",
    "    ax[1].set(title = \"TTF error distributions at best epoch\")\n",
    "    ax[1].set_xlabel(\"Error (seconds)\", fontsize = 20)\n",
    "    ax[1].set_ylabel(\"Density\", fontsize = 20)\n",
    "    ax[1].hist(best_mtd_plot, alpha = 0.3, label = \"training set\", bins = 100, density = True, range = (-16, 16))\n",
    "    ax[1].annotate(text, xy =(-15, 0.1))\n",
    "    ax[1].legend()\n",
    "\n",
    "\n",
    "    plt.gcf()\n",
    "    plt.savefig('Models/' + model_name + '/' + model_name + \"_plot.jpg\")\n",
    "    plt.show()\n",
    "\n",
    "    model_features = {\"N_samples\": N_samples, \"N_train\": N_train, \"N_valid\": N_valid, \\\n",
    "                      \"overlap_rate\": overlap_rate, \"learning_rate\": learning_rate, \"num_epochs\": num_epochs, \\\n",
    "                     \"seed\": seed, \"batch_size\": batch_size, \"train_losses\": train_losses, \"valid_losses\": valid_losses, \\\n",
    "                     \"valid_differences\": valid_differences, \"best_mtd\": best_mtd, \"best_mvd\": best_mvd, \"min_tl\": min_tl, \\\n",
    "                      \"min_vl\": min_vl}\n",
    "\n",
    "    pickle.dump(model_features, open('Models/' + model_name + '/' + model_name + \".p\", \"wb\" ))\n",
    "    \n",
    "    model_features_display = {\"N_samples\": N_samples, \"N_train\": N_train, \"N_valid\": N_valid, \\\n",
    "                  \"overlap_rate\": overlap_rate, \"learning_rate\": learning_rate, \"num_epochs\": num_epochs, \\\n",
    "                 \"seed\": seed, \"batch_size\": batch_size}\n",
    "    \n",
    "    pickle.dump(model_features_display, open('Models/' + model_name + '/' + model_name + \"display.p\", \"wb\" ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# *1D Convolutional Neural Network*\n",
    "\n",
    "This section contains two 1D CNN architecture that we designed during this project, with their respective hyperparameters and the training loop.\n",
    "To use the LSTM instead, you can skip this section and go to the next."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Hyperparameters specific to 1D CNN\n",
    "\n",
    "learning_rate=0.0001\n",
    "num_epochs=50\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bOtetCjsDj8t"
   },
   "outputs": [],
   "source": [
    "### 1D Convolutional Network #1\n",
    "\n",
    "class NN(nn.Module):\n",
    "    \n",
    "    def __init__ (self,num_classes=1):\n",
    "        super(NN, self).__init__()\n",
    "        self.fc1=nn.Conv1d(in_channels=1, out_channels=1, \n",
    "                             kernel_size=10, padding=1, stride=5)\n",
    "        self.fc2=nn.AdaptiveMaxPool1d(200)\n",
    "        self.fct3=nn.Linear(200,50)\n",
    "        self.fct4=nn.Linear(50,num_classes)\n",
    "    \n",
    "    def forward (self,x):\n",
    "        x=self.fc1(x)\n",
    "        x=self.fc2(x)\n",
    "        x=self.fct3(x)\n",
    "        x=self.fct4(x)\n",
    "        return x\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### 1D Convolutional Network #2\n",
    "\n",
    "class NN_2(nn.Module):\n",
    "    def __init__(self,num_classes=1):\n",
    "        super(NN_2, self).__init__()\n",
    "        self.seq_1=nn.Sequential(nn.Conv1d(in_channels = 1, out_channels = 16, kernel_size = 10, stride = 5),\n",
    "                                 nn.ReLU(),\n",
    "                                 nn.MaxPool1d(kernel_size = 2),\n",
    "                                 nn.Conv1d(in_channels = 16, out_channels = 32, kernel_size = 10, stride = 5, \\\n",
    "                                           padding = 'valid'),\n",
    "                                 nn.ReLU(),\n",
    "                                 nn.MaxPool1d(kernel_size = 2),\n",
    "                                 nn.Conv1d(in_channels = 32, out_channels = 64, kernel_size = 10, stride = 5, \\\n",
    "                                           padding = 'valid'),\n",
    "                                 nn.ReLU()\n",
    "                                )\n",
    "        \n",
    "        self.seq_2=nn.Sequential(nn.Dropout(p = 0.4),\n",
    "                                 nn.Linear(in_features = 64, out_features = 1)\n",
    "                                )\n",
    "    def forward(self,x):\n",
    "        x=self.seq_1(x)\n",
    "        x = th.mean(x, dim = 2, keepdim = True)\n",
    "        x = th.squeeze(x)\n",
    "        x=self.seq_2(x)\n",
    "        return (x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hH24ktcBDvrq",
    "outputId": "55f4b5f1-f013-4e34-da59-49e8d0d19e99"
   },
   "outputs": [],
   "source": [
    "### Initialize Network, Define Loss and optimizer\n",
    "\n",
    "model = NN_2()\n",
    "#print(model)\n",
    "model.to(device)\n",
    "loss=nn.MSELoss()\n",
    "optimizer=th.optim.Adam(model.parameters(),learning_rate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Plot Model Graph \n",
    "\n",
    "#from torchviz import make_dot\n",
    "#yhat = model(X_valid.to(device))\n",
    "#make_dot(yhat, params=dict(list(model.named_parameters()))).render(\"cnn_2\", format=\"png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jIl7K2gvZd3U"
   },
   "outputs": [],
   "source": [
    "### Check accuracy during training\n",
    "\n",
    "def check_accuracy(X_valid, Y_valid, model):\n",
    "    \n",
    "    device = 'cuda'\n",
    "    N_valid = X_valid.shape[0]\n",
    "    real_ttf = th.squeeze(Y_valid)\n",
    "    loss_fn = nn.MSELoss()\n",
    "    \n",
    "    with th.no_grad():\n",
    "\n",
    "        scores = th.squeeze(model(X_valid.to(device))).cpu()\n",
    "        valid_loss_ = loss_fn(scores, real_ttf).cpu()\n",
    "        valid_differences = scores[:] - real_ttf[:]\n",
    "        valid_differences = valid_differences.cpu().numpy()\n",
    "        valid_loss = valid_loss_.item()\n",
    "    \n",
    "    return valid_loss, valid_differences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QVS0pEmAD4P4",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "### Train network\n",
    "\n",
    "th.cuda.empty_cache()\n",
    "\n",
    "train_losses, valid_losses = [], []\n",
    "\n",
    "best_mvd, best_mtd = [], [] # validation, training differences at best epoch\n",
    "\n",
    "min_vl = 1000\n",
    "min_tl = 1000\n",
    "\n",
    "N_train_per_epoch = n_batch*batch_size \n",
    "\n",
    "start = time.perf_counter()\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "\n",
    "    th.cuda.empty_cache()\n",
    "    \n",
    "    # Generating Random Batches every epoch\n",
    "    random_idx = np.random.randint(0, N_train, N_train_per_epoch)\n",
    "    X_train_batch = th.reshape(X_train[random_idx], [n_batch, batch_size, 1, patch_size])\n",
    "    Y_train_batch = th.reshape(Y_train[random_idx], [n_batch, batch_size, 1, 1])\n",
    "    \n",
    "    running_loss = 0\n",
    "    _loss = 0\n",
    "\n",
    "    for ids in range (n_batch):\n",
    "        \n",
    "        # Set gradients to 0\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Model Prediction\n",
    "        predicted_ttfs = th.squeeze(model(X_train_batch[ids].cuda())).cpu()\n",
    "        # Loss function\n",
    "        _loss = loss(predicted_ttfs, th.squeeze(Y_train_batch[ids]))\n",
    "        \n",
    "        del predicted_ttfs\n",
    "        gc.collect()\n",
    "        \n",
    "        # Backpropagation\n",
    "        _loss.backward()\n",
    "        running_loss += _loss.item()\n",
    "        \n",
    "        # Updating Model Weights\n",
    "        optimizer.step()\n",
    "        \n",
    "    optimizer.zero_grad()\n",
    "    th.cuda.empty_cache()\n",
    "    \n",
    "    model.eval()\n",
    "    with th.no_grad():\n",
    "    \n",
    "        # Check Accuracy with Validation Set\n",
    "        \n",
    "        Y_valid_pred = th.squeeze(model(X_valid.cuda())).cpu()\n",
    "        valid_loss = loss(Y_valid_pred, Y_valid)\n",
    "        valid_differences = Y_valid_pred[:] - Y_valid[:]\n",
    "        valid_differences = valid_differences.numpy()\n",
    "        \n",
    "        valid_losses.append(valid_loss.item())\n",
    "        train_losses.append(running_loss / n_batch)\n",
    "        \n",
    "        if valid_loss < min_vl:\n",
    "            \n",
    "            min_vl = valid_loss\n",
    "            min_tl = running_loss\n",
    "            \n",
    "            best_mvd = valid_differences\n",
    "            \n",
    "            Y_final = th.squeeze(model(X_train[:int(N_train/10)].to(device))).cpu()\n",
    "            \n",
    "            best_mtd = Y_train[:int(N_train/10)] - Y_final[:]\n",
    "            best_mtd = best_mtd.cpu().detach().numpy()\n",
    "\n",
    "    model.train()\n",
    "    \n",
    "    if epoch%10 == 0:\n",
    "        print(\"Epoch: {}\\t\".format(epoch),\n",
    "                \"train Loss: {:.5f}.. \".format(train_losses[-1]),\n",
    "                \"valid Loss: {:.5f}.. \".format(valid_losses[-1])) \n",
    "\n",
    "print(\"---------- Best : {:.3f}\".format(min(valid_losses)), \" at epoch \" \n",
    "    , np.fromiter(valid_losses, dtype=float).argmin(), \" / \", epoch + 1)\n",
    "\n",
    "\n",
    "end = time.perf_counter()\n",
    "print(\"\\ntime elapsed: {:.3f}\".format(end - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = input(\"Choose a name for the model: \")\n",
    "os.mkdir('Models/' + model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_and_save_results(train_losses, valid_losses, best_mvd, best_mtd, min_tl, min_vl, mm, model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eqi1BoEdwan0"
   },
   "source": [
    "# *LSTM*\n",
    "\n",
    "LSTM is a type of recurrent neural network and is a relevant architecture to treat temporal sequences. \n",
    "\n",
    "Pytorch Notes:\n",
    "\n",
    "\"Before getting to the example, note a few things. Pytorch’s LSTM expects all of its inputs to be 3D tensors. The semantics of the axes of these tensors is important. The first axis is the sequence itself, the second indexes instances in the mini-batch, and the third indexes elements of the input.\"\n",
    "\n",
    "https://pytorch.org/tutorials/beginner/nlp/sequence_models_tutorial.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Hyperparameters specific to LSTM\n",
    "\n",
    "num_epochs = 1000 #1000 epochs\n",
    "learning_rate = 0.001 #0.001 lr\n",
    "hidden_size = 1 #number of features in hidden state\n",
    "num_layers = 1 #number of stacked lstm layers => should stay at 1 unless you want to combine two LSTMs together\n",
    "N_sub_patches = 250\n",
    "N_features = 10\n",
    "device = \"cuda\"\n",
    "n_batch = 1 # One batch with all training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "_4ULTZ0TYn8A"
   },
   "outputs": [],
   "source": [
    "def feature_expansion(X, N_sub_patches, L_sub_patch):\n",
    "    \n",
    "    #print(type(X))\n",
    "    N_samples = X.shape[0]\n",
    "    x = X.reshape(N_samples, N_sub_patches, L_sub_patch)\n",
    "    x_mean = np.mean(x, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_std = np.std(x, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_skew = np.array(scipy.stats.skew(x, axis = 2), dtype = np.double).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_kurt = np.array(scipy.stats.kurtosis(x, axis = 2), dtype = np.double).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_min = np.min(x, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_max = np.max(x, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_q1 = np.quantile(x, 0.25, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_med = np.quantile(x, 0.5, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_q3 = np.quantile(x, 0.75, axis = 2).reshape(N_samples, N_sub_patches, 1)\n",
    "    x_iqr = x_q3 - x_q1\n",
    "    \n",
    "    X_rearranged = np.concatenate((x_mean, x_std, x_skew, x_kurt, x_min, x_max, x_q1, x_med, x_q3, x_iqr), axis = 2)\n",
    "    return X_rearranged\n",
    "    \n",
    "\n",
    "def lstm_feature_engineering(X_train, Y_train, X_valid, Y_valid, patch_size, N_sub_patches):\n",
    "\n",
    "    X_train = np.array(th.squeeze(X_train)); Y_train = np.array(Y_train)\n",
    "    X_valid = np.array(th.squeeze(X_valid)); Y_valid = np.array(Y_valid)\n",
    "\n",
    "    L_seq = X_train.shape[1]\n",
    "    L_sub_patch = int(L_seq / N_sub_patches)\n",
    "    N_features = 10 # mean, std, skew, kurt, min, max, quantiles 0.25, 0.5, 0.75, inter-quartile range\n",
    "    \n",
    "    X_train_rearranged = feature_expansion(X_train, N_sub_patches, L_sub_patch)\n",
    "    X_valid_rearranged = feature_expansion(X_valid, N_sub_patches, L_sub_patch)\n",
    "\n",
    "    X_train = th.from_numpy(X_train_rearranged)\n",
    "    X_valid = th.from_numpy(X_valid_rearranged)\n",
    "\n",
    "    Y_train = th.from_numpy(Y_train); Y_valid = th.from_numpy(Y_valid)\n",
    "    #Y_train = th.unsqueeze(Y_train, -1); Y_valid = th.unsqueeze(Y_valid, -1)\n",
    "\n",
    "    return X_train, Y_train, X_valid, Y_valid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "CFL8PvDnjZvs"
   },
   "outputs": [],
   "source": [
    "X_train_, Y_train_, X_valid_, Y_valid_ = lstm_feature_engineering(X_train, Y_train, X_valid, Y_valid, patch_size, N_sub_patches = N_sub_patches)\n",
    "N_features = X_train_.shape[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "VgH9isMrwZkp"
   },
   "outputs": [],
   "source": [
    "#Inspired from https://cnvrg.io/pytorch-lstm/\n",
    "\n",
    "class Lstm(nn.Module):\n",
    "    \n",
    "    def __init__(self, N_features, hidden_size, num_layers, seq_length):\n",
    "        super(Lstm, self).__init__()\n",
    "\n",
    "        self.num_layers = num_layers # number of layers\n",
    "        self.N_features = N_features # number of features\n",
    "        self.hidden_size = hidden_size # hidden state\n",
    "        self.seq_length = seq_length # sequence length\n",
    "\n",
    "        self.lstm = nn.LSTM(input_size=N_features, hidden_size=hidden_size,\n",
    "                          num_layers=num_layers, batch_first=True) #lstm => Input Shape : (N_batch, L_seq, N_feature)\n",
    "    \n",
    "    def forward(self,x):\n",
    "        \n",
    "        output, (_,_) = self.lstm(x.float())\n",
    "        out = output[:,-1,0] # Retrieving predicted time at the end of the training\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "gc9D9u4T1C33"
   },
   "outputs": [],
   "source": [
    "L_seq = N_sub_patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "g7DE7GjCGeRR"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Lstm(\n",
       "  (lstm): LSTM(10, 1, batch_first=True)\n",
       ")"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Lstm(N_features, hidden_size, num_layers, L_seq) #our lstm class\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "P9m0tjrtEzO2"
   },
   "outputs": [],
   "source": [
    "loss = th.nn.MSELoss()    # mean-squared error for regression\n",
    "optimizer = th.optim.Adam(model.parameters(), lr=learning_rate) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "1k1geqipE79l",
    "outputId": "67aa1abe-5aee-49a5-84f3-7d627ed8c6cc",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0\t train Loss: 0.91807..  valid Loss: 1.10350.. \n",
      "Epoch: 10\t train Loss: 0.87599..  valid Loss: 1.05680.. \n",
      "Epoch: 20\t train Loss: 0.83104..  valid Loss: 1.00782.. \n",
      "Epoch: 30\t train Loss: 0.78497..  valid Loss: 0.95279.. \n",
      "Epoch: 40\t train Loss: 0.73718..  valid Loss: 0.89875.. \n",
      "Epoch: 50\t train Loss: 0.68942..  valid Loss: 0.84610.. \n",
      "Epoch: 60\t train Loss: 0.64310..  valid Loss: 0.79642.. \n",
      "Epoch: 70\t train Loss: 0.59853..  valid Loss: 0.75052.. \n",
      "Epoch: 80\t train Loss: 0.55639..  valid Loss: 0.70860.. \n",
      "Epoch: 90\t train Loss: 0.51767..  valid Loss: 0.67017.. \n",
      "Epoch: 100\t train Loss: 0.48204..  valid Loss: 0.63451.. \n",
      "Epoch: 110\t train Loss: 0.44913..  valid Loss: 0.60116.. \n",
      "Epoch: 120\t train Loss: 0.41868..  valid Loss: 0.56993.. \n",
      "Epoch: 130\t train Loss: 0.39056..  valid Loss: 0.54080.. \n",
      "Epoch: 140\t train Loss: 0.36469..  valid Loss: 0.51375.. \n",
      "Epoch: 150\t train Loss: 0.34099..  valid Loss: 0.48873.. \n",
      "Epoch: 160\t train Loss: 0.31938..  valid Loss: 0.46559.. \n",
      "Epoch: 170\t train Loss: 0.29972..  valid Loss: 0.44417.. \n",
      "Epoch: 180\t train Loss: 0.28183..  valid Loss: 0.42426.. \n",
      "Epoch: 190\t train Loss: 0.26551..  valid Loss: 0.40568.. \n",
      "Epoch: 200\t train Loss: 0.25058..  valid Loss: 0.38825.. \n",
      "Epoch: 210\t train Loss: 0.23687..  valid Loss: 0.37182.. \n",
      "Epoch: 220\t train Loss: 0.22421..  valid Loss: 0.35626.. \n",
      "Epoch: 230\t train Loss: 0.21248..  valid Loss: 0.34147.. \n",
      "Epoch: 240\t train Loss: 0.20156..  valid Loss: 0.32733.. \n",
      "Epoch: 250\t train Loss: 0.19136..  valid Loss: 0.31378.. \n",
      "Epoch: 260\t train Loss: 0.18179..  valid Loss: 0.30075.. \n",
      "Epoch: 270\t train Loss: 0.17280..  valid Loss: 0.28819.. \n",
      "Epoch: 280\t train Loss: 0.16435..  valid Loss: 0.27608.. \n",
      "Epoch: 290\t train Loss: 0.15642..  valid Loss: 0.26439.. \n",
      "Epoch: 300\t train Loss: 0.14899..  valid Loss: 0.25313.. \n",
      "Epoch: 310\t train Loss: 0.14204..  valid Loss: 0.24232.. \n",
      "Epoch: 320\t train Loss: 0.13558..  valid Loss: 0.23199.. \n",
      "Epoch: 330\t train Loss: 0.12961..  valid Loss: 0.22215.. \n",
      "Epoch: 340\t train Loss: 0.12413..  valid Loss: 0.21286.. \n",
      "Epoch: 350\t train Loss: 0.11913..  valid Loss: 0.20413.. \n",
      "Epoch: 360\t train Loss: 0.11460..  valid Loss: 0.19598.. \n",
      "Epoch: 370\t train Loss: 0.11052..  valid Loss: 0.18842.. \n",
      "Epoch: 380\t train Loss: 0.10686..  valid Loss: 0.18146.. \n",
      "Epoch: 390\t train Loss: 0.10359..  valid Loss: 0.17506.. \n",
      "Epoch: 400\t train Loss: 0.10068..  valid Loss: 0.16921.. \n",
      "Epoch: 410\t train Loss: 0.09810..  valid Loss: 0.16388.. \n",
      "Epoch: 420\t train Loss: 0.09581..  valid Loss: 0.15904.. \n",
      "Epoch: 430\t train Loss: 0.09380..  valid Loss: 0.15464.. \n",
      "Epoch: 440\t train Loss: 0.09201..  valid Loss: 0.15066.. \n",
      "Epoch: 450\t train Loss: 0.09043..  valid Loss: 0.14706.. \n",
      "Epoch: 460\t train Loss: 0.08903..  valid Loss: 0.14380.. \n",
      "Epoch: 470\t train Loss: 0.08779..  valid Loss: 0.14085.. \n",
      "Epoch: 480\t train Loss: 0.08668..  valid Loss: 0.13819.. \n",
      "Epoch: 490\t train Loss: 0.08569..  valid Loss: 0.13578.. \n",
      "Epoch: 500\t train Loss: 0.08479..  valid Loss: 0.13360.. \n",
      "Epoch: 510\t train Loss: 0.08399..  valid Loss: 0.13162.. \n",
      "Epoch: 520\t train Loss: 0.08325..  valid Loss: 0.12983.. \n",
      "Epoch: 530\t train Loss: 0.08257..  valid Loss: 0.12821.. \n",
      "Epoch: 540\t train Loss: 0.08194..  valid Loss: 0.12674.. \n",
      "Epoch: 550\t train Loss: 0.08135..  valid Loss: 0.12540.. \n",
      "Epoch: 560\t train Loss: 0.08079..  valid Loss: 0.12418.. \n",
      "Epoch: 570\t train Loss: 0.08027..  valid Loss: 0.12306.. \n",
      "Epoch: 580\t train Loss: 0.07978..  valid Loss: 0.12202.. \n",
      "Epoch: 590\t train Loss: 0.07932..  valid Loss: 0.12105.. \n",
      "Epoch: 600\t train Loss: 0.07889..  valid Loss: 0.12014.. \n",
      "Epoch: 610\t train Loss: 0.07848..  valid Loss: 0.11927.. \n",
      "Epoch: 620\t train Loss: 0.07810..  valid Loss: 0.11844.. \n",
      "Epoch: 630\t train Loss: 0.07774..  valid Loss: 0.11766.. \n",
      "Epoch: 640\t train Loss: 0.07740..  valid Loss: 0.11692.. \n",
      "Epoch: 650\t train Loss: 0.07708..  valid Loss: 0.11622.. \n",
      "Epoch: 660\t train Loss: 0.07677..  valid Loss: 0.11556.. \n",
      "Epoch: 670\t train Loss: 0.07647..  valid Loss: 0.11493.. \n",
      "Epoch: 680\t train Loss: 0.07619..  valid Loss: 0.11434.. \n",
      "Epoch: 690\t train Loss: 0.07591..  valid Loss: 0.11379.. \n",
      "Epoch: 700\t train Loss: 0.07564..  valid Loss: 0.11327.. \n",
      "Epoch: 710\t train Loss: 0.07538..  valid Loss: 0.11277.. \n",
      "Epoch: 720\t train Loss: 0.07512..  valid Loss: 0.11231.. \n",
      "Epoch: 730\t train Loss: 0.07487..  valid Loss: 0.11187.. \n",
      "Epoch: 740\t train Loss: 0.07462..  valid Loss: 0.11145.. \n",
      "Epoch: 750\t train Loss: 0.07438..  valid Loss: 0.11106.. \n",
      "Epoch: 760\t train Loss: 0.07414..  valid Loss: 0.11068.. \n",
      "Epoch: 770\t train Loss: 0.07391..  valid Loss: 0.11033.. \n",
      "Epoch: 780\t train Loss: 0.07368..  valid Loss: 0.10998.. \n",
      "Epoch: 790\t train Loss: 0.07345..  valid Loss: 0.10966.. \n",
      "Epoch: 800\t train Loss: 0.07323..  valid Loss: 0.10934.. \n",
      "Epoch: 810\t train Loss: 0.07301..  valid Loss: 0.10904.. \n",
      "Epoch: 820\t train Loss: 0.07279..  valid Loss: 0.10875.. \n",
      "Epoch: 830\t train Loss: 0.07257..  valid Loss: 0.10847.. \n",
      "Epoch: 840\t train Loss: 0.07236..  valid Loss: 0.10819.. \n",
      "Epoch: 850\t train Loss: 0.07215..  valid Loss: 0.10793.. \n",
      "Epoch: 860\t train Loss: 0.07194..  valid Loss: 0.10767.. \n",
      "Epoch: 870\t train Loss: 0.07173..  valid Loss: 0.10741.. \n",
      "Epoch: 880\t train Loss: 0.07153..  valid Loss: 0.10717.. \n",
      "Epoch: 890\t train Loss: 0.07133..  valid Loss: 0.10692.. \n",
      "Epoch: 900\t train Loss: 0.07113..  valid Loss: 0.10668.. \n",
      "Epoch: 910\t train Loss: 0.07094..  valid Loss: 0.10645.. \n",
      "Epoch: 920\t train Loss: 0.07075..  valid Loss: 0.10621.. \n",
      "Epoch: 930\t train Loss: 0.07056..  valid Loss: 0.10598.. \n",
      "Epoch: 940\t train Loss: 0.07037..  valid Loss: 0.10575.. \n",
      "Epoch: 950\t train Loss: 0.07019..  valid Loss: 0.10553.. \n",
      "Epoch: 960\t train Loss: 0.07001..  valid Loss: 0.10530.. \n",
      "Epoch: 970\t train Loss: 0.06984..  valid Loss: 0.10508.. \n",
      "Epoch: 980\t train Loss: 0.06966..  valid Loss: 0.10485.. \n",
      "Epoch: 990\t train Loss: 0.06949..  valid Loss: 0.10463.. \n",
      "---------- Best : 0.104  at epoch  999  /  1000\n",
      "\n",
      "time elapsed: 45.426\n"
     ]
    }
   ],
   "source": [
    "# Training LSTM\n",
    "\n",
    "th.cuda.empty_cache()\n",
    "\n",
    "train_losses, valid_losses = [], []\n",
    "best_mvd, best_mtd = [], []\n",
    "\n",
    "min_vl = 1000\n",
    "min_tl = 1000\n",
    "\n",
    "start = time.perf_counter()\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "    outputs = model.forward(X_train_.cuda()) #forward pass\n",
    "    optimizer.zero_grad() #calculate the gradient, manually setting to 0\n",
    "    # obtain the loss function\n",
    "    _loss = loss(outputs, Y_train_.cuda())\n",
    "    _loss.backward() #calculates the loss of the loss function\n",
    "    running_loss = _loss.item()\n",
    "    optimizer.step() #improve from loss, i.e backprop\n",
    "\n",
    "\n",
    "    model.eval()\n",
    "    with th.no_grad():\n",
    "        \n",
    "        Y_valid_pred = th.squeeze(model(X_valid_.cuda())).cpu()\n",
    "        valid_loss = loss(Y_valid_pred, Y_valid_)\n",
    "        valid_differences = Y_valid_pred[:] - Y_valid_[:]\n",
    "        valid_differences = valid_differences.numpy()\n",
    "        \n",
    "        valid_losses.append(valid_loss.item())\n",
    "        train_losses.append(running_loss / n_batch)\n",
    "        \n",
    "        if valid_loss < min_vl:\n",
    "\n",
    "            min_vl = valid_loss\n",
    "            min_tl = running_loss\n",
    "\n",
    "            best_mvd = valid_differences\n",
    "\n",
    "            Y_final = th.squeeze(model(X_train_.cuda())).cpu()\n",
    "\n",
    "            best_mtd = Y_train_ - Y_final\n",
    "            best_mtd = best_mtd.cpu().detach().numpy()\n",
    "            \n",
    "    model.train() \n",
    "    \n",
    "    if epoch%10 == 0:\n",
    "        print(\"Epoch: {}\\t\".format(epoch),\n",
    "                \"train Loss: {:.5f}.. \".format(train_losses[-1]),\n",
    "                \"valid Loss: {:.5f}.. \".format(valid_losses[-1])) \n",
    "        \n",
    "print(\"---------- Best : {:.3f}\".format(min(valid_losses)), \" at epoch \" \n",
    "    , np.fromiter(valid_losses, dtype=float).argmin(), \" / \", epoch + 1)\n",
    "    \n",
    "end = time.perf_counter()\n",
    "print(\"\\ntime elapsed: {:.3f}\".format(end - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = input(\"Choose a name for the model: \")\n",
    "os.mkdir('Models/' + model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_and_save_results(train_losses, valid_losses, best_mvd, best_mtd, min_tl, min_vl, mm, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_features = pickle.load(open(model_name + \"display.p\", \"rb\" ))\n",
    "print(model_features)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:earthquake-prediction]",
   "language": "python",
   "name": "conda-env-earthquake-prediction-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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}