nbs/GT-TDC-ADME-GNN-PE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Use the Gaussian Network Model for simple Positional Encoding (PC) as a replacement for Laplacian-based PE "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/pawel/anaconda3/envs/gt/lib/python3.7/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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numpy version: 1.21.6\n",
      "Rdkit version: 2022.09.5\n",
      "Torch version: 1.13.1\n",
      "TorchMetrics version: 0.11.4\n"
     ]
    }
   ],
   "source": [
    "# Standard\n",
    "import logging\n",
    "import importlib\n",
    "\n",
    "# Third party\n",
    "import numpy as np\n",
    "import rdkit\n",
    "from rdkit import RDLogger\n",
    "import torch\n",
    "from torch import nn\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from torch_geometric.loader import DataLoader\n",
    "import torchmetrics\n",
    "from torchmetrics import MeanAbsoluteError\n",
    "\n",
    "# GT-PyG\n",
    "import gt_pyg\n",
    "from gt_pyg.data.utils import (\n",
    "    get_tensor_data, \n",
    "    get_node_dim, \n",
    "    get_edge_dim, \n",
    "    get_train_valid_test_data\n",
    ")\n",
    "from gt_pyg.nn.model import GraphTransformerNet\n",
    "\n",
    "\n",
    "# Turn off majority of RDKit warnings\n",
    "lg = RDLogger.logger()\n",
    "lg.setLevel(RDLogger.CRITICAL)\n",
    "\n",
    "\n",
    "# Set a random seed for a reproducibility purposes\n",
    "torch.manual_seed(192837465)\n",
    "\n",
    "# Setup the logger\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "# Log the used versions of RDkit and torch\n",
    "print(f'Numpy version: {np.__version__}')\n",
    "print(f'Rdkit version: {rdkit.__version__}')\n",
    "print(f'Torch version: {torch.__version__}')\n",
    "print(f'TorchMetrics version: {torchmetrics.__version__}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the ADME@TDC data\n",
    "\n",
    "**Note**: To use the code below, make sure that the chosen endpoint is a regression task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Available endpoints:\n",
      "\n",
      "1. caco2_wang\n",
      "2. hia_hou\n",
      "3. pgp_broccatelli\n",
      "4. bioavailability_ma\n",
      "5. lipophilicity_astrazeneca\n",
      "6. solubility_aqsoldb\n",
      "7. bbb_martins\n",
      "8. ppbr_az\n",
      "9. vdss_lombardo\n",
      "10. cyp2d6_veith\n",
      "11. cyp3a4_veith\n",
      "12. cyp2c9_veith\n",
      "13. cyp2d6_substrate_carbonmangels\n",
      "14. cyp3a4_substrate_carbonmangels\n",
      "15. cyp2c9_substrate_carbonmangels\n",
      "16. half_life_obach\n",
      "17. clearance_microsome_az\n",
      "18. clearance_hepatocyte_az\n",
      "19. herg\n",
      "20. ames\n",
      "21. dili\n",
      "22. ld50_zhu\n"
     ]
    }
   ],
   "source": [
    "from tdc import utils\n",
    "names = utils.retrieve_benchmark_names('ADMET_Group')\n",
    "output = \"\\n\".join([f\"{index}. {name}\" for index, name in enumerate(names, start=1)])\n",
    "print(\"Available endpoints:\\n\\n\" + output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Regression endpoints with MAE metric:\n",
    "1. caco2_wang (Best: 0.285 ± 0.005)\n",
    "2. lipophilicity_astrazeneca (Best: 0.535 ± 0.012)\n",
    "3. solubility_aqsoldb (Best: 0.776 ± 0.008)\n",
    "4. ppbr_az (Best: 9.185 ± 0.000)\n",
    "5. ld50_zhu (Best: 0.588 ± 0.005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found local copy...\n",
      "Loading...\n",
      "Done!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of training examples: 1130\n",
      "Number of validation examples: 161\n",
      "Number of test examples: 323\n"
     ]
    }
   ],
   "source": [
    "PE_DIM = 0\n",
    "(tr, va, te) = get_train_valid_test_data('ppbr_az', min_num_atoms=0)\n",
    "tr_dataset = get_tensor_data(tr.Drug.to_list(), tr.Y.to_list(), pe_dim=PE_DIM, gnn=True, pe=False)\n",
    "va_dataset = get_tensor_data(va.Drug.to_list(), va.Y.to_list(), pe_dim=PE_DIM, gnn=True, pe=False)\n",
    "te_dataset = get_tensor_data(te.Drug.to_list(), te.Y.to_list(), pe_dim=PE_DIM, gnn=True, pe=False)\n",
    "NODE_DIM = get_node_dim()\n",
    "EDGE_DIM = get_edge_dim()\n",
    "\n",
    "print(f'Number of training examples: {len(tr_dataset)}')\n",
    "print(f'Number of validation examples: {len(va_dataset)}')\n",
    "print(f'Number of test examples: {len(te_dataset)}')\n",
    "\n",
    "train_loader = DataLoader(tr_dataset, batch_size=64)\n",
    "val_loader = DataLoader(va_dataset, batch_size=512)\n",
    "test_loader = DataLoader(te_dataset, batch_size=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and eval the GT model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch, loss_func):\n",
    "    model.train()\n",
    "    train_mae = MeanAbsoluteError()\n",
    "\n",
    "    for data in train_loader:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # randomly flip sign of eigenvectors\n",
    "        batch_pe = None\n",
    "        (out,_) = model(data.x, data.edge_index, data.edge_attr, batch_pe, data.batch, zero_var=False)\n",
    "        loss = loss_func(out.squeeze(), data.y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        train_mae.update(out.squeeze(), data.y)\n",
    "\n",
    "    return train_mae.compute()\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def test(loader):\n",
    "    model.eval()\n",
    "    test_mae = MeanAbsoluteError()\n",
    "\n",
    "    total_error = 0\n",
    "    for data in loader:\n",
    "        data = data.to(device)\n",
    "       \n",
    "        # randomly flip sign of eigenvectors\n",
    "        batch_pe = None\n",
    "        (out,_) = model(data.x, data.edge_index, data.edge_attr, batch_pe, data.batch)\n",
    "        \n",
    "        test_mae.update(out.squeeze(), data.y)\n",
    "        \n",
    "    return test_mae.compute()\n",
    "\n",
    "train_loss = nn.L1Loss(reduction='mean')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard Graph Transformer setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GraphTransformerNet(\n",
      "  (node_emb): Linear(in_features=76, out_features=128, bias=False)\n",
      "  (edge_emb): Linear(in_features=10, out_features=128, bias=False)\n",
      "  (gt_layers): ModuleList(\n",
      "    (0): GTConv(128, 128, heads=8, aggrs: sum,mean, qkv_bias: False, gate: True)\n",
      "    (1): GTConv(128, 128, heads=8, aggrs: sum,mean, qkv_bias: False, gate: True)\n",
      "    (2): GTConv(128, 128, heads=8, aggrs: sum,mean, qkv_bias: False, gate: True)\n",
      "    (3): GTConv(128, 128, heads=8, aggrs: sum,mean, qkv_bias: False, gate: True)\n",
      "  )\n",
      "  (global_pool): MultiAggregation([\n",
      "    SumAggregation(),\n",
      "    MeanAggregation(),\n",
      "    MinAggregation(),\n",
      "    MaxAggregation(),\n",
      "  ], mode=cat)\n",
      "  (mu_mlp): MLP(\n",
      "    (mlp): Sequential(\n",
      "      (0): Linear(in_features=512, out_features=128, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=128, out_features=1, bias=True)\n",
      "    )\n",
      "  )\n",
      "  (log_var_mlp): MLP(\n",
      "    (mlp): Sequential(\n",
      "      (0): Linear(in_features=512, out_features=128, bias=True)\n",
      "      (1): ReLU()\n",
      "      (2): Linear(in_features=128, out_features=1, bias=True)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "Number of params: 1003 k\n",
      "Epoch: 01, Loss: 67.9054, Val: 51.2259, Test: 42.9332\n",
      "Epoch: 02, Loss: 28.9047, Val: 23.3430, Test: 23.5090\n",
      "Epoch: 03, Loss: 15.0898, Val: 26.3209, Test: 26.0182\n",
      "Epoch: 04, Loss: 12.4090, Val: 18.6056, Test: 17.2357\n",
      "Epoch: 05, Loss: 12.0454, Val: 17.3376, Test: 16.6983\n",
      "Epoch: 06, Loss: 11.9575, Val: 21.5594, Test: 21.1579\n",
      "Epoch: 07, Loss: 11.4733, Val: 14.8946, Test: 15.3955\n",
      "Epoch: 08, Loss: 10.9486, Val: 12.7779, Test: 13.7182\n",
      "Epoch: 09, Loss: 11.3676, Val: 16.3640, Test: 16.5855\n",
      "Epoch: 10, Loss: 11.0898, Val: 25.0348, Test: 24.4669\n",
      "Epoch: 11, Loss: 11.0134, Val: 14.0643, Test: 14.5464\n",
      "Epoch: 12, Loss: 10.8510, Val: 16.2039, Test: 16.9654\n",
      "Epoch: 13, Loss: 10.9855, Val: 26.5054, Test: 26.3783\n",
      "Epoch: 14, Loss: 10.7232, Val: 12.2520, Test: 13.4538\n",
      "Epoch: 15, Loss: 10.8379, Val: 12.1771, Test: 12.7884\n",
      "Epoch: 16, Loss: 10.8779, Val: 10.4116, Test: 11.7527\n",
      "Epoch: 17, Loss: 10.8205, Val: 22.0192, Test: 22.0707\n",
      "Epoch: 18, Loss: 10.4186, Val: 12.3866, Test: 13.6080\n",
      "Epoch: 19, Loss: 10.2508, Val: 12.6924, Test: 11.5388\n",
      "Epoch: 20, Loss: 10.2397, Val: 15.8327, Test: 16.9361\n",
      "Epoch: 21, Loss: 10.2107, Val: 8.9558, Test: 9.8679\n",
      "Epoch: 22, Loss: 10.2598, Val: 36.9397, Test: 35.6801\n",
      "Epoch: 23, Loss: 10.3430, Val: 9.3078, Test: 9.5391\n",
      "Epoch: 24, Loss: 10.0940, Val: 20.8000, Test: 20.7961\n",
      "Epoch: 25, Loss: 9.8928, Val: 18.7864, Test: 17.4517\n",
      "Epoch: 26, Loss: 9.7673, Val: 32.5481, Test: 32.1844\n",
      "Epoch: 27, Loss: 9.4030, Val: 16.5057, Test: 17.0697\n",
      "Epoch: 28, Loss: 9.2019, Val: 9.1004, Test: 10.6287\n",
      "Epoch: 29, Loss: 8.9756, Val: 11.4093, Test: 12.0364\n",
      "Epoch: 30, Loss: 8.8116, Val: 9.9293, Test: 11.3834\n",
      "Epoch: 31, Loss: 8.6773, Val: 8.8113, Test: 10.1029\n",
      "Epoch: 32, Loss: 8.7540, Val: 9.6385, Test: 10.8604\n",
      "Epoch: 33, Loss: 8.5177, Val: 9.4990, Test: 10.3666\n",
      "Epoch: 34, Loss: 8.4901, Val: 12.4248, Test: 12.6609\n",
      "Epoch: 35, Loss: 8.4158, Val: 8.4130, Test: 9.4064\n",
      "Epoch: 36, Loss: 8.1959, Val: 9.8287, Test: 11.0842\n",
      "Epoch: 37, Loss: 8.2287, Val: 8.9260, Test: 9.7113\n",
      "Epoch: 38, Loss: 8.1167, Val: 11.6102, Test: 10.3065\n",
      "Epoch: 39, Loss: 8.1587, Val: 11.9287, Test: 12.7901\n",
      "Epoch: 40, Loss: 8.1090, Val: 9.4885, Test: 10.0828\n",
      "Epoch: 41, Loss: 8.0388, Val: 9.7082, Test: 10.8527\n",
      "Epoch: 42, Loss: 8.0383, Val: 9.7585, Test: 9.8192\n",
      "Epoch: 43, Loss: 7.8738, Val: 9.9404, Test: 11.0548\n",
      "Epoch: 44, Loss: 7.7051, Val: 13.2197, Test: 13.1937\n",
      "Epoch: 45, Loss: 7.6354, Val: 9.6716, Test: 10.3944\n",
      "Epoch: 46, Loss: 7.3542, Val: 17.2025, Test: 17.1773\n",
      "Epoch: 47, Loss: 7.2231, Val: 7.5942, Test: 8.5401\n",
      "Epoch: 48, Loss: 7.1903, Val: 15.4758, Test: 15.2600\n",
      "Epoch: 49, Loss: 7.0789, Val: 7.5622, Test: 8.5897\n",
      "Epoch: 50, Loss: 7.0212, Val: 9.4274, Test: 9.8206\n",
      "\n",
      "Model's performance on the test set\n",
      "===================================\n",
      "MAE=8.589730262756348\n",
      "Epoch=49\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = GraphTransformerNet(node_dim_in=NODE_DIM,\n",
    "                            edge_dim_in=EDGE_DIM,\n",
    "                            pe_in_dim=PE_DIM,\n",
    "                            num_gt_layers=4, \n",
    "                            hidden_dim=128,\n",
    "                            num_heads=8,\n",
    "                            norm='bn',\n",
    "                            gate=True,\n",
    "                            qkv_bias=False,\n",
    "                            gt_aggregators=['sum', 'mean'],\n",
    "                            aggregators=['sum', 'mean', 'min', 'max'],\n",
    "                            dropout=0.1,\n",
    "                            act='relu').to(device)\n",
    "\n",
    "if int(torch.__version__.split('.')[0]) >= 2:\n",
    "    model = torch_geometric.compile(model) \n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\n",
    "scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5,\n",
    "                              min_lr=0.00001)\n",
    "\n",
    "print(model)\n",
    "print(f\"Number of params: {model.num_parameters()//1000} k\")\n",
    "\n",
    "best_epoch = 0\n",
    "best_validation_loss = np.inf\n",
    "test_set_mae = np.inf\n",
    "for epoch in range(1, 51):\n",
    "    tr_loss = train(epoch, loss_func=train_loss)\n",
    "    va_loss = test(val_loader)\n",
    "    te_loss = test(test_loader)\n",
    "    scheduler.step(va_loss)\n",
    "    print(f'Epoch: {epoch:02d}, Loss: {tr_loss:.4f}, Val: {va_loss:.4f}, '\n",
    "          f'Test: {te_loss:.4f}')\n",
    "    if va_loss < best_validation_loss:\n",
    "        best_epoch = epoch\n",
    "        best_validation_loss = va_loss\n",
    "        test_set_mae = te_loss\n",
    "        \n",
    "print(\"\\nModel's performance on the test set\\n\"\n",
    "        \"===================================\\n\"\n",
    "        f'MAE={test_set_mae}\\n'\n",
    "        f'Epoch={best_epoch}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}