Chenrui_Li.ipynb

{
 "metadata": {
  "kernelspec": {
   "language": "python",
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.13",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "sourceId": 68978,
     "databundleVersionId": 7709659,
     "sourceType": "competition"
    }
   ],
   "dockerImageVersionId": 30665,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook",
   "isGpuEnabled": true
  }
 },
 "nbformat_minor": 4,
 "nbformat": 4,
 "cells": [
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# !pip install ax-platform"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from ax.plot.slice import plot_slice\n",
    "from ax.service.ax_client import AxClient\n",
    "from ax.service.utils.instantiation import ObjectiveProperties\n",
    "from ax.utils.notebook.plotting import render\n",
    "from captum.attr import LayerIntegratedGradients\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# !conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=12.1 -c pytorch -c nvidia"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "print(torch.__version__)\n",
    "print(\"CUDA Available:\", torch.cuda.is_available())"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# # 1. Observe the data"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# one_train_pssm = pd.read_csv('deep-learning-for-msc-202324/train/1A0A_3_A_train.csv')\n",
    "# one_train_pssm"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# seqs_test = pd.read_csv('deep-learning-for-msc-202324/seqs_test.csv')\n",
    "# seqs_test"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# seqs_train = pd.read_csv('deep-learning-for-msc-202324/seqs_train.csv')\n",
    "# seqs_train"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# labels_train = pd.read_csv('deep-learning-for-msc-202324/labels_train.csv')\n",
    "# labels_train"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 2. Define the dataset and data loaders"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "class ProteinDataset(Dataset):\n",
    "    \"\"\"\n",
    "    It is used to load and preprocess the protein sequence and label data. It also reads the PSSM data from the provided files.\n",
    "    \n",
    "    :param seq_file_path: The file path to the sequence data\n",
    "    :type seq_file_path: str\n",
    "    :param pssm_files_path: The directory path to the PSSM files\n",
    "    :type pssm_files_path: str\n",
    "    :param label_file_path: The file path to the label data\n",
    "    :type label_file_path: str\n",
    "    :param indices: The indices of the data to select\n",
    "    :type indices: list\n",
    "    \n",
    "    :return: The protein dataset\n",
    "    :rtype: Dataset\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, seq_file_path, pssm_files_path, label_file_path=None, indices=None):\n",
    "        # Define the amino acid and structure mappings\n",
    "        self.amino_acid_to_ix = {'A': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6, 'H': 7, 'I': 8, 'K': 9, 'L': 10,\n",
    "                                 'M': 11, 'N': 12, 'P': 13, 'Q': 14, 'R': 15, 'S': 16, 'T': 17, 'V': 18, 'W': 19,\n",
    "                                 'Y': 20}\n",
    "        self.struct_to_ix = {'H': 0, 'E': 1, 'C': 2}\n",
    "\n",
    "        # Read the sequence data, and select a subset of the data if indices are provided(train set and val set)\n",
    "        self.seq_data = pd.read_csv(seq_file_path)\n",
    "        self.real_lan = [len(seq) for seq in self.seq_data['SEQUENCE']]\n",
    "        if indices is not None:\n",
    "            self.seq_data = self.seq_data.iloc[indices].reset_index(drop=True)\n",
    "\n",
    "        # Read the PSSM files path\n",
    "        self.pssm_files_path = pssm_files_path\n",
    "\n",
    "        # Read the label data, and select a subset of the data if indices are provided\n",
    "        if label_file_path:\n",
    "            self.label_data = pd.read_csv(label_file_path)\n",
    "            if indices is not None:\n",
    "                self.label_data = self.label_data.iloc[indices].reset_index(drop=True)\n",
    "        else:\n",
    "            self.label_data = None\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.seq_data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        It is used to get the sequence, label, and PSSM data for a given index.\n",
    "        \n",
    "        :param idx: The index of the data to retrieve \n",
    "        :type idx: int\n",
    "        :return: The sequence, label, and PSSM data\n",
    "        :rtype: dict\n",
    "        \"\"\"\n",
    "        pdb_id = self.seq_data.iloc[idx, 0]\n",
    "        sequence = self.seq_data.iloc[idx, 1]\n",
    "        sequence_encoded = [self.amino_acid_to_ix[aa] for aa in sequence]\n",
    "\n",
    "        if self.label_data is not None:\n",
    "            label_sequence = self.label_data[self.label_data['PDB_ID'] == pdb_id].iloc[0, 1]\n",
    "            label_encoded = [self.struct_to_ix[label] for label in label_sequence]\n",
    "        else:\n",
    "            label_encoded = [-1] * len(sequence)\n",
    "\n",
    "        pssm_file_path = os.path.join(self.pssm_files_path,\n",
    "                                      f\"{pdb_id}_test.csv\" if self.label_data is None else f\"{pdb_id}_train.csv\")\n",
    "        pssm_data = pd.read_csv(pssm_file_path).iloc[:, 2:].to_numpy()\n",
    "\n",
    "        return {\n",
    "            'pdb_id': pdb_id,\n",
    "            'sequence': torch.tensor(sequence_encoded, dtype=torch.long),\n",
    "            'labels': torch.tensor(label_encoded, dtype=torch.long),\n",
    "            'pssm': torch.tensor(pssm_data, dtype=torch.float),\n",
    "            'real_len': self.real_lan[idx]\n",
    "        }"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def protein_collate_fn(batch):\n",
    "    \"\"\"\n",
    "    It is used to process and prepare batched data during data loading. The DataLoader object can accept a function through the collate_fn parameter to customize how to combine multiple samples into a batch. This is particularly useful when dealing with sequence data, as the lengths of sequence data are often not consistent and need to be padded or otherwise processed to ensure that all data in a batch have consistent dimensions.\n",
    "    :param batch: The batch of data to process\n",
    "    :type batch: int\n",
    "    :return: A dictionary containing the processed batched data\n",
    "    :rtype: dict\n",
    "    \"\"\"\n",
    "    sequences, labels, pssms, pdb_ids, real_lens = zip(\n",
    "        *[(sample['sequence'], sample['labels'], sample['pssm'], sample['pdb_id'], sample['real_len']) for sample in\n",
    "          batch])\n",
    "\n",
    "    # Pad sequence and label data\n",
    "    sequences_padded = pad_sequence(sequences, batch_first=True)\n",
    "    labels_padded = pad_sequence(labels, batch_first=True)\n",
    "\n",
    "    # Pad PSSM data\n",
    "    pssms_padded = pad_sequence(pssms, batch_first=True, padding_value=-1)  # Use -1 to pad the PSSM data\n",
    "\n",
    "    return {'sequence': sequences_padded,\n",
    "            'labels': labels_padded,\n",
    "            'pssm': pssms_padded,\n",
    "            'pdb_id': pdb_ids,\n",
    "            'real_len': torch.tensor(real_lens)}"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def get_data_loaders(batch_size):\n",
    "    \"\"\"\n",
    "    It is used to create training and validation data loaders for the protein secondary structure prediction model. The function loads the sequence data, label data, and PSSM data, and then creates a training and validation split of the data. The function then creates data loaders for the training and validation data.\n",
    "    :param batch_size: batch size for the data loaders\n",
    "    :type batch_size: int\n",
    "    :return: a tuple containing the training and validation data loaders\n",
    "    :rtype: tuple\n",
    "    \"\"\"\n",
    "    seqs_file_path = 'kaggle/input/deep-learning-for-msc-202324/seqs_train.csv'\n",
    "    label_file_path = 'kaggle/input/deep-learning-for-msc-202324/labels_train.csv'\n",
    "    pssm_files_path = 'kaggle/input/deep-learning-for-msc-202324/train/'\n",
    "\n",
    "    # load the sequence data and generate a list of PDB IDs\n",
    "    seqs_data = pd.read_csv(seqs_file_path)\n",
    "    # pdb_ids = seqs_data['PDB_ID'].tolist()\n",
    "    indices = range(len(seqs_data))\n",
    "\n",
    "    # split the PDB_ID list into training and validation sets\n",
    "    pdb_ids_train, pdb_ids_val = train_test_split(indices, test_size=0.2, random_state=10)\n",
    "\n",
    "    # create training and validation data loaders\n",
    "    dataset_train = ProteinDataset(seqs_file_path, pssm_files_path, label_file_path, pdb_ids_train)\n",
    "    dataset_val = ProteinDataset(seqs_file_path, pssm_files_path, label_file_path, pdb_ids_val)\n",
    "\n",
    "    # create data loaders for training and validation\n",
    "    dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, collate_fn=protein_collate_fn,\n",
    "                                  pin_memory=True)\n",
    "    dataloader_val = DataLoader(dataset_val, batch_size=batch_size, shuffle=False, collate_fn=protein_collate_fn,\n",
    "                                pin_memory=True)\n",
    "\n",
    "    return dataloader_train, dataloader_val"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 3. Design the model(Fully Convolutional Networks, FCN)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "class ProteinSecondaryStructureFCNwithEmbedding(nn.Module):\n",
    "    \"\"\"\n",
    "    It is used to define a fully convolutional neural network (FCN) model for protein secondary structure prediction. The model uses an embedding layer to embed the amino acid sequence data, and then concatenates the embedded sequence data with the PSSM data before applying convolutional layers to learn features from the data. The model outputs a sequence of predictions for the secondary structure of the protein.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embedding_dim=10, dropout_rate=0.5):\n",
    "        num_amino_acids = 20\n",
    "        num_pssm_features = 20\n",
    "        num_classes = 3\n",
    "        super(ProteinSecondaryStructureFCNwithEmbedding, self).__init__()\n",
    "        self.embedding = nn.Embedding(num_embeddings=num_amino_acids + 1,\n",
    "                                      embedding_dim=embedding_dim)\n",
    "\n",
    "        # convolutional layers\n",
    "        self.layer1 = nn.Sequential(nn.Conv1d(embedding_dim + num_pssm_features, 64, kernel_size=5, padding=2),\n",
    "                                    nn.ReLU(), nn.BatchNorm1d(64), nn.Dropout(dropout_rate))\n",
    "        self.layer2 = nn.Sequential(nn.Conv1d(64, 128, kernel_size=5, padding=2), nn.ReLU(), nn.BatchNorm1d(128),\n",
    "                                    nn.Dropout(dropout_rate))\n",
    "        self.layer3 = nn.Sequential(nn.Conv1d(128, 256, kernel_size=5, padding=2), nn.ReLU(), nn.BatchNorm1d(256),\n",
    "                                    nn.Dropout(dropout_rate))\n",
    "        self.layer4 = nn.Sequential(nn.Conv1d(256, 512, kernel_size=5, padding=2), nn.ReLU(), nn.BatchNorm1d(512),\n",
    "                                    nn.Dropout(dropout_rate))\n",
    "        self.output_layer = nn.Conv1d(512, num_classes, kernel_size=1)\n",
    "\n",
    "    def forward(self, sequence, pssm):\n",
    "        # sequence: [batch_size, seq_len]\n",
    "        # pssm: [batch_size, seq_len, num_pssm_features]\n",
    "\n",
    "        # embed the sequence\n",
    "        embedded_sequence = self.embedding(sequence)  # [batch_size, seq_len, embedding_dim]\n",
    "\n",
    "        # adjust the dimensions of the embedded sequence to be concatenated with the PSSM\n",
    "        embedded_sequence = embedded_sequence.permute(0, 2, 1)  # [batch_size, embedding_dim, seq_len]\n",
    "\n",
    "        # concatenate the embedded sequence and PSSM data\n",
    "        combined_input = torch.cat((embedded_sequence, pssm.permute(0, 2, 1)),\n",
    "                                   dim=1)  # [batch_size, embedding_dim+num_pssm_features, seq_len]\n",
    "\n",
    "        # apply the network layers\n",
    "        out = self.layer1(combined_input)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = self.layer4(out)\n",
    "        out = self.output_layer(out)\n",
    "\n",
    "        return out.permute(0, 2, 1)  # [batch_size, seq_len, num_classes]"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def train_model(batch_size=4, num_epochs=10, embedding_dim=10, dropout_rate=0.5, lr=0.0001):\n",
    "    \"\"\"\n",
    "    It is used to train the protein secondary structure prediction model. The function creates the data loaders, initializes the model, loss function, and optimizer, and then trains the model for the specified number of epochs. The function returns the trained model and the training and validation performance metrics.\n",
    "    :param batch_size: batch size for the data loaders\n",
    "    :type batch_size: int\n",
    "    :param num_epochs: the number of epochs to train the model\n",
    "    :type num_epochs: int\n",
    "    :param embedding_dim: the dimension of the embedding layer\n",
    "    :type embedding_dim: int\n",
    "    :param dropout_rate: the dropout rate for the model\n",
    "    :type dropout_rate: float\n",
    "    :param lr: the learning rate for the optimizer\n",
    "    :type lr: float\n",
    "    :return: a dictionary containing the trained model and the training and validation performance metrics\n",
    "    :rtype: dict\n",
    "    \"\"\"\n",
    "    train_loader, val_loader = get_data_loaders(batch_size)\n",
    "    model = ProteinSecondaryStructureFCNwithEmbedding(embedding_dim=embedding_dim, dropout_rate=dropout_rate).to(device)\n",
    "    criterion = nn.CrossEntropyLoss(ignore_index=-1)  # Cross entropy loss function for classification problems\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)  # Adam optimizer\n",
    "\n",
    "    # initialize lists to store performance metrics\n",
    "    metrics = {\n",
    "        'train_loss': [],\n",
    "        'train_accuracy': [],\n",
    "        'train_precision': [],\n",
    "        'train_recall': [],\n",
    "        'train_f1': [],\n",
    "        'val_loss': [],\n",
    "        'val_accuracy': [],\n",
    "        'val_precision': [],\n",
    "        'val_recall': [],\n",
    "        'val_f1': []\n",
    "    }\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()  # Set the model to training mode\n",
    "\n",
    "        train_losses, train_true, train_pred = [], [], []\n",
    "\n",
    "        # Training loop\n",
    "        for batch in tqdm(train_loader, desc=f\"Training Epoch {epoch + 1}/{num_epochs}\"):\n",
    "            sequences = batch['sequence'].to(device)\n",
    "            labels = batch['labels'].to(device)\n",
    "            pssms = batch['pssm'].to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(sequences, pssms)\n",
    "            loss = criterion(outputs.reshape(-1, 3), labels.reshape(-1))\n",
    "\n",
    "            # Backward pass and optimization\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Track the training loss\n",
    "            train_losses.append(loss.item())\n",
    "            _, predicted = torch.max(outputs, 2)\n",
    "            train_true.extend(labels.view(-1).cpu().numpy())\n",
    "            train_pred.extend(predicted.view(-1).cpu().numpy())\n",
    "\n",
    "        # Calculate the training performance metrics\n",
    "        metrics['train_loss'].append(np.mean(train_losses))\n",
    "        metrics['train_accuracy'].append(accuracy_score(train_true, train_pred))\n",
    "        metrics['train_precision'].append(precision_score(train_true, train_pred, average='macro', zero_division=0))\n",
    "        metrics['train_recall'].append(recall_score(train_true, train_pred, average='macro', zero_division=0))\n",
    "        metrics['train_f1'].append(f1_score(train_true, train_pred, average='macro', zero_division=0))\n",
    "\n",
    "        print(\n",
    "            f'Epoch [{epoch + 1}/{num_epochs}], Train Loss: {metrics[\"train_loss\"][-1]:.4f}, Train Accuracy: {metrics[\"train_accuracy\"][-1]:.4f}, Train Precision: {metrics[\"train_precision\"][-1]:.4f}, Train Recall: {metrics[\"train_recall\"][-1]:.4f}, Train F1: {metrics[\"train_f1\"][-1]:.4f}')\n",
    "\n",
    "        torch.cuda.empty_cache()  # clean up CUDA memory\n",
    "\n",
    "        # Validation loop\n",
    "        model.eval()  # Set the model to evaluation mode\n",
    "        val_losses, val_true, val_pred = [], [], []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch in val_loader:\n",
    "                sequences = batch['sequence'].to(device)\n",
    "                labels = batch['labels'].to(device)\n",
    "                pssms = batch['pssm'].to(device)\n",
    "\n",
    "                # Forward pass\n",
    "                outputs = model(sequences, pssms)\n",
    "                loss = criterion(outputs.reshape(-1, 3), labels.reshape(-1))\n",
    "\n",
    "                # Track the validation loss\n",
    "                val_losses.append(loss.item())\n",
    "                _, predicted = torch.max(outputs, 2)\n",
    "                val_true.extend(labels.view(-1).cpu().numpy())\n",
    "                val_pred.extend(predicted.view(-1).cpu().numpy())\n",
    "\n",
    "        # Calculate the validation performance metrics\n",
    "        metrics['val_loss'].append(np.mean(val_losses))\n",
    "        metrics['val_accuracy'].append(accuracy_score(val_true, val_pred))\n",
    "        metrics['val_precision'].append(precision_score(val_true, val_pred, average='macro', zero_division=0))\n",
    "        metrics['val_recall'].append(recall_score(val_true, val_pred, average='macro', zero_division=0))\n",
    "        metrics['val_f1'].append(f1_score(val_true, val_pred, average='macro', zero_division=0))\n",
    "\n",
    "        print(\n",
    "            f'Epoch [{epoch + 1}/{num_epochs}], Validation Loss: {metrics[\"val_loss\"][-1]:.4f}, Validation Accuracy: {metrics[\"val_accuracy\"][-1]:.4f}, Validation Precision: {metrics[\"val_precision\"][-1]:.4f}, Validation Recall: {metrics[\"val_recall\"][-1]:.4f}, Validation F1: {metrics[\"val_f1\"][-1]:.4f}')\n",
    "\n",
    "        # torch.cuda.empty_cache()  # clean up CUDA memory\n",
    "\n",
    "    return {'model': model, 'metrics': metrics}"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 4. Train the model and evaluate the performance"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "torch.cuda.is_available()"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "device"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# # Ensure you have defined and instantiated your DataLoader here\n",
    "# train_rsults = train_model(batch_size=32, num_epochs=10, embedding_dim=15, dropout_rate=0.234673, lr=0.000143)"
   ],
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-03-16T17:34:38.637421Z",
     "iopub.execute_input": "2024-03-16T17:34:38.637937Z"
    },
    "trusted": true
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# display(train_rsults['model'])"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# # Save the trained model\n",
    "# torch.save(train_rsults['model'].state_dict(), 'prediction_model_default_para.pth')"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def plot_loss_accuracy_history(metrics):\n",
    "    \"\"\"\n",
    "    It is used to plot the loss and accuracy history of the model during training and validation.\n",
    "    :param metrics: A dictionary containing the training and validation performance metrics\n",
    "    :type metrics: dict\n",
    "    \"\"\"\n",
    "    train_loss_history = metrics['train_loss']\n",
    "    val_loss_history = metrics['val_loss']\n",
    "    train_accuracy_history = metrics['train_accuracy']\n",
    "    val_accuracy_history = metrics['val_accuracy']\n",
    "\n",
    "    epochs = range(1, len(train_loss_history) + 1)\n",
    "\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(epochs, train_loss_history, label='Train Loss')\n",
    "    plt.plot(epochs, val_loss_history, label='Validation Loss')\n",
    "    plt.title('Loss History')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(epochs, train_accuracy_history, label='Train Accuracy')\n",
    "    plt.plot(epochs, val_accuracy_history, label='Validation Accuracy')\n",
    "    plt.title('Accuracy History')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# display(train_rsults['metrics'])"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "# plot_loss_accuracy_history(train_rsults['metrics'])"
   ],
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 5. Hyperparameter optimization with Ax"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def train_evaluate(parameterization):\n",
    "    \"\"\"\n",
    "    It is used to train and evaluate the protein secondary structure prediction model with the given hyperparameters.\n",
    "    :param parameterization: A dictionary containing the hyperparameters to use for training and evaluation\n",
    "    :type parameterization: dict\n",
    "    :return: A dictionary containing the validation performance metrics\n",
    "    :rtype: dict\n",
    "    \"\"\"\n",
    "    # extract the hyperparameters from the parameterization\n",
    "    learning_rate = parameterization[\"lr\"]\n",
    "    dropout_rate = parameterization[\"dropout_rate\"]\n",
    "    batch_size = parameterization[\"batch_size\"]\n",
    "    embedding_dim = parameterization[\"embedding_dim\"]\n",
    "\n",
    "    # use the extracted hyperparameters to train and evaluate the model\n",
    "    metrics = \\\n",
    "        train_model(batch_size=batch_size, embedding_dim=embedding_dim, dropout_rate=dropout_rate, lr=learning_rate)[\n",
    "            'metrics']\n",
    "\n",
    "    val_loss = metrics['val_loss'][-1]\n",
    "    val_accuracy = metrics['val_accuracy'][-1]\n",
    "    val_precision = metrics['val_precision'][-1]\n",
    "    val_recall = metrics['val_recall'][-1]\n",
    "    val_f1 = metrics['val_f1'][-1]\n",
    "\n",
    "    return {\"val_loss\": (val_loss, 0.0), \"val_accuracy\": (val_accuracy, 0.0), \"val_precision\": (val_precision, 0.0),\n",
    "            \"val_recall\": (val_recall, 0.0), \"val_f1\": (val_f1, 0.0)}"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "ax_client = AxClient()\n",
    "ax_client.create_experiment(\n",
    "    name=\"protein_structure_prediction_experiment\",\n",
    "    parameters=[\n",
    "        {\"name\": \"lr\", \"type\": \"range\", \"bounds\": [5e-4, 3e-3], \"log_scale\": True},\n",
    "        {\"name\": \"dropout_rate\", \"type\": \"range\", \"bounds\": [0.15, 0.2]},\n",
    "        {\"name\": \"batch_size\", \"type\": \"choice\", \"values\": [32, 64]},\n",
    "        {\"name\": \"embedding_dim\", \"type\": \"choice\", \"values\": [4, 5, 6, 7]},\n",
    "    ],\n",
    "    objectives={\"val_loss\": ObjectiveProperties(minimize=True)},\n",
    "    tracking_metric_names=[\"val_accuracy\", \"val_precision\", \"val_recall\", \"val_f1\"],\n",
    ")\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "times_ax = 10\n",
    "for i in range(times_ax):\n",
    "    print(f\"Running optimization iteration {i + 1}/{times_ax}...\")\n",
    "    parameters, trial_index = ax_client.get_next_trial()\n",
    "    ax_client.complete_trial(trial_index=trial_index, raw_data=train_evaluate(parameters))"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "best_parameters_with_eval_value = ax_client.get_best_parameters()\n",
    "display(best_parameters_with_eval_value)"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "best_parameters = best_parameters_with_eval_value[0]"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Show the best parameters after ax"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "best_parameters"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "results_df = ax_client.get_trials_data_frame()\n",
    "results_df"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "After 70+ times of loop optimization, the best parameters are:\n",
    "{'lr': 0.0015518165052243,\n",
    " 'dropout_rate': 0.1858316733276795,\n",
    " 'batch_size': 64,\n",
    " 'embedding_dim': 7}"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# results_df.to_csv('10_trials_2_results.csv', index=False)"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Plotting to aid in parameterisation"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "render(ax_client.get_contour_plot(param_x=\"lr\", param_y=\"dropout_rate\", metric_name=\"val_accuracy\"))"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "render(ax_client.get_optimization_trace())"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "render(plot_slice(\n",
    "    model=ax_client.generation_strategy.model,\n",
    "    param_name='lr',\n",
    "    metric_name='val_accuracy'\n",
    "))"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "render(plot_slice(\n",
    "    model=ax_client.generation_strategy.model,\n",
    "    param_name='dropout_rate',\n",
    "    metric_name='val_accuracy'\n",
    "))"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 6. Train the model with the best hyperparameters and make predictions of the test set"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def test_train(best_parameters):\n",
    "    batch_size = best_parameters['batch_size']\n",
    "    embedding_dim = best_parameters['embedding_dim']\n",
    "    dropout_rate = best_parameters['dropout_rate']\n",
    "    lr = best_parameters['lr']\n",
    "    num_epochs = 100\n",
    "\n",
    "    best_results = train_model(batch_size=batch_size, num_epochs=num_epochs, embedding_dim=embedding_dim,\n",
    "                               dropout_rate=dropout_rate, lr=lr)\n",
    "\n",
    "    # save the trained model\n",
    "    filename = f\"{batch_size}_{num_epochs}_{embedding_dim}_{dropout_rate}_{lr}\"\n",
    "    torch.save(best_results['model'].state_dict(), f'{filename}.pth')\n",
    "\n",
    "    return best_results, filename"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Assist in the best parameters"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "file_path = '/all_ax.csv'\n",
    "if os.path.exists(file_path):\n",
    "    all_ax = pd.read_csv(file_path)\n",
    "    all_ax\n",
    "\n",
    "    max_val_accuracy_idx = all_ax['val_accuracy'].idxmax()\n",
    "    max_val_accuracy_idx\n",
    "    \n",
    "    max_val_accuracy_row = all_ax.loc[max_val_accuracy_idx]\n",
    "    max_val_accuracy_row\n",
    "\n",
    "    print(\n",
    "        f'Best parameters: \\n lr:\\t{max_val_accuracy_row[\"lr\"]}, dropout_rate:\\t{max_val_accuracy_row[\"dropout_rate\"]}, batch_size:\\t{max_val_accuracy_row[\"batch_size\"]}, embedding_dim:\\t{max_val_accuracy_row[\"embedding_dim\"]}')\n",
    "\n",
    "    best_parameters = {'lr': max_val_accuracy_row['lr'],\n",
    "                       'dropout_rate': max_val_accuracy_row['dropout_rate'],\n",
    "                       'batch_size': int(max_val_accuracy_row['batch_size']),\n",
    "                       'embedding_dim': int(max_val_accuracy_row['embedding_dim'])}\n",
    "else:\n",
    "    best_parameters = {'lr': 0.0015518165052243, 'dropout_rate': 0.1858316733276795, 'batch_size': 64,\n",
    "                       'embedding_dim': 7}\n",
    "\n",
    "print(best_parameters)"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "best_results, filename = test_train(best_parameters)"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def predict_and_save_test_set_modified(model, batch_size, filename):\n",
    "    seqs_test_path = 'kaggle/input/deep-learning-for-msc-202324/seqs_test.csv'\n",
    "    pssm_test_path = 'kaggle/input/deep-learning-for-msc-202324/test/'\n",
    "    output_csv_path = f'kaggle/output/{filename}.csv'\n",
    "\n",
    "    struct_to_ix = {'H': 0, 'E': 1, 'C': 2}\n",
    "    ix_to_struct = {v: k for k, v in struct_to_ix.items()}\n",
    "\n",
    "    # initialize the test dataset and data loader\n",
    "    dataset_test = ProteinDataset(seqs_test_path, pssm_test_path)\n",
    "    dataloader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False, collate_fn=protein_collate_fn,\n",
    "                                 pin_memory=True)\n",
    "\n",
    "    submission = [['ID', 'STRUCTURE']]\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_index, batch in enumerate(dataloader_test):\n",
    "            sequences = batch['sequence'].to(device)\n",
    "            pssms = batch['pssm'].to(device)\n",
    "            pdb_ids = batch['pdb_id']\n",
    "            real_lens = batch['real_len']\n",
    "\n",
    "            print(f\"Batch {batch_index + 1}/{len(dataloader_test)}\")\n",
    "\n",
    "            outputs = model(sequences, pssms)\n",
    "            _, predicted = torch.max(outputs, 2)\n",
    "\n",
    "            # split the batch predictions into individual sequences\n",
    "            for i, pdb_id in enumerate(pdb_ids):\n",
    "                real_len = real_lens[i].item()\n",
    "                # print(f\"{pdb_id} with {real_len} residues\")\n",
    "\n",
    "                for j in range(real_len):\n",
    "                    residue_prediction = predicted[i][j].cpu().item()\n",
    "                    submission.append([f\"{pdb_id}_{j + 1}\", ix_to_struct[residue_prediction]])\n",
    "                    # print(f\"{pdb_id}_{j+1} with {ix_to_struct[residue_prediction]}\")\n",
    "\n",
    "    # save the predictions to a CSV file\n",
    "    with open(output_csv_path, 'w') as f:\n",
    "        for line in submission:\n",
    "            f.write(','.join(line) + '\\n')\n",
    "\n",
    "    print(f\"Total predictions made: {len(submission) - 1}\")\n",
    "    print(f'file saved to {output_csv_path}')\n",
    "\n",
    "    return model, dataloader_test"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "filename"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "test_model, dataloader_test = predict_and_save_test_set_modified(best_results['model'], best_parameters['batch_size'],\n",
    "                                                                 filename)"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "plot_loss_accuracy_history(best_results['metrics'])"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 7. Explain the model predictions with Captum"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def layer_attributions(model, dataloader):\n",
    "    model.eval()\n",
    "    # initialize the LayerIntegratedGradients object, specifying the model and the layer of interest\n",
    "    lig = LayerIntegratedGradients(model, model.embedding)\n",
    "\n",
    "    for sample in tqdm(dataloader, desc=\"Calculating attributions\"):\n",
    "\n",
    "    sequences = sample['sequence'].to(device).long()\n",
    "    pssms = sample['pssm'].to(device).float()\n",
    "\n",
    "    # set the baseline for the sequence data\n",
    "    baseline = torch.zeros_like(sequences)\n",
    "\n",
    "    # focus on the first class\n",
    "    target = torch.tensor([0] * sequences.size(0)).to(device)\n",
    "\n",
    "    # calculate the attributions\n",
    "    attributions, delta = lig.attribute(inputs=(sequences, pssms),\n",
    "                                        baselines=(baseline, pssms),\n",
    "                                        target=target,\n",
    "                                        return_convergence_delta=True)\n",
    "\n",
    "    # convert the attributions and delta tensors to numpy arrays\n",
    "    attributions = attributions.cpu().detach().numpy()\n",
    "    delta = delta.cpu().detach().numpy()\n",
    "\n",
    "    return attributions, delta"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "attributions, delta = layer_attributions(test_model, dataloader_test)"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def visualize_attributions(attributions_seq, attributions_pssm, amino_acids=None):\n",
    "    \"\"\"\n",
    "    It is used to visualize the attributions of the model predictions for the amino acid sequence and PSSM features.\n",
    "    :param attributions_seq: The attributions of the model predictions for the amino acid sequence\n",
    "    :type attributions_seq: numpy.ndarray\n",
    "    :param attributions_pssm: The attributions of the model predictions for the PSSM features\n",
    "    :type attributions_pssm: numpy.ndarray\n",
    "    :param amino_acids: The amino acids in the sequence\n",
    "    :type amino_acids: list\n",
    "    \"\"\"\n",
    "    # only visualize the attributions for the first sample in the batch\n",
    "    sample_attributions_seq = attributions_seq[0]\n",
    "    sample_attributions_pssm = attributions_pssm[0]\n",
    "\n",
    "    fig, ax = plt.subplots(2, 1, figsize=(12, 8))\n",
    "\n",
    "    # plot the attributions for the amino acid sequence\n",
    "    ax[0].bar(range(len(sample_attributions_seq)), sample_attributions_seq)\n",
    "    ax[0].set_title('Attribution scores for amino acid sequences')\n",
    "    ax[0].set_xlabel('amino acid position')\n",
    "    ax[0].set_ylabel('attributable score')\n",
    "    if amino_acids:\n",
    "        ax[0].set_xticks(range(len(amino_acids)))\n",
    "        ax[0].set_xticklabels(amino_acids, rotation=45, ha=\"right\")\n",
    "\n",
    "    # plot the attributions for the PSSM features\n",
    "    cax = ax[1].matshow(sample_attributions_pssm, aspect='auto', cmap='viridis')\n",
    "    fig.colorbar(cax, ax=ax[1])\n",
    "    ax[1].set_title('Attribution scores for PSSM characteristics')\n",
    "    ax[1].set_xlabel('features')\n",
    "    ax[1].set_ylabel('amino acid position')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "visualize_attributions(attributions['seq'], attributions['pssm'])"
   ],
   "metadata": {
    "collapsed": false
   },
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Since it may take too long to run from scratch, I've put some output images in this place, which can be viewed by clicking on them:\n",
    "\n",
    "https://imgur.com/a/MtaC1b3\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "![64_100_7_0.1858316733276795_0.0015518165052243.png](https://i.imgur.com/I3EMIF1.png \"64_100_7_0.1858316733276795_0.0015518165052243.png\")\n",
    "\n",
    "![lr_trend](https://i.imgur.com/ZY5wkac.png \"lr_trend\")\n",
    "\n",
    "![lr&dropout_rate](https://i.imgur.com/nRN5T15.png \"lr&dropout_rate\")\n",
    "\n",
    "![dropout_rate_tredn](https://i.imgur.com/H3uQVPk.png \"dropout_rate_tredn\")\n",
    "\n",
    "![embedding_dim](https://i.imgur.com/soHrEwR.png \"embedding_dim\")\n",
    "\n",
    "![64_150_7_0.1858316733276795_0.0015518165052243](https://i.imgur.com/3gQye4Z.png)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ]
}