{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "N47Ue-53WqAG" }, "source": [ "# P2PXML - Deep Geometric Framework to Predict Antibody-Antigen Binding Affinity" ] }, { "cell_type": "markdown", "metadata": { "id": "35qWIH9dWzu7" }, "source": [ "List of resources:\n", "\n", "* Paper: https://www.biorxiv.org/content/10.1101/2024.06.09.598103v1\n", "* Project Page: https://drug-discovery-entc.github.io/p2pxml/\n", "* Codes: https://github.com/Drug-Discovery-ENTC/p2pxml/\n", "* Our dataset: https://zenodo.org/records/11531319\n", "\n", "The codes, including this notebook, are released under MIT license (https://github.com/Drug-Discovery-ENTC/p2pxml/blob/main/LICENSE) and the dataset is released under Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License (https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)." ] }, { "cell_type": "markdown", "metadata": { "id": "sA_FHFqAY25r" }, "source": [ "Citation:\n", "\n", "```bibtex\n", "@article{bandara2024deep,\n", " title={Deep Geometric Framework to Predict Antibody-Antigen Binding Affinity},\n", " author={Bandara, Nuwan Sriyantha and Premathilaka, Dasun and Chandanayake, Sachini and Hettiarachchi, Sahan and Varenthirarajah, Vithurshan and Munasinghe, Aravinda and Madhawa, Kaushalya and Charles, Subodha},\n", " journal={bioRxiv},\n", " pages={2024--06},\n", " year={2024},\n", " publisher={Cold Spring Harbor Laboratory}\n", "}\n", "```" ] }, { "cell_type": "markdown", "metadata": { "id": "VnvPPD74XZBK" }, "source": [ "*Disclaimer: This is a minimal working demonstration and the results may be inaccurate. Even though our models are trained on our curated dataset, which is the largest and most generalized publicly-available dataset for antibody-antigen binding affinity prediction to-date in the literature (to the best of our knowledge), it is still biased towards certain antigen variants such as HIV and SARS-CoV-2 due to their abundance in terms of number of data points in the dataset.*" ] }, { "cell_type": "markdown", "metadata": { "id": "3Yam7nw1w2vB" }, "source": [ "### Run all the cells" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "cellView": "form", "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-07-30T02:29:43.140859Z", "iopub.status.busy": "2026-07-30T02:29:43.140678Z", "iopub.status.idle": "2026-07-30T02:29:43.143751Z", "shell.execute_reply": "2026-07-30T02:29:43.143172Z" }, "id": "0LGtmIUycmB_", "outputId": "2d00a0bd-9a18-4f2f-9b5d-5d4cd2dd481d" }, "outputs": [], "source": [ "# Dependencies are already installed in the evo_bio environment.\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "cellView": "form", "execution": { "iopub.execute_input": "2026-07-30T02:29:43.146347Z", "iopub.status.busy": "2026-07-30T02:29:43.146161Z", "iopub.status.idle": "2026-07-30T02:30:01.976519Z", "shell.execute_reply": "2026-07-30T02:30:01.975099Z" }, "id": "z72fVm0ybIX5" }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/public/home/scnb9biwet/.conda/envs/evo_bio/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] } ], "source": [ "# @title\n", "import pandas as pd\n", "import numpy as np\n", "from tqdm import tqdm\n", "import math\n", "\n", "import torch\n", "import torch.nn as nn\n", "from torch.optim import AdamW\n", "from torch import Tensor\n", "import torch.nn.functional as F\n", "from torch.nn import Parameter\n", "from torch.nn import Sequential, Linear, ReLU, MultiheadAttention, Dropout, LayerNorm, AvgPool1d\n", "from torchmetrics.functional import mean_absolute_error\n", "import torch.optim as optim\n", "from torch.utils.data import Dataset, DataLoader\n", "from torch.optim.lr_scheduler import LambdaLR\n", "from tqdm import tqdm\n", "\n", "from torch.nn.init import zeros_,xavier_normal_\n", "\n", "import os\n", "\n", "import networkx as nx\n", "import torch_geometric.data as Data\n", "from torch_geometric.loader import DataLoader\n", "from torch_geometric.nn import GCNConv, global_mean_pool, GATConv\n", "from torch_geometric.transforms import NormalizeScale\n", "from torch_geometric.data import Batch\n", "from torchmetrics.functional import mean_absolute_error\n", "from torch.utils.data import random_split\n", "\n", "import matplotlib.pyplot as plt\n", "\n", "from biopandas.pdb import PandasPdb\n", "import periodictable\n", "from Bio import SeqIO\n", "from Bio.PDB import PDBParser\n", "from Bio.SeqUtils import seq1\n", "\n", "from sklearn.model_selection import train_test_split\n", "\n", "import warnings\n", "warnings.filterwarnings(\"ignore\")\n", "\n", "import logging, sys" ] }, { "cell_type": "markdown", "metadata": { "id": "Yyu_GMViwOen" }, "source": [ "**Note**\n", "\n", "Here, input the **.pdb** files for both antibody and antigen. *Note that the current maximum FASTA sequence lengths for antibody and antigen are 669 and 3102 respectively.*\n", "\n", "Example PDB files for an antibody-antigen pair can be found at https://github.com/Drug-Discovery-ENTC/p2pxml/tree/main/data" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "cellView": "form", "colab": { "base_uri": "https://localhost:8080/", "height": 184 }, "execution": { "iopub.execute_input": "2026-07-30T02:30:01.979875Z", "iopub.status.busy": "2026-07-30T02:30:01.979328Z", "iopub.status.idle": "2026-07-30T02:30:02.077205Z", "shell.execute_reply": "2026-07-30T02:30:02.076465Z" }, "id": "vJY-76IGngwE", "outputId": "6da824c5-99a0-462c-cc6f-0f54752321c8" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " Ab Ag log(IC50)\n", "0 10-1074 0013095_2_11 0.0\n" ] } ], "source": [ "import os\n", "import shutil\n", "import pandas as pd\n", "\n", "ROOT = os.environ[\"P2PXML_ROOT\"]\n", "SMOKE_DIR = os.path.join(ROOT, \"scripts\", \"smoke\")\n", "\n", "os.makedirs(os.path.join(SMOKE_DIR, \"antibodies\"), exist_ok=True)\n", "os.makedirs(os.path.join(SMOKE_DIR, \"antigens\"), exist_ok=True)\n", "os.makedirs(os.path.join(SMOKE_DIR, \"graph_data\"), exist_ok=True)\n", "\n", "antibody_name = \"10-1074\"\n", "antigen_name = \"0013095_2_11\"\n", "\n", "shutil.copy2(\n", " os.path.join(ROOT, \"conf\", \"data\", antibody_name + \".pdb\"),\n", " os.path.join(SMOKE_DIR, \"antibodies\", antibody_name + \".pdb\"),\n", ")\n", "shutil.copy2(\n", " os.path.join(ROOT, \"conf\", \"data\", antigen_name + \".pdb\"),\n", " os.path.join(SMOKE_DIR, \"antigens\", antigen_name + \".pdb\"),\n", ")\n", "\n", "test_df = pd.DataFrame({\n", " \"Ab\": [antibody_name],\n", " \"Ag\": [antigen_name],\n", " \"log(IC50)\": [0.0],\n", "})\n", "\n", "print(test_df)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "cellView": "form", "execution": { "iopub.execute_input": "2026-07-30T02:30:02.079727Z", "iopub.status.busy": "2026-07-30T02:30:02.079511Z", "iopub.status.idle": "2026-07-30T02:30:02.121823Z", "shell.execute_reply": "2026-07-30T02:30:02.121107Z" }, "id": "kigJrq_0XkxU" }, "outputs": [], "source": [ "# @title\n", "max_antibody_sequence_length = 669\n", "max_antigen_sequence_length = 3102\n", "\n", "class ProteinDataset(Dataset):\n", " def __init__(self, df, max_antibody_sequence_length, max_antigen_sequence_length, save_dir=os.path.join(SMOKE_DIR, 'graph_data')):\n", " self.sequences = df['Ab'].values\n", " self.viruses = df['Ag'].values\n", " self.labels = df['log(IC50)'].values\n", " self.max_antibody_sequence_length = max_antibody_sequence_length\n", " self.max_antigen_sequence_length = max_antigen_sequence_length\n", " self.save_dir = save_dir\n", " os.makedirs(save_dir, exist_ok=True)\n", "\n", " def __len__(self):\n", " return len(self.sequences)\n", "\n", " def __getitem__(self, idx):\n", " sequence = self.sequences[idx]\n", " virus = self.viruses[idx]\n", " label = self.labels[idx]\n", " label = torch.tensor(label, dtype=torch.float64)\n", "\n", " antibody = self.load_or_generate_graph(sequence, self.max_antibody_sequence_length, 'antibodies')\n", " antigen = self.load_or_generate_graph(virus, self.max_antigen_sequence_length, 'antigens')\n", "\n", " if antibody is not None and antigen is not None:\n", " return antibody, antigen, label\n", " else:\n", " return self.__getitem__((idx + 1) % len(self)) # Ensure idx is within bounds\n", "\n", " def load_or_generate_graph(self, pdb_file, max_sequence_length, graph_type):\n", " graph_path = os.path.join(self.save_dir, f'{graph_type}_{pdb_file}.pt')\n", "\n", " if os.path.exists(graph_path):\n", " return torch.load(graph_path)\n", "\n", " if graph_type == 'antigens':\n", " graph_constructed = self.pdb_to_graph_virus(pdb_file, max_sequence_length)\n", " else:\n", " graph_constructed = self.pdb_to_graph_antibody(pdb_file, max_sequence_length)\n", "\n", " if graph_constructed is not None:\n", " torch.save(graph_constructed, graph_path)\n", " return graph_constructed\n", "\n", " def pdb_to_graph_virus(self, pdb_file, max_antigen_sequence_length):\n", " return self.pdb_to_graph(pdb_file, max_antigen_sequence_length, 'antigens')\n", "\n", " def pdb_to_graph_antibody(self, pdb_file, max_antibody_sequence_length):\n", " return self.pdb_to_graph(pdb_file, max_antibody_sequence_length, 'antibodies')\n", "\n", " def pdb_to_graph(self, pdb_file, max_sequence_length, graph_type):\n", " seq_length = max_sequence_length\n", " amino_acids = list(\"ACDEFGHIKLMNPQRSTVWY\")\n", " aa_to_index = {aa: i for i, aa in enumerate(amino_acids)}\n", " pdbparser = PDBParser()\n", "\n", " try:\n", " structure = pdbparser.get_structure(pdb_file, os.path.join(SMOKE_DIR, graph_type, pdb_file + '.pdb'))\n", " chains = {chain.id: seq1(''.join(residue.resname for residue in chain)) for chain in structure.get_chains()}\n", " full_sequence = ''\n", " for value in chains.values():\n", " full_sequence += value\n", " full_sequence = full_sequence.replace(\"X\", \"\")\n", "\n", " if len(full_sequence) > seq_length:\n", " print(f\"Exceeds max length {graph_type}\")\n", " return None\n", "\n", " indices = [aa_to_index[aa] for aa in full_sequence]\n", " encoded = F.one_hot(torch.tensor(indices), num_classes=len(amino_acids)).float()\n", " padded_encoded = F.pad(encoded.flatten(), (0, max(seq_length * len(amino_acids) - encoded.flatten().shape[0], 0)))\n", "\n", " ppdb = PandasPdb()\n", " ppdb.read_pdb(os.path.join(SMOKE_DIR, graph_type, pdb_file + '.pdb'))\n", " coords = ppdb.df['ATOM'][['x_coord', 'y_coord', 'z_coord']].values\n", " atomic_nums = ppdb.df['ATOM']['element_symbol'].apply(lambda symbol: periodictable.elements.symbol(symbol).number).values\n", "\n", " graph = nx.Graph()\n", " num_atoms = len(coords)\n", " for i in range(num_atoms):\n", " graph.add_node(i, x=coords[i][0], y=coords[i][1], z=coords[i][2], atomic_number=atomic_nums[i])\n", " for i in range(num_atoms):\n", " for j in range(i + 1, num_atoms):\n", " dist = ((coords[i] - coords[j]) ** 2).sum() ** 0.5\n", " if dist < 5:\n", " bond_strength = 1 / dist\n", " graph.add_edge(i, j, distance=dist, bond_strength=bond_strength)\n", "\n", " edge_attrs = {}\n", "\n", " for u, v, data in graph.edges(data=True):\n", " edge_attrs[(u, v)] = [data['distance'], data['bond_strength']]\n", " edge_attrs[(v, u)] = [data['distance'], data['bond_strength']]\n", "\n", " data = Data.Data(\n", " x=torch.tensor(list(nx.get_node_attributes(graph, 'x').values())).to(torch.float64),\n", " y_coord=torch.tensor(list(nx.get_node_attributes(graph, 'y').values())).to(torch.float64),\n", " z_coord=torch.tensor(list(nx.get_node_attributes(graph, 'z').values())).to(torch.float64),\n", " pos=torch.tensor(coords).to(torch.float64),\n", " edge_index=torch.tensor(list(graph.edges)).to(torch.float64).t().contiguous(),\n", " edge_attr=torch.tensor([edge_attrs[e] for e in graph.edges()]).to(torch.float64),\n", " z=torch.tensor(list(nx.get_node_attributes(graph, 'atomic_number').values())).to(torch.float64),\n", " seq=padded_encoded.to(torch.float64),\n", " y = torch.tensor([0.0]).to(torch.float64)\n", " )\n", " return data\n", "\n", " except Exception as e:\n", " print(e)\n", " return None\n", "\n", "test_ds = ProteinDataset(test_df, max_antibody_sequence_length, max_antigen_sequence_length)\n", "test_loader = DataLoader(test_ds, batch_size=1, shuffle=False)" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "cellView": "form", "execution": { "iopub.execute_input": "2026-07-30T02:30:02.123856Z", "iopub.status.busy": "2026-07-30T02:30:02.123653Z", "iopub.status.idle": "2026-07-30T02:30:02.155676Z", "shell.execute_reply": "2026-07-30T02:30:02.154975Z" }, "id": "apGATVycfjcM" }, "outputs": [], "source": [ "# @title\n", "class SelfAttention(nn.Module):\n", " def __init__(self, embed_dim, num_heads=16):\n", " super(SelfAttention, self).__init__()\n", " self.embed_dim = embed_dim\n", " self.num_heads = num_heads\n", " if embed_dim % num_heads != 0:\n", " raise ValueError(f\"embedding dimension = {embed_dim} should be divisible by number of heads = {num_heads}\")\n", " self.head_dim = embed_dim // num_heads\n", " self.query_dense = nn.Linear(embed_dim, embed_dim)\n", " self.key_dense = nn.Linear(embed_dim, embed_dim)\n", " self.value_dense = nn.Linear(embed_dim, embed_dim)\n", " self.combine_heads = nn.Linear(embed_dim, embed_dim)\n", "\n", " def forward(self, inputs):\n", " query = self.query_dense(inputs)\n", " key = self.key_dense(inputs)\n", " value = self.value_dense(inputs)\n", " query = query.view(-1, self.num_heads, self.head_dim)\n", " key = key.view(-1, self.num_heads, self.head_dim)\n", " value = value.view(-1, self.num_heads, self.head_dim)\n", " query = query.permute(1, 0, 2)\n", " key = key.permute(1, 0, 2)\n", " value = value.permute(1, 0, 2)\n", " dot_product = torch.matmul(query, key.permute(0, 2, 1))\n", " scaled_dot_product = dot_product / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))\n", " attention_weights = torch.softmax(scaled_dot_product, dim=-1)\n", " output = torch.matmul(attention_weights, value)\n", " output = output.permute(1, 0, 2)\n", " output = output.view(-1, self.embed_dim)\n", " output = self.combine_heads(output)\n", " return output\n", "\n", "class TransformerBlock(nn.Module):\n", " def __init__(self, embed_dim, num_heads, dense_dim=1024, dropout_rate=0.1):\n", " super(TransformerBlock, self).__init__()\n", " self.attention = SelfAttention(embed_dim, num_heads)\n", " self.dropout1 = nn.Dropout(dropout_rate)\n", " self.norm1 = nn.LayerNorm(embed_dim, eps=1e-6)\n", " self.dense1 = nn.Linear(embed_dim, dense_dim)\n", " self.dropout2 = nn.Dropout(dropout_rate)\n", " self.norm2 = nn.LayerNorm(embed_dim, eps=1e-6)\n", " self.dense2 = nn.Linear(dense_dim, embed_dim)\n", "\n", " def forward(self, inputs):\n", " attention_output = self.attention(inputs)\n", " attention_output = self.dropout1(attention_output)\n", " output1 = self.norm1(inputs + attention_output)\n", " dense_output = self.dense1(output1)\n", " dense_output = self.dropout2(dense_output)\n", " output2 = self.norm2(output1 + dense_output)\n", " output = self.dense2(output2)\n", " return output\n", "\n", "class CrossAttention(nn.Module):\n", " def __init__(self, dim, input_shape):\n", " super(CrossAttention, self).__init__()\n", " self.dim = dim\n", " self.input_shape = input_shape\n", "\n", " self.Wq = nn.Parameter(torch.Tensor(input_shape[0][-1], self.dim))\n", " self.Wk = nn.Parameter(torch.Tensor(input_shape[1][-1], self.dim))\n", " self.Wv = nn.Parameter(torch.Tensor(input_shape[1][-1], self.dim))\n", "\n", " nn.init.xavier_uniform_(self.Wq)\n", " nn.init.xavier_uniform_(self.Wk)\n", " nn.init.xavier_uniform_(self.Wv)\n", "\n", " def forward(self, inputs):\n", " x, y = inputs\n", "\n", " Q = torch.matmul(x, self.Wq)\n", " K = torch.matmul(y, self.Wk)\n", " V = torch.matmul(y, self.Wv)\n", "\n", " attn_weights = torch.matmul(Q, K.t()) / torch.sqrt(torch.tensor(self.dim, dtype=torch.float64))\n", " attn_weights = F.softmax(attn_weights, dim=-1)\n", " attn_output = attn_weights * V\n", " output = torch.cat([x, attn_output], dim=-1)\n", "\n", " return output\n", "\n", "class CombinedModel(nn.Module):\n", " def __init__(self, hidden_channels=128, num_layers=16):\n", " super(CombinedModel, self).__init__()\n", "\n", " self.cross_attn_1 = CrossAttention(128, [(1024,),(1024,)])\n", " self.cross_attn_2 = CrossAttention(128,[(1024,),(1024,)])\n", " self.cross_attn = CrossAttention(128,[(1152,),(1152,)])\n", " self.self_atten_1 = nn.AdaptiveAvgPool1d(1)\n", " self.self_atten_2 = nn.AdaptiveAvgPool1d(1)\n", " self.cross_pooling = nn.AdaptiveAvgPool1d(1)\n", " self.dense = nn.Linear(1280, 256)\n", " self.output_layer1 = nn.Linear(2816, 128)\n", " self.output_layer = nn.Linear(128, 1)\n", "\n", " self.input_1 = nn.Linear(input_shape_1[0], 1024)\n", " self.self_attn_1 = SelfAttention(1024)\n", " self.transformer_1 = TransformerBlock(1024, 4)\n", " self.pooling_1 = nn.AdaptiveAvgPool1d(1)\n", " self.dense_1 = nn.Linear(1024, 1024)\n", " self.dropout_1 = nn.Dropout(p=0.05)\n", "\n", " self.input_2 = nn.Linear(input_shape_2[0], 1024)\n", " self.self_attn_2 = SelfAttention(1024)\n", " self.transformer_2 = TransformerBlock(1024, 4)\n", " self.pooling_2 = nn.AdaptiveAvgPool1d(1)\n", " self.dense_2 = nn.Linear(1024, 1024)\n", " self.dropout_2 = nn.Dropout(p=0.05)\n", "\n", " self.num_layers = num_layers\n", "\n", " self.convs1 = nn.ModuleList()\n", " self.convs1.append(GCNConv(4, hidden_channels))\n", " for _ in range(num_layers - 1):\n", " self.convs1.append(GCNConv(hidden_channels, hidden_channels))\n", "\n", " self.convs2 = nn.ModuleList()\n", " self.convs2.append(GCNConv(4, hidden_channels))\n", " for _ in range(num_layers - 1):\n", " self.convs2.append(GCNConv(hidden_channels, hidden_channels))\n", "\n", " self.cross_att = GATConv(hidden_channels, hidden_channels, heads=2)\n", "\n", " self.lin1 = nn.Linear(2816, hidden_channels)\n", " self.lin2 = nn.Linear(hidden_channels, 1)\n", "\n", " self.transform = NormalizeScale()\n", "\n", " def forward(self, data_batch_1, data_batch_2):\n", " x1 = data_batch_1.x.double()\n", " edge_index_1 = data_batch_1.edge_index\n", " z1 = data_batch_1.z\n", " y1_coord = data_batch_1.y_coord\n", " z1_coord = data_batch_1.z_coord\n", "\n", " concatenated_x1 = torch.stack([x1, z1, y1_coord, z1_coord], dim=1)\n", "\n", " # Apply graph normalization\n", " data_batch_1 = self.transform(data_batch_1)\n", " x1 = data_batch_1.x\n", "\n", " for i in range(self.num_layers):\n", " concatenated_x1 = self.convs1[i](concatenated_x1, edge_index_1.to(torch.int64))\n", " concatenated_x1 = F.relu(concatenated_x1.double())\n", "\n", " # Process second graph\n", " x2 = data_batch_2.x\n", " edge_index_2 = data_batch_2.edge_index\n", " z2 = data_batch_2.z\n", " y2_coord = data_batch_2.y_coord\n", " z2_coord = data_batch_2.z_coord\n", "\n", " concatenated_x2 = torch.stack([x2, z2, y2_coord, z2_coord], dim=1)\n", "\n", " # Apply graph normalization\n", " data_batch_2 = self.transform(data_batch_2)\n", " x2 = data_batch_2.x\n", "\n", " for i in range(self.num_layers):\n", " concatenated_x2 = self.convs2[i](concatenated_x2, edge_index_2.to(torch.int64))\n", " concatenated_x2 = F.relu(concatenated_x2)\n", "\n", " # Cross-attention block\n", " x1 = self.cross_att(concatenated_x1, edge_index_1.to(torch.int64))\n", " x2 = self.cross_att(concatenated_x2, edge_index_2.to(torch.int64))\n", "\n", " # Concatenate all tensors along the last dimension\n", " x = torch.cat([\n", " global_mean_pool(x1, data_batch_1.batch),\n", " global_mean_pool(x2, data_batch_2.batch)], dim=1)\n", "\n", " input_11 = self.input_1(data_batch_1.seq)\n", " self_attn_1 = self.self_attn_1(input_11)\n", " transformer_1 = self.transformer_1(self_attn_1)\n", " pooling_1 = self.pooling_1(transformer_1.transpose(0, 1)).squeeze(dim=1)\n", " dense_1 = self.dense_1(pooling_1)\n", " dropout_1 = self.dropout_1(dense_1)\n", "\n", " input_22 = self.input_2(data_batch_2.seq)\n", " self_attn_2 = self.self_attn_2(input_22)\n", " transformer_2 = self.transformer_2(self_attn_2)\n", " pooling_2 = self.pooling_2(transformer_2.transpose(0, 1)).squeeze(dim=1)\n", " dense_2 = self.dense_2(pooling_2)\n", " dropout_2 = self.dropout_2(dense_2)\n", "\n", " input_shape = [(dropout_1.shape[-1],), (dropout_2.shape[-1],)]\n", " cross_attn_1 = self.cross_attn_1([dropout_1, dropout_2])\n", " cross_attn_2 = self.cross_attn_2([self.self_attn_1(input_11), self.self_attn_2(input_22)])\n", " cross_pooling = self.cross_pooling(cross_attn_2.transpose(0, 1)).squeeze(dim=1)\n", " cross_attn = self.cross_attn([cross_attn_1, cross_pooling])\n", " cross_atten = F.tanh(self.dense(cross_attn))\n", "\n", " self_atten_1 = self.self_atten_1(self_attn_1.transpose(0, 1)).squeeze(dim=1)\n", " self_atten_2 = self.self_atten_2(self_attn_2.transpose(0, 1)).squeeze(dim=1)\n", " attention_scores = torch.cat([self_atten_1, self_atten_2, cross_atten], dim=-1)\n", "\n", " x_2 = torch.cat([self.pooling_2(x.transpose(0, 1)).squeeze(dim=1), attention_scores], dim = -1)\n", "\n", " attention_scores = torch.cat([attention_scores, self.pooling_2(x.transpose(0, 1)).squeeze(dim=1)], dim = -1)\n", "\n", " output_layer1 = F.tanh(self.output_layer1(attention_scores)) #transformer\n", " output_layer = self.output_layer(output_layer1)\n", "\n", " x = F.relu(self.lin1(x_2))\n", " x = self.lin2(x)\n", "\n", " return x, output_layer" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "cellView": "form", "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-07-30T02:30:02.157730Z", "iopub.status.busy": "2026-07-30T02:30:02.157535Z", "iopub.status.idle": "2026-07-30T02:30:02.172317Z", "shell.execute_reply": "2026-07-30T02:30:02.171636Z" }, "id": "wWCGMOsynHU8", "outputId": "dc83398f-ea31-4e0e-93ce-c0460ebd7fe5" }, "outputs": [], "source": [ "# Use the local weight/model_weights.pth file.\n" ] }, { "cell_type": "markdown", "metadata": { "id": "FtQaxzI349rF" }, "source": [ "Running on GPU (\"cuda\") device is advised" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "cellView": "form", "colab": { "base_uri": "https://localhost:8080/" }, "execution": { "iopub.execute_input": "2026-07-30T02:30:02.174211Z", "iopub.status.busy": "2026-07-30T02:30:02.174009Z", "iopub.status.idle": "2026-07-30T02:31:26.642376Z", "shell.execute_reply": "2026-07-30T02:31:26.640985Z" }, "id": "Sdn3Mq5PbxRb", "outputId": "66b69bc7-fe5b-424c-ebae-4995c9cf2011" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Inference device: cuda\n", "Predicted binding affinity for the uploaded antibody-antigen pair (in IC50): 0.11984269454374391\n" ] } ], "source": [ "# @title\n", "input_shape_1 = (int(max_antibody_sequence_length*20),)\n", "input_shape_2 = (int(max_antigen_sequence_length*20),)\n", "\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "print(\"Inference device:\", device)\n", "com_model = CombinedModel().to(device)\n", "com_model = com_model.to(torch.float64)\n", "\n", "alpha = 0.45\n", "beta = 0.55\n", "gamma = 0.05\n", "\n", "com_model_cp = torch.load(os.path.join(ROOT, 'weight', 'model_weights.pth'), map_location=device, weights_only=False)\n", "com_model_epoch = com_model_cp['epoch']\n", "com_model.load_state_dict(com_model_cp['model_state_dict'])\n", "\n", "# Evaluate the model\n", "com_model.eval()\n", "test_loss = 0.0\n", "test_mae = 0.0\n", "total_samples = 0\n", "\n", "with torch.no_grad():\n", " for input in test_loader:\n", " input_1 = input[0].to(device)\n", " input_2 = input[1].to(device)\n", " target = input[2].to(device)\n", " batch_size = input_1.size(0)\n", "\n", " if input_1.pos is not None:\n", " input_1 = NormalizeScale()(input_1)\n", " else:\n", " print(\"Data does not have position information, skipping normalization.\")\n", "\n", " if input_2.pos is not None:\n", " input_2 = NormalizeScale()(input_2)\n", " else:\n", " print(\"Data does not have position information, skipping normalization.\")\n", "\n", " output_gnn, output_tranf = com_model(input_1, input_2)\n", " print(f\"Predicted binding affinity for the uploaded antibody-antigen pair (in IC50): {10**(alpha*output_gnn.item()+beta*output_tranf.item()+gamma*np.abs(output_gnn.item() - output_tranf.item()))}\") #" ] }, { "cell_type": "markdown", "metadata": { "id": "SUyJQvdLMI3j" }, "source": [ "**Notes:**\n", "\n", "- Check that the runtime type is set to GPU at \"Runtime\" -> \"Change runtime type\".\n", "- Try to restart the session \"Runtime\" -> \"Restart session and run all\".\n", "- Check your input PDB files.\n", "- Our current model only supports proteins with amino-acid sequence lengths upto 669 (for antibodies) and 3102 (for antigens).\n", "- If you encounter any bugs, please report the issue to https://github.com/Drug-Discovery-ENTC/p2pxml/issues or email the corresponding author (Nuwan) at pmnsribandara@gmail.com" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "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.11.15" } }, "nbformat": 4, "nbformat_minor": 0 }