{
"cells": [
{
"cell_type": "markdown",
"id": "e6a6649b",
"metadata": {},
"source": [
"# Train an NNAlign model to predict both binding core and affinity of MHC class II ligands\n",
"The notebook consists of the following sections:\n",
"\n",
"0. Module imports, define functions, set constants\n",
"1. Load Data\n",
"2. Build Model\n",
"3. Select Hyper-paramerters\n",
"4. Compile Model\n",
"5. Train Model\n",
"6. Evaluation\n",
"\n",
"## Exercise\n",
"\n",
"The exercise is to modify the conventional Feed-Forward Neural Network (FFNN) so that it uses an NNAlign-like forward algorithm to identify a common binding core within each MHC class II ligand. In our case, the binding core length is fixed to 9 amino acids.\n",
"\n",
"Unlike the `DRB1_0101_cores` dataset, where the 9-mer binding cores were already identified and provided, the `DRB1_0101` dataset contains the original full-length MHC class II peptide ligands. These ligands have variable lengths, typically between 12 and 21 amino acids, and the position of the binding core is not given to the model.\n",
"\n",
"To handle this, for each peptide we generate all possible continuous 9-mer candidate cores. Each candidate core is encoded and passed through the neural network. The peptide-level prediction is then obtained by selecting the candidate core with the highest prediction score, using a max operation over all cores belonging to the same peptide.\n",
"\n",
"In this way, the model learns not only the binding pattern itself, but also which 9-mer segment of each peptide is most likely to act as the binding core.\n",
"\n",
"#### Performance evaluation\n",
"\n",
"Run the notebook and compare the performance of this NNAlign-like model with the conventional padded FFNN on the `DRB1_0101` dataset. Also compare the results with the model trained on `DRB1_0101_cores` with the conventional padded FFNN, where the binding cores were provided beforehand."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "539a3d32",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import pandas as pd\n",
"import numpy as np\n",
"import math\n",
"import matplotlib.pyplot as plt\n",
"from sklearn.metrics import roc_auc_score, roc_curve"
]
},
{
"cell_type": "markdown",
"id": "000ffcb6",
"metadata": {},
"source": [
"# Utility functions"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "60665494",
"metadata": {},
"outputs": [],
"source": [
"def load_blosum(filename):\n",
" \"\"\"\n",
" Load the BLOSUM matrix used for peptide encoding.\n",
" \"\"\"\n",
"\n",
" aa = ['A', 'R', 'N' ,'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X']\n",
" df = pd.read_csv(filename, sep='\\s+', comment='#', index_col=0)\n",
"\n",
" blosum_matrix = torch.tensor(df.loc[aa, aa].to_numpy(), dtype=torch.float32)\n",
"\n",
" aa_to_idx = {aa:i for i, aa in enumerate(aa)}\n",
"\n",
" return blosum_matrix, aa_to_idx\n",
"\n",
"\n",
"def load_peptide_target(filename):\n",
" \"\"\"\n",
" Read amino acid sequence of peptides and corresponding target values from text file.\n",
" \"\"\"\n",
" df = pd.read_csv(filename, sep='\\s+', usecols=[0,1], names=['peptide','target'])\n",
" return df.sort_values(by='target', ascending=False).reset_index(drop=True)\n",
"\n",
"\n",
"def encode_peptides(X_in, blosum_file, core_len=9):\n",
" \"\"\"\n",
" Encode all possible fixed-length candidate cores from each peptide using BLOSUM.\n",
"\n",
" Returns the encoded cores, one target per peptide, and an index mapping each core\n",
" to its original peptide.\n",
" \"\"\"\n",
" \n",
"\n",
" blosum50, aa_to_idx = load_blosum(blosum_file)\n",
"\n",
" X_idx = []\n",
" pep_idx = []\n",
" for i, peptide in enumerate(X_in[\"peptide\"]):\n",
" for j in range(len(peptide)-core_len+1):\n",
" core = peptide[j:j+core_len]\n",
" X_idx.append([aa_to_idx.get(aa, -1) for aa in core])\n",
" pep_idx.append(i)\n",
"\n",
" X_idx = torch.tensor(X_idx, dtype=torch.long)\n",
"\n",
" X_out = blosum50[X_idx]\n",
"\n",
" return X_out, torch.tensor(X_in[\"target\"].to_numpy(), dtype=torch.float32).reshape(-1, 1), torch.tensor(pep_idx, dtype=torch.long).reshape(-1, 1)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "dd4d9543",
"metadata": {},
"outputs": [],
"source": [
"def plot_losses(train_losses, valid_losses, n_epochs):\n",
" # Plotting the losses \n",
" fig,ax = plt.subplots(1,1, figsize=(9,5))\n",
" ax.plot(range(n_epochs), train_losses, label='Train loss', c='b')\n",
" ax.plot(range(n_epochs), valid_losses, label='Valid loss', c='m')\n",
" ax.legend()\n",
" plt.show()"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "aa3347ef",
"metadata": {},
"outputs": [],
"source": [
"def xavier_initialization_normal(layer_weights):\n",
" return nn.init.xavier_normal_(layer_weights)\n",
"\n",
"def random_initialization_normal(layer_weights):\n",
" return nn.init.normal_(layer_weights)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "47d6ec44",
"metadata": {},
"outputs": [],
"source": [
"def train_network(net, x_train, y_train, pep_idx_train, optimizer):\n",
" \"\"\"\n",
" Trains the network for a single epoch, running the forward and backward pass, and compute and return the loss.\n",
" \"\"\"\n",
"\n",
" net.train()\n",
"\n",
" loss_fn = nn.MSELoss()\n",
" \n",
" # Forward pass\n",
" a2 = net(x_train, pep_idx_train)\n",
"\n",
" # backward pass\n",
" loss = loss_fn(a2, y_train)\n",
" optimizer.zero_grad()\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" return loss.item()\n",
"\n",
"\n",
"def eval_network(net, x_valid, y_valid, pep_idx_valid):\n",
" \"\"\"\n",
" Evaluates the network ; Note that we do not update weights (no backward pass)\n",
" \"\"\"\n",
"\n",
" net.eval()\n",
"\n",
" loss_fn = nn.MSELoss()\n",
"\n",
" with torch.no_grad():\n",
" a2 = net(x_valid, pep_idx_valid)\n",
" loss = loss_fn(a2, y_valid)\n",
"\n",
" return loss.item()"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "f69a7b53",
"metadata": {},
"outputs": [],
"source": [
"def save_ffnn_model(filepath, model):\n",
" \"\"\"\n",
" Save PyTorch FFNN model weights and architecture metadata.\n",
" \"\"\"\n",
"\n",
" if not filepath.endswith(\".pt\"):\n",
" filepath = filepath + \".pt\"\n",
"\n",
" dict_to_save = {\n",
" \"input_size\": model.in_layer.in_features,\n",
" \"hidden_size\": model.in_layer.out_features,\n",
" \"output_size\": model.out_layer.out_features,\n",
" \"state_dict\": model.state_dict()\n",
" }\n",
"\n",
" torch.save(dict_to_save, filepath)\n",
"\n",
" print(f\"Saved FFNN model at {filepath}\")\n",
"\n",
"\n",
"def load_ffnn_model(filepath, model=None):\n",
" \"\"\"\n",
" Load PyTorch FFNN model weights and architecture metadata.\n",
" \"\"\"\n",
"\n",
" loaded_dict = torch.load(filepath, map_location=\"cpu\")\n",
"\n",
" if model is None:\n",
" model = SimpleFFNN(\n",
" input_size=loaded_dict[\"input_size\"],\n",
" hidden_size=loaded_dict[\"hidden_size\"],\n",
" output_size=loaded_dict[\"output_size\"]\n",
" )\n",
"\n",
" assert model.in_layer.in_features == loaded_dict[\"input_size\"], \\\n",
" f\"Input size mismatch. Model: {model.in_layer.in_features}, loaded: {loaded_dict['input_size']}\"\n",
"\n",
" assert model.in_layer.out_features == loaded_dict[\"hidden_size\"], \\\n",
" f\"Hidden size mismatch. Model: {model.in_layer.out_features}, loaded: {loaded_dict['hidden_size']}\"\n",
"\n",
" assert model.out_layer.out_features == loaded_dict[\"output_size\"], \\\n",
" f\"Output size mismatch. Model: {model.out_layer.out_features}, loaded: {loaded_dict['output_size']}\"\n",
"\n",
" model.load_state_dict(loaded_dict[\"state_dict\"])\n",
"\n",
" print(\n",
" f\"Model loaded successfully from {filepath}\\n\"\n",
" f\"with architecture: \"\n",
" f\"input_size={loaded_dict['input_size']}, \"\n",
" f\"hidden_size={loaded_dict['hidden_size']}, \"\n",
" f\"output_size={loaded_dict['output_size']}\"\n",
" )\n",
"\n",
" return model"
]
},
{
"cell_type": "markdown",
"id": "97275e6b",
"metadata": {},
"source": [
"# Model creation and training"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "d2555122",
"metadata": {},
"outputs": [],
"source": [
"class SimpleFFNN(nn.Module):\n",
"\n",
"\n",
" def __init__(self, input_size, hidden_size, output_size, initialization_function=xavier_initialization_normal):\n",
"\n",
" super().__init__()\n",
"\n",
" self.in_layer = nn.Linear(input_size, hidden_size)\n",
" self.out_layer = nn.Linear(hidden_size, output_size)\n",
" self.hidden_act = nn.ReLU()\n",
" self.out_act = nn.Sigmoid()\n",
"\n",
" initialization_function(self.in_layer.weight)\n",
" initialization_function(self.out_layer.weight)\n",
"\n",
" print(f'Input -> Hidden Layer Weight Matrix Shape: {self.in_layer.weight.shape}',\n",
" f'First Layer Bias Weights Vector Shape: {self.in_layer.bias.shape}',\n",
" f'Hidden -> Output layer Weight Matrix Shape: {self.out_layer.weight.shape}',\n",
" f'Second Layer Bias Weights Vector Shape: {self.out_layer.bias.shape}', sep=\"\\n\")\n",
"\n",
" \n",
" def forward(self, x, pep_idx):\n",
"\n",
" z1 = self.in_layer(x)\n",
" a1 = self.hidden_act(z1)\n",
" z2 = self.out_layer(a1)\n",
" a2 = self.out_act(z2)\n",
"\n",
" # Number of peptides in the batch\n",
" n_peptides = int(pep_idx.max().item()) + 1\n",
"\n",
" a2_best = torch.full(\n",
" size=(n_peptides, a2.shape[1]),\n",
" fill_value=-float(\"inf\"),\n",
" dtype=a2.dtype\n",
" )\n",
"\n",
" # For each core prediction in a2, assign it to its peptide group in pep_idx\n",
" # If several cores belong to the same peptide, keep the maximum score\n",
" a2_best.scatter_reduce_(\n",
" dim=0,\n",
" index=pep_idx,\n",
" src=a2,\n",
" reduce=\"amax\",\n",
" include_self=True\n",
" )\n",
" \n",
" return a2_best"
]
},
{
"cell_type": "markdown",
"id": "c2446b13",
"metadata": {},
"source": [
"# Data loading and encoding"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "9e4d5c6b",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Preview of the dataframe ; Peptides have to be *encoded* to BLOSUM matrices\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" peptide | \n",
" target | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" MVKVLDAVRGSPA | \n",
" 1.0 | \n",
"
\n",
" \n",
" | 1 | \n",
" KRFLYKKLPSVEGLHA | \n",
" 1.0 | \n",
"
\n",
" \n",
" | 2 | \n",
" AERLFKAMKGLGTRDN | \n",
" 1.0 | \n",
"
\n",
" \n",
" | 3 | \n",
" AEIEYAMAYSKAAFER | \n",
" 1.0 | \n",
"
\n",
" \n",
" | 4 | \n",
" APPVRHLIATQLLS | \n",
" 1.0 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" count | \n",
"
\n",
" \n",
" | len | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | 12 | \n",
" 348 | \n",
"
\n",
" \n",
" | 13 | \n",
" 358 | \n",
"
\n",
" \n",
" | 14 | \n",
" 363 | \n",
"
\n",
" \n",
" | 15 | \n",
" 430 | \n",
"
\n",
" \n",
" | 16 | \n",
" 406 | \n",
"
\n",
" \n",
" | 17 | \n",
" 424 | \n",
"
\n",
" \n",
" | 18 | \n",
" 418 | \n",
"
\n",
" \n",
" | 19 | \n",
" 431 | \n",
"
\n",
" \n",
" | 20 | \n",
" 404 | \n",
"
\n",
" \n",
" | 21 | \n",
" 418 | \n",
"
\n",
" \n",
"
\n",
"