molecular_property_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "59qfiDHHeyin"
   },
   "source": [
    "# Molecular Toxicity Prediction\n",
    "This project focuses on predicting the toxicity of molecules by utilizing graph classification techniques in PyTorch. By representing molecules as graphs, where nodes correspond to atoms and edges to chemical bonds, we can analyze their atomic composition and arrangement to determine their biological behaviors and properties. Using high-throughput experimental data provided in three files (train_data.pt, valid_data.pt, and test_data.pt), we aim to classify molecular properties and interactions, including activation or inhibition, binding affinity, and dose-response relationships. This notebook establishes a baseline for molecular toxicity prediction, leveraging PyTorch 2.1.0 for implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JTCASsyypP4K"
   },
   "source": [
    "# Install package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "G1peiPgqo5IX",
    "outputId": "e7e622b4-bf1c-4257-d174-0bdcb90cdee5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch_geometric in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (2.4.0)\n",
      "Requirement already satisfied: tqdm in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (4.66.1)\n",
      "Requirement already satisfied: numpy in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (1.25.2)\n",
      "Requirement already satisfied: scipy in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (1.11.3)\n",
      "Requirement already satisfied: jinja2 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (3.1.2)\n",
      "Requirement already satisfied: requests in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (2.31.0)\n",
      "Requirement already satisfied: pyparsing in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (3.0.4)\n",
      "Requirement already satisfied: scikit-learn in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (1.3.2)\n",
      "Requirement already satisfied: psutil>=5.8.0 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from torch_geometric) (5.9.5)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from jinja2->torch_geometric) (2.1.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from requests->torch_geometric) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from requests->torch_geometric) (3.3)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from requests->torch_geometric) (1.26.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from requests->torch_geometric) (2022.12.7)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from scikit-learn->torch_geometric) (1.3.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from scikit-learn->torch_geometric) (3.2.0)\n",
      "Requirement already satisfied: rdkit-pypi in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (2022.9.5)\n",
      "Requirement already satisfied: numpy in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from rdkit-pypi) (1.25.2)\n",
      "Requirement already satisfied: Pillow in /Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages (from rdkit-pypi) (10.0.1)\n"
     ]
    }
   ],
   "source": [
    "# New these two packages\n",
    "!pip install torch_geometric\n",
    "!pip install rdkit-pypi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cQLM3rMTpu6r"
   },
   "source": [
    "# Some tutorials.\n",
    "\n",
    "\n",
    "\n",
    "1.   Pytorch geometric package: https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html\n",
    "2.   PyTorch Geometric for Graph-Based Molecular Property Prediction using MoleculeNet benchmark: https://medium.com/@nikopavl4/pytorch-geometric-for-graph-based-molecular-property-prediction-using-moleculenet-benchmark-41e36369d3c6\n",
    "3. Graph neural networks for graph classification. https://colab.research.google.com/drive/1I8a0DfQ3fI7Njc62__mVXUlcAleUclnb?usp=sharing\n",
    "4. Related github repository on molecular property predictions. https://github.com/yifeiwang15/MotifConv/tree/main/MCM_for_molecule_benchmarks\n",
    "\n",
    "\n",
    "## What are node and edge features in a molecule.\n",
    "\n",
    "### Node features:\n",
    "\n",
    "**Atomic number**: Number of protons in the nucleus of an atom. It’s characteristic of a chemical element and determines its place in the periodic table.\n",
    "\n",
    "**Chirality**: A molecule is chiral if it is distinguishable from its mirror image by any combination of rotations, translations, and some conformational changes. Different types of chirality exist depending on the molecule and the arrangement of the atoms.\n",
    "\n",
    "**Degree**: Number of directly-bonded neighbors of the atom.\n",
    "Formal charge: Charge assigned to an atom. It reflects the electron count associated with the atom compared to the isolated neutral atom.\n",
    "\n",
    "**Number of H**: Total number of hydrogen atoms on the atom.\n",
    "Number of radical e: Number of unpaired electrons of the atom.\n",
    "\n",
    "**Hybridization**: Atom’s hybridization.\n",
    "\n",
    "**Is aromatic**: Whether it is included in a cyclic structure with pi bonds. This type of structure tends to be very stable in comparison with other geometric arrangements of the same atoms.\n",
    "\n",
    "**Is in ring**: Whether it is included in a ring (a simple cycle of atoms and bonds in a molecule).\n",
    "\n",
    "### Edge features:\n",
    "\n",
    "**Bond type:**: Whether the bond is single, double, triple, or aromatic.\n",
    "\n",
    "**Stereo Type:** Stereo configuration of the bond.\n",
    "\n",
    "**Is conjugated**: Whether or not the bond is considered to be conjugated.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0054Ib-4Vfaj"
   },
   "source": [
    "# Dataset preparation and train-valid splitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "atIc86zFnj0c"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/defnecoban/opt/miniconda3/lib/python3.9/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": [
    "import torch\n",
    "import torch_geometric\n",
    "import numpy as np\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.datasets import MoleculeNet\n",
    "import pickle\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wqVa56ajntas",
    "outputId": "6a4b393f-2369-4aa7-9a04-40edc8ac3bfd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of training set: 6264\n",
      "Size of validation set: 783\n",
      "Size of test set: 784\n"
     ]
    }
   ],
   "source": [
    "# Load datasets. The training and validation sets contain both molecules and their property labels. The test set only contain molecules.\n",
    "# There are 12 property tasks for prediction. Some properties labels are missing (i.e., nan). You can ignore them.\n",
    "train_dataset = torch.load(\"train_data.pt\")\n",
    "valid_dataset = torch.load(\"valid_data.pt\")\n",
    "test_dataset = torch.load(\"test_data.pt\")\n",
    "\n",
    "print(f'Size of training set: {len(train_dataset)}')\n",
    "print(f'Size of validation set: {len(valid_dataset)}')\n",
    "print(f'Size of test set: {len(test_dataset)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dQRRylPsVYKn",
    "outputId": "8230bbf8-db6d-4bf5-fd28-cd754afb2679"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data(x=[11, 9], edge_index=[2, 20], edge_attr=[20, 3], y=[1, 12], smiles='CC(O)(P(=O)(O)O)P(=O)(O)O')\n",
      "Get node feature matrix:\n",
      "tensor([[ 6,  0,  4,  5,  3,  0,  4,  0,  0],\n",
      "        [ 6,  0,  4,  5,  0,  0,  4,  0,  0],\n",
      "        [ 8,  0,  2,  5,  1,  0,  4,  0,  0],\n",
      "        [15,  0,  4,  5,  0,  0,  4,  0,  0],\n",
      "        [ 8,  0,  1,  5,  0,  0,  3,  0,  0],\n",
      "        [ 8,  0,  2,  5,  1,  0,  4,  0,  0],\n",
      "        [ 8,  0,  2,  5,  1,  0,  4,  0,  0],\n",
      "        [15,  0,  4,  5,  0,  0,  4,  0,  0],\n",
      "        [ 8,  0,  1,  5,  0,  0,  3,  0,  0],\n",
      "        [ 8,  0,  2,  5,  1,  0,  4,  0,  0],\n",
      "        [ 8,  0,  2,  5,  1,  0,  4,  0,  0]])\n",
      "torch.Size([11, 9])\n",
      "Get edge index matrix:\n",
      "tensor([[ 0,  1,  1,  1,  1,  2,  3,  3,  3,  3,  4,  5,  6,  7,  7,  7,  7,  8,\n",
      "          9, 10],\n",
      "        [ 1,  0,  2,  3,  7,  1,  1,  4,  5,  6,  3,  3,  3,  1,  8,  9, 10,  7,\n",
      "          7,  7]])\n",
      "torch.Size([2, 20])\n",
      "Get edge attribute matrix:\n",
      "tensor([[1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [2, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [2, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [2, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [2, 0, 0],\n",
      "        [1, 0, 0],\n",
      "        [1, 0, 0]])\n",
      "torch.Size([20, 3])\n",
      "Get molecular property labels:\n",
      "tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])\n",
      "torch.Size([1, 12])\n"
     ]
    }
   ],
   "source": [
    "# one graph example\n",
    "g = train_dataset[0]\n",
    "print(g)\n",
    "\n",
    "print(\"Get node feature matrix:\")\n",
    "print(g.x)\n",
    "print(g.x.shape) # (num_of_nodes, num_of_node_features)\n",
    "\n",
    "print(\"Get edge index matrix:\")\n",
    "print(g.edge_index)\n",
    "print(g.edge_index.shape) # (2, num_of_edges)\n",
    "\n",
    "print(\"Get edge attribute matrix:\")\n",
    "print(g.edge_attr)\n",
    "print(g.edge_attr.shape) # (num_of_edges, num_of_edge_features)\n",
    "\n",
    "print(\"Get molecular property labels:\")\n",
    "print(g.y)\n",
    "print(g.y.shape) # (1, 12)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cilgcBXdCNS3"
   },
   "source": [
    "As we can observe, we have 11 nodes (rows) and each node has 9 features (columns). However, the features provided by Moleculenet are discrete and of type long, so we need to convert them first to continuous embeddings in order to feed them in any ML model.\n",
    "\n",
    "For example, the first column indicates the atomic number of a node, where 1 represents Hydrogen, 6 represents Carbon, 8 for Oxygen, according to periodic table of elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AZLIUrqobxY3",
    "outputId": "aafde21f-b0ec-4df3-99f5-8f25b2966f42"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages/torch_geometric/deprecation.py:22: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead\n",
      "  warnings.warn(out)\n"
     ]
    }
   ],
   "source": [
    "# Example of preparing data loaders.\n",
    "# You can use any batch size and see what happens in model performance.\n",
    "\n",
    "from torch_geometric.data import DataLoader\n",
    "\n",
    "batch_size=32\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Wdzj4S9Jds3m",
    "outputId": "9d22b922-9c93-43a6-9263-1ebf2579f701"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DataBatch(x=[563, 9], edge_index=[2, 1156], edge_attr=[1156, 3], y=[32, 12], smiles=[32], batch=[563], ptr=[33])\n"
     ]
    }
   ],
   "source": [
    "# Example of creating one mini-batch\n",
    "# See more info about mini-batch in pytorch geometric in https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html\n",
    "batch = next(iter(train_loader))\n",
    "print(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mc2_Mytlhn5P"
   },
   "source": [
    "# Visualization of molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 167
    },
    "id": "BNoZJPN0ic5S",
    "outputId": "cfa92541-b0ef-4029-f4da-67d0b58ea952"
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAcIAAACWCAIAAADCEh9HAAAABmJLR0QA/wD/AP+gvaeTAAAVc0lEQVR4nO3de1hUBf7H8fdwVS6CFwQEVLK8oFlp3tZsu5Bum+ua/XqSUnd9tov+Vn+2XtJMS0ksNS3Ntq3d7Wk1N2l7tr2ktatWu2qlWd4AM0kRzQTEREERYc7vDyYROjOgZy4wfF5Pf4xnvpzz6XD4cGbmzGAzDAMREblSAb4OICLStKlGRUQsUY2KiFiiGhURsUQ1KiJiSZCvA4i3VFRQVMTp00RFERNDcLCvA/mjCxcoKqKkhFataN9eO7mZ0NloM7BpE8OH06YNiYmkpJCQQNu2jBzJli11J/ftIyCAgADy881X9eGHjoGKCk+nbmK2bOHuu2nbloQEUlJITKR1a4YPZ9OmupPffOPYh1lZ5qvascMxcOKEp1OLW6hG/VpVFQ8/TGoq69Zx/jw33EBqKr16UVbG3//OkCFMm4bdXjNvGDX/map3oBmy25k6lSFD+NvfOHuW3r1JTeW66ygvZ906UlN55BGqqmrmtZP9jmrUr02fzu9/j83GY49RWMgXX7BhA3v3cuwYDz0EsGwZ6em+TtnEzZ/P888DTJ7Mt9+yezcbNrBrFwUFTJkC8OqrzJjh24ziWYb4q23bDJvNAGPJEvOBadMMMIKCjH37HEuysw0wwMjLM/+STZscA+XlHsnc5GRnG4GBBhhz55oPzJ9vgGGzGZ995lhy5IhjH+7ZY/4l27c7BgoLPZJZ3E1no/5rxQoMg169mDrVfCAjg8REKitZudK7yfzIypVUVZGczJNPmg/Mnk3XrhgGK1Z4N5l4j2rUTxkG778P8MADBDj5LoeGkpYG8N573gvmZ6p33ejRBDm56CUoiDFjAN5/X891+ivVqJ/Ky6O4GKBfP1djAwYAHDzId995I5WfKS4mLw+gf39XY9U7uajI6fUP0sTpulE/VVjouNGhg6uxhATHjaIiWreuWT5smPk1j2VlbknnJ65gJ3fqVLP8nnsIDTWZP3fOLenEa1Sjfqq83HGjRQtXY2Fhjht1+nH/fg9kuhKjR48G1q5d6+sgZi72neud3LKl40adnXzggAcyiQ+oRv1UZKTjhuvzxzNnHDeiomot/+ILEhNN5rdsYdQo6+kaLjMzk0Zbo61aOW643smlpY4bdXbyf/5Djx4m87t2MXSo9XTiNapRPxUf77iRl0evXk7Hqp/aCwykfftay9u0ISbGZL5OETRzsbEEBGC3k5fHoEFOx6p3Mpd8U6q1bm2+k6Oj3ZZQvEIvMfmp+HjH6eQnn7ga+/hjgJQUIiK8kcrPREbSvTs0bCd37kxsrDdSidepRv3XiBEAa9Y4ff/7mTNkZgKMHOm9VH6meidnZjp9XF9ezpo1oJ3sz1Sj/mvyZIKDOXyYOXNM7jUMHn2U4mLCw5kwwevh/MXEibRsSWEh06ebD8yaxdGjhIYyaZJ3k4n3qEb9V/fuPPUUwJIljB/P4cM1dx04wL338tprAMuW1XO9jrjQsSOLFgH87neMHk1ubs1deXmMHcvy5QALFtCli28SiufpJSa/Nns2FRUsWMDrr/P663TuTNu2FBZy5AhAUBDPPcfDD/s6ZRM3eTLl5cyeTWYmmZl07EhMTM2V+YGBzJ/v9FxV/IJq1K/ZbMyfz6hRLFvG+++Tl+f42e7QgZ/+lOnT6dat1nxICMnJgNPPG27Z0jFgs3kyd1MzYwbDh/Pcc7z3Hvn5jncrxcZy551Mm1b3SomgIMc+DAkxX1uLFo6BwECPphZ3sRl6n2/zUVbGyZO0bVtz1X2jZ7PZgKZ0lJ49S3ExbdoQHu7rKOIlqlFp1JpejUrzo5eYmpOFC4mL48UXfZ3Dr734InFxLFzo6xziParR5uTMGQoKat6bKJ5QWkpBQc27bKUZUI2KiFiiGhURsUQ1KiJiiWpURMQS1aiIiCWqURERS1SjIiKWqEZFRCxRjYqIWKIaFRGxRDUqImKJalRExBLVqIiIJapRERFLVKMiIpaoRkVELFGNiohYohptTgZGMaMvfSJ8ncOv9YlgRl8GRvk6h3iParQ5iSgh/HPC9UdEPCm8lPDPiSjxdQ7xHtWoiIglqlEREUtUoyIilqhGRUQsUY2KiFiiGhURsUQ1KiJiiWpURMQS1aiIiCWqURERS1SjIiKWqEZFRCxRjYqIWKIaFRGxRDUqImKJalRExBLVqIiIJUG+DiBe1GUYIZF0vsXXOfxa51u4LYOkH/k6h3iPzTAMX2cQccpmswE6SqUx09no5SspYcMGPvuMwkICAoiLY9AgbruNsDBfJzOzOYPzZ2otiYilQz+SBmOz+SiT3/n4Oc6eqLUkrB0dbqTTzdga8fNmZ8/ywQd88gnHj2O30749/fuTmkqU/h7f5dHZ6OWw21m8mGee4fTpune1b09GBg8+6ItYLi2Np/S4yfL2vRj1BrHXeT3Q5WkaZ6MvduXkAZPlba5h1GoSBng9UAO8+ipz51JYWHd5VBSPP85jj+m3bMOpRhvMMBg3jjfeABg8mLQ0rrqKqir27eNPfyI7G2DmTJ591rcx66qu0X7/63hK9Nx3FOxh52tUniOsHRN2E9nBxwldako1ev14rrkToLyEomy++AMVpYS24pGdtL7K1xFrmzWLRYsArruOsWPp3h2bjdxc/vxntm0DGDOGVavUpA2kGm2wlSuZPBmbjRUrmDSp1l1VVUyfzgsvALzzDiNHAnz2GX/845VvrlUYN5+1EBduz6BlW0eN/vw1rh9fc1fBbn4/gKrzDH6M1EWWtuJhTalGf7KcAf9Xs7D4K17tS0UpN07krt9yrphNT1jayn/DOG3hkPjVr+jXD+Cddxg1CmD2bJ5+moBLnnYwDBYvZtYsgJUr+fWvreRtPvTcaMOUl5OeDvDII3U7FAgMZNkysrLYuJEnnuDnP8dm48ABXnnlyrfYIZ7wb6/8y4GbZtKyrfldsdfR9S72/ZX8LZY2IS607UrK/7DrdcdOPn+azy0cD8CaeI5ZOCRuvpl+/TAMnngC4Gc/IyOj7ozNxsyZ5OSwahXp6Tz4IKGhV77FZkM12jDvvktREQEBzJxpPmCzMWcOGzeSk8P27QwYQP/+vPzylW+xZTA9L1z5lwNh7VzdG5kAUF5iaRPiWvVOPl8CENaOuywcD0DnYM5ZOCT69wfYto19+wAef9zp5Jw5rF5NYSHr1jnOW8Ul1WjDbNkC0KMHnTs7nbn5ZqKiKClh61YGDODqq7n6am/lu3xFOQARcb7O4dccOzkeICSSGydYWtuNbkjkOJKjoxk40OnMNdfQtSv797N1q2q0IVSjDfP11wDdu7uasdno3p1t28jN9U6oK5e1lkMfAHT7GUBVBaePumvdVYbt8Ck3P5V58OBBd62qU7Qt0Oa+eK0SCQwxv+urf7L/7wBdh7ttc9ZVH5w9etTz8lFKCvv3c8Ds8gP5AdVow5w6BRAdXc9Y69YA333n8TyXK+dtTuwHKD9FwW6OfgoQ34e+jwCc+JLfue3KJ1vLdl1mnqh/7nJ06dLFXauqWtSOc+6LN2E3sb0dt7/6J2eOAZw/TWEW+ZsBYlIYOMVtm7OupASo/8rQRnskN0qq0YapfjXTbq9nrKoKIDDQ43ku14H1HFhf88/AUHqPYegSgloABIUSneyuTdlDopOTI921tkOHDgHJye6LFxkdEOq2eARd8grMwY0c3Fjzz4Bgeo1m2FJC3Lc56xp4JFcPNMIjuVFSjTZM9Xlovb+cT56E73+TNyrX/cLxLu8W0UTEEncDoa1q7m3bjSlue9QcBAcnumtljgue3Pig3oN63kfybQAtogiLIb4PLep7+OJ9Tf1IbpRUow3Towf/+AdZWa5mKisdr4GmpHgn1GXo/ONa142KJyT9iL4P+zpEfXr0AMjJoarK1cnm3r0APXt6KVUTpxptmJtuYtEicnP58kunLzRt2sTZswBDhgBs2cLy5Ve+xegIhpZe+ZcDP32J8PaW1iBuVFbIemtXs/87glMWDokpU7jpJm66CaCsjI8+4vbbzSezsjh0CL4/kqU+qtGG+clPSEjgm29YuJBVq0wGDMNxMfPAgfTqBZCfz9tvX/kWO8STaO3y+zsWW/pyca8LZeRYOB6A9dYuv7/7boDrr6dPH774goULue0289frFywA6NyZ1NQr31xzohptmKAgFixg/HjeeIOePet+cENlJY8+yubNBATUvKd+yBDeeuvKtxgaxNWVljKHx1r6cnGv8FjutXA8ADcEcd7CIXHxQtFnn2XYMD74gKlTWbKEoEtKoPpsIDMTYOFCvcTUQKrRBvvlL/n0U155hVmz+Otfuf9+rr4au93xzrmcHGw2nnmGH//YMZ+URFKSTxNLYxIcRsq9ltbgrqfc77iD9HTmzuWFF9iwgXHj6NGDgAAOHGDNGnbsAJg8mbQ0N23P/6lGL8fLL5OSwvz5bN/O9u217kpKYskS7rvPR8lELsecOSQnM3Mm2dl139/crh3p6Ux038UWzYA+4enylZaycSPbtnH8OAEBJCQweDC33NJIP8Th8H+pqiAmpZF/IJ4zTeMTno58zIWztOtGqyb1+KO8nI8+YutWjh3DbicujoEDSU0lPNzXyZoY1ag0ak2jRqV5a8R/4UDcbucfWXU7WW/6Oodfy3qTVbez08JHzUpToxptTk7mcugDTuX5OodfO5XHoQ842eg/nkbcRzUqImKJalRExBLVqIiIJapRERFLVKMiIpaoRkVELFGNiohYohoVEbFENSoiYolqVETEEtWoiIglqlEREUv0sc3SqO2v/hNsIo2YalQata5btvg6gkg99KBeRMQS1aiIiCWqURERS1SjzUlxGwqGcDLK1zn82skoCoZQ3MbXOcR7VKPNyc6TvLyZ7BJf5/Br2SW8vJmdJ32dQ7xHNSoiYolqVETEEtWoiIglqlEREUtUoyIilqhGRUQsUY2KiFiiGhURsUQ1KiJiiWpURMQS1aiIiCWqURERS1SjIiKWqEZFRCxRjYqIWKIaFRGxRDUqImKJ/sBycxIZSWwsERG+znE5YmN9neAyRUQQG0tkpK9ziPfYDMPwdQYRkSZMD+qbk9JSjhyhrMzXORqsqoqCAgoKqKrydZQGKyvjyBFKS32dQ7xHNdoMbNtGWhoxMURG0rEjERHExTFuHLt21Z3cv5/oaKKjOXLEfFWbNzsGKio8GPjcOZYuZcAAQkOJiyMujtBQBg5k6VLOnas7PGEC0dE89JDTtd19N9HRPPGEBwMDu3YxbhxxcURE0LEjkZHExJCWxvbtdSePHXPsw5wc81Xt3OkYKC72bGZxE9WoX7PbmTaNQYNYu5biYrp0oW9fOnWioIDVq+nbl3nzas1XVVFSQkkJdrv5Ci9ccAx47rmgL7+kZ0+mT2f7dsLD6d2ba68lLIxt25g+nV69+PLLWvNlZZSUuDrFLi2lpMSkf93oqafo25fVqykoIDmZvn1JTubECdauZeBApk+vtbvsdsc+dHaKXVlZz3dBGhnVqF+bO5dlyzAMJkzgyBFyc9mxg7w8cnO57z7sdubPZ9EiX6e8RGEht97KoUN07Mjbb3PiBLt3s2ePo5Li4zl4kFtvpbDQ10EvsWgR6enY7fziFxw6xMGD7NjBwYMcPsz48RgGS5cyZ46vU4onGeKvdu40AgIMMObNMx+YMMEAIyTEOHDAsSQ72wADjLw88y/ZtMkxUF7ukcxpaQYYsbFGfr7Jvbm5Rps2BhgPPFCzcMwYA4y0NKfrTE01wPjNb9yf1jCMr74ygoMNMKZONR+YMcMAIzDQ2LXLseTIEcc+3LPH/Eu2b3cMFBZ6JLO4m85G/dfy5djtdOvm9FRo6VLi4qio4KWXvJvMiaNH+ctfADIySEoyGejShfR0gLfe4ptvvJrNmZUruXCBpCQWLjQfWLCA5GSqqlixwrvJxHtUo/5r/XqAMWMIDDQfCAsjLQ1g3TrvpXLhX/+ispIWLRg92unM2LGEhHDhAv/+txeTOVe960aPJjTUfCAkhLFj4ftvh/gj1aifOnzY8QTigAGuxgYNAsjN5dQpb6RybccOgJ49CQ93OtOqFT171gz71nff8fXX0LCdfPy40+sfpInTu5j8VEGB40ZCgquxxEQAw6CoiOjomuX33GN+elVS4q6AJqozV0dyISGBnTtr/gerbdjA4MHm89nZ7ghn5mIG15kv3ltYWOvJinHjCAszmdc1p02NatRPnT3ruNGypauxiz/GdX50P//cA5nqU525RYt6xqrPVesEPnGCEyc8E8s5izv5h9ftStOkGvVTF9/TffFH3dTFH+yoqFrLN2+mQweT+U8+YcwY6+nMVWd2HZjvM9cJPGIEzz9vPj9uHFu3Wk9nwuJOXr+ebt1M5vfuZeRI6+nEa1SjfuriJ3rk5zueTDSVnw8QEEBMTK3lSUl06mQyn5fnpnxmqjPX+wRidea4uFoLw8O56irzedenila0b4/NhmGQn8/AgU7HqgPzg49ZSUw0z6w3LzU1eonJTyUmEh8P8Omnrsaq7+3evVF8ItGNNwLk5HD6tNOZU6fYtw+gf38vpXIhKoquXaFhOzkpyfEdEb+jGvVfw4cDrFlDZaX5wNmzZGYCjBjhvVQuDBtGcDAVFbz5ptOZ1auprCQ0lKFDvZjMueqdnJlJebn5QEUFb7wBjWYniweoRv3X5MkEBvL112RkmA/MmkVBAS1aMHGid5M5ER/PffcBPPkkR4+aDBw+zPz5AGPG1H0WwlcmTiQkhGPHnL7HIT2dw4cJDmbSJO8mE+9Rjfqva69l1iyAefOYMqXWC9nffsv48bz4IsAzz9Cxo28S/tCyZcTGUljIkCGsX1/z2Rx2O+++y5AhFBeTkMDixT5NeYkuXXj6aYClS3noIY4fr7mrqIhJkxy/w+bMoXt33yQUz9NLTH4tPZ2yMl54gRUreOklevYkJobjx9m3D7udgADmzePRR32d8hIxMXz4IXfeSV4ed91F+/Zccw3AV19RVARw1VW89x5t2vg2Zi0zZnDmDBkZ/OEPvPYaKSmO3wQ5OVRVYbMxYwZz5/o6pXiQatSvBQTw/PPccw9Ll7JxI3v2OJZHRzN0KDNn0qdPrfmgINq1A5y+fzQkxDFgs3kqc48eZGezfDlr17J3b82HOfXuzf33M3ly3UvWW7WiXTtatXK6wuho2rXz4J9Osdl4+mlGjGDxYjZsICuLrCyAyEjuuIMZM+q+iB8Y6NiHQU5++oKDHQMBerDYNOiPiDQbFy5QVERJCdHRxMQ4/RluVM6ccTxMjotrFNcS1KuykqIiTp0iKoqYGIKDfR1IvEE1KiJiiR41iIhYohoVEbFENSoiYolqVETEEtWoiIgl/w9FfvijJS+84AAAAK16VFh0cmRraXRQS0wgcmRraXQgMjAyMi4wOS41AAB4nHu/b+09BiDgZYAAJiDmBmIuIG5gZGNIANKMzGwMCkCahQPCZeQHc1k5GDSAFDMTTJhIaW4GRgZGJqCpDMwsHEzMrAxA4xnZGdg5OJjYORnYuRhEQA4R7wMphroK7CQgOGAPIq9vvrsE6DoVBjg4sB+JVoWwHQ6gyiOrgZgDpWHqHYBmqiLE4ertYerFAAS3Hbpt8FUpAAAA83pUWHRNT0wgcmRraXQgMjAyMi4wOS41AAB4nI1SSQ6DMBC75xX+AChD2HJkU1VVkKql/UPv/b86AcGAqCISDmPjTMwYBb8e7e3zxbqSVilABx5rLd5Ga616+AJ1d7kOaMaqXpjGvYbxCSKQxrT32mp0/cIQGuhYTwsUZ3OxMuvJhHWR0MdiERo4YaNAxxT3v+8Pwow7RubE1bkXrh1Ddxc7ZUBYsskzHi03PGORA3E4ZZHT2ygDwm5od4nOGdduaCVjvxOJkgGMBOZhKrEwQCbDJ4a5jNjDQuboz5YyLQawMhRiSNuPnwjaWt8a9Xj5kblWPyLilIUuFqvtAAAAWHpUWHRTTUlMRVMgcmRraXQgMjAyMi4wOS41AAB4nHN21vDX1AjQsAWS/pr+mnCWQo2GgY6hnqmOtYGOAYjQBXPAJJCvawwmQTwkCTALqgKsACEPk9asAQBkZBdAp6OUhAAAAABJRU5ErkJggg==",
      "text/plain": [
       "<rdkit.Chem.rdchem.Mol at 0x7f99718c0dd0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# visualize one 2D molecule.\n",
    "from rdkit import Chem\n",
    "Chem.MolFromSmiles(g.smiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M3O_MZj_TjJ7"
   },
   "source": [
    "#  Building Graph Convolutional Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "0Agk424bTnmZ"
   },
   "outputs": [],
   "source": [
    "# Atom encoder\n",
    "\n",
    "class AtomEncoder(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super(AtomEncoder, self).__init__()\n",
    "\n",
    "        self.embeddings = torch.nn.ModuleList()\n",
    "\n",
    "        for i in range(9):\n",
    "            self.embeddings.append(torch.nn.Embedding(100, hidden_channels))\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        for embedding in self.embeddings:\n",
    "            embedding.reset_parameters()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.long()\n",
    "        \n",
    "        if x.dim() == 1:\n",
    "            x = x.unsqueeze(1)\n",
    "\n",
    "        out = 0\n",
    "        for i in range(x.size(1)):\n",
    "            out += self.embeddings[i](x[:, i])\n",
    "        return out\n",
    "\n",
    "\n",
    "# A simple graph neural network model\n",
    "\n",
    "from torch_geometric.nn import GCNConv\n",
    "from torch_geometric.nn import global_mean_pool as gap\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Linear\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels, num_node_features, num_classes):\n",
    "        super(GCN, self).__init__()\n",
    "        torch.manual_seed(42)\n",
    "        self.emb = AtomEncoder(hidden_channels=hidden_channels)\n",
    "        self.conv1 = GCNConv(hidden_channels,hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.lin = Linear(hidden_channels, num_classes)\n",
    "\n",
    "    def forward(self, batch):\n",
    "        x , edge_index, batch_size = batch.x, batch.edge_index, batch.batch\n",
    "        \n",
    "        x = self.emb(x)\n",
    "\n",
    "        # 1. Obtain node embeddings\n",
    "        x = self.conv1(x, edge_index)\n",
    "        \n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index)\n",
    "        \n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index)\n",
    "\n",
    "        # 2. Readout layer\n",
    "        x = gap(x, batch_size)  # [batch_size, hidden_channels]\n",
    "        # 3. Apply a final classifier\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kyi8ZC9YUCUS",
    "outputId": "112420f1-d6a0-4cec-b235-22f811683e12"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 12])\n"
     ]
    }
   ],
   "source": [
    "# create a model\n",
    "model = GCN(32, 9, 12)\n",
    "\n",
    "# prediction\n",
    "out = model(batch)\n",
    "print(out.shape) #(num_of_graph, num_of_task)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QIAfufrpuKAm"
   },
   "source": [
    "# Start training the GCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "vCQE_TIoUw4x"
   },
   "outputs": [],
   "source": [
    "# loss function and optimizer\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "criterion = nn.BCEWithLogitsLoss(reduction = \"none\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "6SL3fUw4VD38"
   },
   "outputs": [],
   "source": [
    "# train and eval function\n",
    "from sklearn.metrics import roc_auc_score\n",
    "def train(model, device, loader, optimizer):\n",
    "    model.train()\n",
    "\n",
    "    for step, batch in enumerate(loader):\n",
    "        batch = batch.to(device)\n",
    "        pred = model(batch)\n",
    "        y = batch.y.view(pred.shape).to(torch.float64)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        ## ignore nan targets (unlabeled) when computing training loss.\n",
    "        is_labeled = batch.y == batch.y\n",
    "        loss = criterion(pred.to(torch.float32)[is_labeled], batch.y.to(torch.float32)[is_labeled]).mean()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "def eval(model, device, loader):\n",
    "    model.eval()\n",
    "    y_true = []\n",
    "    y_pred = []\n",
    "    # For every batch in test loader\n",
    "    for batch in loader:\n",
    "\n",
    "        batch = batch.to(device)\n",
    "        if batch.x.shape[0] == 1:\n",
    "            pass\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                pred = model(batch)\n",
    "\n",
    "            y_true.append(batch.y.view(pred.shape))\n",
    "            y_pred.append(pred)\n",
    "\n",
    "    y_true = torch.cat(y_true, dim = 0).numpy()\n",
    "    y_pred = torch.cat(y_pred, dim = 0).numpy()\n",
    "    # Compute the ROC - AUC score and store as history\n",
    "    rocauc_list = []\n",
    "\n",
    "    for i in range(y_true.shape[1]):\n",
    "        #AUC is only defined when there is at least one positive data.\n",
    "        if np.sum(y_true[:,i] == 1) > 0 and np.sum(y_true[:,i] == 0) > 0:\n",
    "            # ignore nan values\n",
    "            is_labeled = y_true[:,i] == y_true[:,i]\n",
    "            rocauc_list.append(roc_auc_score(y_true[is_labeled,i], y_pred[is_labeled,i]))\n",
    "\n",
    "    if len(rocauc_list) == 0:\n",
    "        raise RuntimeError('No positively labeled data available. Cannot compute ROC-AUC.')\n",
    "\n",
    "    return {'rocauc': sum(rocauc_list)/len(rocauc_list)}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yYM7x48bZugM",
    "outputId": "39bdfa93-9de4-4510-e911-122dc9eb1437"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start training...\n",
      "====epoch 1\n",
      "{'Train': {'rocauc': 0.6934363319535813}, 'Validation': {'rocauc': 0.6638609055311179}}\n",
      "====epoch 2\n",
      "{'Train': {'rocauc': 0.710108537571113}, 'Validation': {'rocauc': 0.6722307325644348}}\n",
      "====epoch 3\n",
      "{'Train': {'rocauc': 0.7291279938561535}, 'Validation': {'rocauc': 0.702837336240314}}\n",
      "====epoch 4\n",
      "{'Train': {'rocauc': 0.7400275323165998}, 'Validation': {'rocauc': 0.7133099678195949}}\n",
      "====epoch 5\n",
      "{'Train': {'rocauc': 0.7485362178464104}, 'Validation': {'rocauc': 0.7198246707268581}}\n",
      "====epoch 6\n",
      "{'Train': {'rocauc': 0.7539066699475044}, 'Validation': {'rocauc': 0.7270453404123515}}\n",
      "====epoch 7\n",
      "{'Train': {'rocauc': 0.7612414990122806}, 'Validation': {'rocauc': 0.7272809293202961}}\n",
      "====epoch 8\n",
      "{'Train': {'rocauc': 0.7601087416648205}, 'Validation': {'rocauc': 0.732634011619033}}\n",
      "====epoch 9\n",
      "{'Train': {'rocauc': 0.7727335522869624}, 'Validation': {'rocauc': 0.7337990298386572}}\n",
      "====epoch 10\n",
      "{'Train': {'rocauc': 0.7729444654809609}, 'Validation': {'rocauc': 0.7324153677006873}}\n",
      "====epoch 11\n",
      "{'Train': {'rocauc': 0.7780783370436288}, 'Validation': {'rocauc': 0.7358012524254217}}\n",
      "====epoch 12\n",
      "{'Train': {'rocauc': 0.7775126564199368}, 'Validation': {'rocauc': 0.7369433831567168}}\n",
      "====epoch 13\n",
      "{'Train': {'rocauc': 0.7837555737604083}, 'Validation': {'rocauc': 0.7385679653186857}}\n",
      "====epoch 14\n",
      "{'Train': {'rocauc': 0.7865744360368496}, 'Validation': {'rocauc': 0.7371508383333443}}\n",
      "====epoch 15\n",
      "{'Train': {'rocauc': 0.7849227859273694}, 'Validation': {'rocauc': 0.7390215266466634}}\n",
      "====epoch 16\n",
      "{'Train': {'rocauc': 0.7896515362294622}, 'Validation': {'rocauc': 0.7415701227894959}}\n",
      "====epoch 17\n",
      "{'Train': {'rocauc': 0.7896109366316888}, 'Validation': {'rocauc': 0.7404580393502275}}\n",
      "====epoch 18\n",
      "{'Train': {'rocauc': 0.7911390657300692}, 'Validation': {'rocauc': 0.7412262604331513}}\n",
      "====epoch 19\n",
      "{'Train': {'rocauc': 0.7927013017260576}, 'Validation': {'rocauc': 0.7437383397658966}}\n"
     ]
    }
   ],
   "source": [
    "# Training\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(\"Start training...\")\n",
    "for epoch in range(1, 20):\n",
    "    print(\"====epoch \" + str(epoch))\n",
    "\n",
    "    # training\n",
    "    train(model, device, train_loader, optimizer)\n",
    "\n",
    "    # evaluating\n",
    "    train_acc = eval(model, device, train_loader)\n",
    "    val_acc = eval(model, device, val_loader)\n",
    "    print({'Train': train_acc, 'Validation': val_acc})\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameter Tuning for GCN "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_roc_auc(model, device, loader):\n",
    "    # Set model to evaluation mode\n",
    "    model.eval()\n",
    "    \n",
    "    # Initialize lists for true labels and predictions\n",
    "    y_true = []  # Ground truth\n",
    "    y_pred = []  # Predictions\n",
    "    \n",
    "    # Disable gradient computation for evaluation\n",
    "    with torch.no_grad():\n",
    "        # Iterate through data loader batches\n",
    "        for batch in loader:\n",
    "            # Move batch data to the specified device\n",
    "            batch = batch.to(device)\n",
    "            \n",
    "            # Get model predictions and store true labels and predictions\n",
    "            pred = model(batch)\n",
    "            y_true.append(batch.y.view(pred.shape))\n",
    "            y_pred.append(pred)\n",
    "\n",
    "    # Concatenate lists into numpy arrays\n",
    "    y_true = torch.cat(y_true, dim=0).numpy()\n",
    "    y_pred = torch.cat(y_pred, dim=0).numpy()\n",
    "\n",
    "    # Calculate ROC AUC scores for each class with labels\n",
    "    rocauc_list = []\n",
    "    for i in range(y_true.shape[1]):\n",
    "        if np.sum(y_true[:, i] == 1) > 0 and np.sum(y_true[:, i] == 0) > 0:\n",
    "            is_labeled = y_true[:, i] == y_true[:, i]\n",
    "            rocauc_list.append(roc_auc_score(y_true[is_labeled, i], y_pred[is_labeled, i]))\n",
    "\n",
    "    # Calculate and return the average ROC AUC score (if any)\n",
    "    return sum(rocauc_list) / len(rocauc_list) if len(rocauc_list) > 0 else 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hidden: 16, LR: 0.01, Dropout: 0.1, ROC AUC: 0.7237088554419236\n",
      "Hidden: 16, LR: 0.01, Dropout: 0.3, ROC AUC: 0.7237088554419236\n",
      "Hidden: 16, LR: 0.01, Dropout: 0.5, ROC AUC: 0.7237088554419236\n",
      "Hidden: 16, LR: 0.001, Dropout: 0.1, ROC AUC: 0.6725667023422709\n",
      "Hidden: 16, LR: 0.001, Dropout: 0.3, ROC AUC: 0.6725667023422709\n",
      "Hidden: 16, LR: 0.001, Dropout: 0.5, ROC AUC: 0.6725667023422709\n",
      "Hidden: 16, LR: 0.0001, Dropout: 0.1, ROC AUC: 0.612605538371896\n",
      "Hidden: 16, LR: 0.0001, Dropout: 0.3, ROC AUC: 0.612605538371896\n",
      "Hidden: 16, LR: 0.0001, Dropout: 0.5, ROC AUC: 0.612605538371896\n",
      "Hidden: 32, LR: 0.01, Dropout: 0.1, ROC AUC: 0.7271076268727348\n",
      "Hidden: 32, LR: 0.01, Dropout: 0.3, ROC AUC: 0.7271076268727348\n",
      "Hidden: 32, LR: 0.01, Dropout: 0.5, ROC AUC: 0.7271076268727348\n",
      "Hidden: 32, LR: 0.001, Dropout: 0.1, ROC AUC: 0.7093087743752432\n",
      "Hidden: 32, LR: 0.001, Dropout: 0.3, ROC AUC: 0.7093087743752432\n",
      "Hidden: 32, LR: 0.001, Dropout: 0.5, ROC AUC: 0.7093087743752432\n",
      "Hidden: 32, LR: 0.0001, Dropout: 0.1, ROC AUC: 0.6424193930059094\n",
      "Hidden: 32, LR: 0.0001, Dropout: 0.3, ROC AUC: 0.6424193930059094\n",
      "Hidden: 32, LR: 0.0001, Dropout: 0.5, ROC AUC: 0.6424193930059094\n",
      "Hidden: 64, LR: 0.01, Dropout: 0.1, ROC AUC: 0.7298426226065535\n",
      "Hidden: 64, LR: 0.01, Dropout: 0.3, ROC AUC: 0.7298426226065535\n",
      "Hidden: 64, LR: 0.01, Dropout: 0.5, ROC AUC: 0.7298426226065535\n",
      "Hidden: 64, LR: 0.001, Dropout: 0.1, ROC AUC: 0.7317036535922322\n",
      "Hidden: 64, LR: 0.001, Dropout: 0.3, ROC AUC: 0.7317036535922322\n",
      "Hidden: 64, LR: 0.001, Dropout: 0.5, ROC AUC: 0.7317036535922322\n",
      "Hidden: 64, LR: 0.0001, Dropout: 0.1, ROC AUC: 0.6723412905834102\n",
      "Hidden: 64, LR: 0.0001, Dropout: 0.3, ROC AUC: 0.6723412905834102\n",
      "Hidden: 64, LR: 0.0001, Dropout: 0.5, ROC AUC: 0.6723412905834102\n",
      "Best ROC AUC: 0.7317036535922322\n",
      "Best Hyperparameters: {'hidden_channels': 64, 'learning_rate': 0.001, 'dropout_rate': 0.1}\n"
     ]
    }
   ],
   "source": [
    "# Hyperparameters to tune\n",
    "hidden_channels_options = [16, 32, 64]  # Possible values for the number of hidden channels\n",
    "learning_rate_options = [0.01, 0.001, 0.0001]  # Possible learning rates\n",
    "dropout_rate_options = [0.1, 0.3, 0.5]  # Possible dropout rates\n",
    "\n",
    "best_roc_auc = 0  # Initialize the best ROC AUC score\n",
    "best_params = {}  # Initialize a dictionary to store the best hyperparameters\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else torch.device(\"cpu\"))  # Set the device to GPU if available, otherwise use CPU\n",
    "\n",
    "# Loop through different combinations of hyperparameters\n",
    "for hidden_channels in hidden_channels_options:\n",
    "    for learning_rate in learning_rate_options:\n",
    "        for dropout_rate in dropout_rate_options:\n",
    "            # Create a Graph Convolutional Network (GCN) model with specified hyperparameters\n",
    "            model = GCN(hidden_channels, 9, 12).to(device)\n",
    "            optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # Use Adam optimizer with the chosen learning rate\n",
    "            criterion = nn.BCEWithLogitsLoss(reduction=\"none\")  # Define loss criterion\n",
    "\n",
    "            # Train the model for a fixed number of epochs\n",
    "            for epoch in range(5):  # You can adjust the number of epochs here\n",
    "                train(model, device, train_loader, optimizer)  # Train the model on the training data\n",
    "                eval(model, device, val_loader)  # Evaluate the model on the validation data\n",
    "\n",
    "            # Calculate the ROC AUC score on the validation data\n",
    "            roc_auc = calculate_roc_auc(model, device, val_loader)\n",
    "\n",
    "            # Print the results for the current hyperparameter combination\n",
    "            print(f\"Hidden: {hidden_channels}, LR: {learning_rate}, Dropout: {dropout_rate}, ROC AUC: {roc_auc}\")\n",
    "\n",
    "            # Update the best ROC AUC score and hyperparameters if a better score is achieved\n",
    "            if roc_auc > best_roc_auc:\n",
    "                best_roc_auc = roc_auc\n",
    "                best_params = {'hidden_channels': hidden_channels, 'learning_rate': learning_rate, 'dropout_rate': dropout_rate}\n",
    "\n",
    "# Print the best ROC AUC score and the corresponding hyperparameters\n",
    "print(\"Best ROC AUC:\", best_roc_auc)\n",
    "print(\"Best Hyperparameters:\", best_params)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reinitializing the Model and Getting Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_hidden_channels = 64\n",
    "optimal_learning_rate = 0.01\n",
    "optimal_dropout_rate = 0.1\n",
    "\n",
    "model = GCN(optimal_hidden_channels, 9, 12).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====epoch 0\n",
      "Epoch 0: Train Loss: {'rocauc': 0.7457529316677588}, Validation Loss: {'rocauc': 0.7092484831924949}\n",
      "====epoch 1\n",
      "Epoch 1: Train Loss: {'rocauc': 0.7567118942697841}, Validation Loss: {'rocauc': 0.7208981503403864}\n",
      "====epoch 2\n",
      "Epoch 2: Train Loss: {'rocauc': 0.76436481935777}, Validation Loss: {'rocauc': 0.7246265288172676}\n",
      "====epoch 3\n",
      "Epoch 3: Train Loss: {'rocauc': 0.7624918112474862}, Validation Loss: {'rocauc': 0.7237942781537923}\n",
      "====epoch 4\n",
      "Epoch 4: Train Loss: {'rocauc': 0.7662901673796431}, Validation Loss: {'rocauc': 0.7318023629858145}\n",
      "====epoch 5\n",
      "Epoch 5: Train Loss: {'rocauc': 0.7757893239012738}, Validation Loss: {'rocauc': 0.7348679278312346}\n",
      "====epoch 6\n",
      "Epoch 6: Train Loss: {'rocauc': 0.7377744249893684}, Validation Loss: {'rocauc': 0.6763912659001776}\n",
      "====epoch 7\n",
      "Epoch 7: Train Loss: {'rocauc': 0.7730230301392789}, Validation Loss: {'rocauc': 0.731976426965232}\n",
      "====epoch 8\n",
      "Epoch 8: Train Loss: {'rocauc': 0.7762995459161681}, Validation Loss: {'rocauc': 0.731495434542962}\n",
      "====epoch 9\n",
      "Epoch 9: Train Loss: {'rocauc': 0.7848025128186468}, Validation Loss: {'rocauc': 0.7423215727792588}\n",
      "====epoch 10\n",
      "Epoch 10: Train Loss: {'rocauc': 0.7861398625428206}, Validation Loss: {'rocauc': 0.7437038979119942}\n",
      "====epoch 11\n",
      "Epoch 11: Train Loss: {'rocauc': 0.7787701115479195}, Validation Loss: {'rocauc': 0.7331123696060647}\n",
      "====epoch 12\n",
      "Epoch 12: Train Loss: {'rocauc': 0.7838210081183483}, Validation Loss: {'rocauc': 0.7320751355024905}\n",
      "====epoch 13\n",
      "Epoch 13: Train Loss: {'rocauc': 0.7900067506781007}, Validation Loss: {'rocauc': 0.735107202660869}\n",
      "====epoch 14\n",
      "Epoch 14: Train Loss: {'rocauc': 0.7852183583036876}, Validation Loss: {'rocauc': 0.7419733931842564}\n",
      "====epoch 15\n",
      "Epoch 15: Train Loss: {'rocauc': 0.7902510062418483}, Validation Loss: {'rocauc': 0.742653486082394}\n",
      "====epoch 16\n",
      "Epoch 16: Train Loss: {'rocauc': 0.7853744845735031}, Validation Loss: {'rocauc': 0.733254150058175}\n",
      "====epoch 17\n",
      "Epoch 17: Train Loss: {'rocauc': 0.792503630331109}, Validation Loss: {'rocauc': 0.7470677772149479}\n",
      "====epoch 18\n",
      "Epoch 18: Train Loss: {'rocauc': 0.7773483035774312}, Validation Loss: {'rocauc': 0.7333295621337915}\n",
      "====epoch 19\n",
      "Epoch 19: Train Loss: {'rocauc': 0.7820713429562374}, Validation Loss: {'rocauc': 0.7207288865585353}\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=optimal_learning_rate)\n",
    "criterion = nn.BCEWithLogitsLoss(reduction=\"none\")\n",
    "\n",
    "num_epochs = 20  \n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    print(\"====epoch \" + str(epoch))\n",
    "    train(model, device, train_loader, optimizer)  # Assuming you have a train function\n",
    "    train_loss = eval(model, device, train_loader)  # Assuming you have an evaluate function\n",
    "    val_loss = eval(model, device, val_loader)\n",
    "    print(f\"Epoch {epoch}: Train Loss: {train_loss}, Validation Loss: {val_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Valid ROC AUC: 0.7207288865585353\n"
     ]
    }
   ],
   "source": [
    "roc_auc = calculate_roc_auc(model, device, val_loader)  # Assuming you have a calculate_roc_auc function\n",
    "print(f\"Valid ROC AUC: {roc_auc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Outputting predictions\n",
    "def generate_predictions(model, device, loader):\n",
    "    model.eval()\n",
    "    predictions = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            logits = model(batch)\n",
    "            probs = torch.sigmoid(logits)\n",
    "            predictions.append(probs.cpu().numpy())\n",
    "\n",
    "    return np.concatenate(predictions, axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "predictions = generate_predictions(model, device, test_loader)\n",
    "\n",
    "# Convert to DataFrame for easier CSV writing\n",
    "prediction_df = pd.DataFrame(predictions)\n",
    "prediction_df.to_csv(\"test_output.csv\", index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Attention Network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "I99LMlSa4qSK"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.nn import GATConv\n",
    "from torch_geometric.nn import global_mean_pool\n",
    "import torch.nn.functional as F\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/defnecoban/opt/miniconda3/lib/python3.9/site-packages/torch_geometric/deprecation.py:22: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead\n",
      "  warnings.warn(out)\n"
     ]
    }
   ],
   "source": [
    "def remove_nan_instances(dataset):\n",
    "    # List comprehension filters out instances where any label is NaN\n",
    "    return [data for data in dataset if not torch.isnan(data.y).any()]\n",
    "\n",
    "# Remove instances with NaN labels from the training dataset\n",
    "cleaned_train_dataset = remove_nan_instances(train_dataset)\n",
    "\n",
    "# Remove instances with NaN labels from the validation dataset\n",
    "cleaned_val_dataset = remove_nan_instances(valid_dataset)\n",
    "\n",
    "# Create data loaders for the cleaned training and validation datasets\n",
    "train_loader = DataLoader(cleaned_train_dataset, batch_size=32, shuffle=True)\n",
    "val_loader = DataLoader(cleaned_val_dataset, batch_size=32, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GAT(torch.nn.Module):\n",
    "    def __init__(self, num_features, num_classes):\n",
    "        super(GAT, self).__init__()\n",
    "        self.encoder = AtomEncoder(hidden_channels=32)\n",
    "        \n",
    "        self.conv1 = GATConv(32, 8, heads=8, dropout=0.6)\n",
    "        # The output of the first layer will be 64 * 8 due to 8 attention heads\n",
    "        self.conv2 = GATConv(64, 64, heads=1, concat=False, dropout=0.6)\n",
    "        # Adding another GAT layer here. \n",
    "        self.conv3 = GATConv(64, 32, heads=1, concat=False, dropout=0.6)\n",
    "        # The final linear layer now takes the output of the last GAT layer as input\n",
    "        self.lin = torch.nn.Linear(32, num_classes)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        \n",
    "        x = x.long()\n",
    "        x = self.encoder(x)\n",
    "         \n",
    "\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = self.conv1(x, edge_index)\n",
    "        \n",
    "        x = F.elu(x)\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "         \n",
    "        x = F.elu(x)\n",
    "        # Adding the forward pass through the 3rd new layer\n",
    "        x = F.dropout(x, p=0.6, training=self.training)\n",
    "        x = self.conv3(x, edge_index)\n",
    "        \n",
    "        x = global_mean_pool(x, batch)  # Pool node features to graph-level features\n",
    "        x = self.lin(x)  # Apply a linear layer to get the final predictions\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9\n"
     ]
    }
   ],
   "source": [
    "# Initialize the GAT model\n",
    "num_features = train_dataset[0].num_node_features\n",
    "num_classes = 12  # Assuming 12 molecular property tasks as mentioned\n",
    "print(num_features)\n",
    "\n",
    "# Set the device to GPU if available, else CPU\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "model = GAT(num_features=num_features, num_classes=num_classes)\n",
    "model.to(device)\n",
    "\n",
    "# Set up the optimizer and loss function\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)\n",
    "criterion = torch.nn.BCEWithLogitsLoss(reduction=\"none\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training function\n",
    "def train(model, device, loader, optimizer, criterion):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "\n",
    "        #print(f\"Max edge index: {batch.edge_index.max()}, Number of nodes: {batch.num_nodes}\")\n",
    "        assert batch.edge_index.max() < batch.num_nodes\n",
    "        \n",
    "        # Ensure the data is in the correct format\n",
    "        batch.x = batch.x.float()  # Node features should be floats\n",
    "        batch.y = batch.y.float()  # Target labels should be floats for BCEWithLogitsLoss\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        out = model(batch)\n",
    "        \n",
    "        # Calculate loss - ensure out and batch.y are compatible types\n",
    "        loss = criterion(out, batch.y).mean()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    return total_loss / len(loader)\n",
    "\n",
    "# Evaluation function\n",
    "def evaluate(model, device, loader, criterion):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "\n",
    "            batch.x = batch.x.float() \n",
    "            batch.y = batch.y.float()\n",
    "\n",
    "            out = model(batch)\n",
    "            loss = criterion(out, batch.y).mean()  # Adjust loss calculation based on your data structure\n",
    "            total_loss += loss.item()\n",
    "    return total_loss / len(loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial loss: 0.821220338344574\n"
     ]
    }
   ],
   "source": [
    "# Function to compute initial loss on the validation dataset\n",
    "def compute_initial_loss(model, device, loader, criterion):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        batch = next(iter(loader)).to(device)\n",
    "        batch.x = batch.x.float()\n",
    "        batch.y = batch.y.float()\n",
    "        out = model(batch)\n",
    "        loss = criterion(out, batch.y).mean()\n",
    "        return loss.item()\n",
    "\n",
    "# Compute and print the initial loss on the validation dataset\n",
    "initial_loss = compute_initial_loss(model, device, val_loader, criterion)\n",
    "print(f\"Initial loss: {initial_loss}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No NaN found in the data\n",
      "No NaN found in the data\n"
     ]
    }
   ],
   "source": [
    "def check_nan_in_data(loader): #Making sure there isn't nan values \n",
    "    for batch in loader:\n",
    "        if torch.isnan(batch.x).any() or torch.isnan(batch.y).any():\n",
    "            print(\"NaN found in batch.x\" if torch.isnan(batch.x).any() else \"NaN found in batch.y\")\n",
    "            break\n",
    "    else:\n",
    "        print(\"No NaN found in the data\")\n",
    "\n",
    "check_nan_in_data(train_loader)  # Check training data\n",
    "check_nan_in_data(val_loader)    # Check validation data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect model predictions and true labels from a data loader\n",
    "def collect_predictions_and_labels(model, device, loader):\n",
    "    model.eval()\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "\n",
    "            batch.x = batch.x.float()\n",
    "            batch.y = batch.y.float()\n",
    "\n",
    "            preds = model(batch)\n",
    "            all_preds.append(preds.cpu().numpy())\n",
    "            all_labels.append(batch.y.cpu().numpy())\n",
    "\n",
    "    return np.concatenate(all_preds, axis=0), np.concatenate(all_labels, axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====epoch 1\n",
      "Epoch: 1, Training Loss: 0.2341, Validation Loss: 0.1979, rocauc: 0.4277\n",
      "====epoch 2\n",
      "Epoch: 2, Training Loss: 0.1436, Validation Loss: 0.1409, rocauc: 0.4593\n",
      "====epoch 3\n",
      "Epoch: 3, Training Loss: 0.1315, Validation Loss: 0.1386, rocauc: 0.4776\n",
      "====epoch 4\n",
      "Epoch: 4, Training Loss: 0.1280, Validation Loss: 0.1533, rocauc: 0.4844\n",
      "====epoch 5\n",
      "Epoch: 5, Training Loss: 0.1251, Validation Loss: 0.1326, rocauc: 0.4917\n",
      "====epoch 6\n",
      "Epoch: 6, Training Loss: 0.1230, Validation Loss: 0.1415, rocauc: 0.4978\n",
      "====epoch 7\n",
      "Epoch: 7, Training Loss: 0.1232, Validation Loss: 0.1354, rocauc: 0.5230\n",
      "====epoch 8\n",
      "Epoch: 8, Training Loss: 0.1206, Validation Loss: 0.1333, rocauc: 0.5330\n",
      "====epoch 9\n",
      "Epoch: 9, Training Loss: 0.1205, Validation Loss: 0.1513, rocauc: 0.5282\n"
     ]
    }
   ],
   "source": [
    "# Training loop\n",
    "for epoch in range(1,10):\n",
    "    print(\"====epoch \" + str(epoch))\n",
    "    train_loss = train(model, device, train_loader, optimizer, criterion)\n",
    "    val_loss = evaluate(model, device, val_loader, criterion)\n",
    "\n",
    "    # Collect predictions and labels from the validation set\n",
    "    preds, labels = collect_predictions_and_labels(model, device, val_loader)\n",
    "\n",
    "    # Calculate ROC AUC for each task\n",
    "    roc_auc_scores = []\n",
    "    for i in range(labels.shape[1]):\n",
    "        if len(np.unique(labels[:, i])) == 2:\n",
    "            roc_auc = roc_auc_score(labels[:, i], preds[:, i])\n",
    "            roc_auc_scores.append(roc_auc)\n",
    "\n",
    "    # Calculate average ROC AUC\n",
    "    average_roc_auc = np.mean(roc_auc_scores) if roc_auc_scores else 0\n",
    "\n",
    "\n",
    "    print(f'Epoch: {epoch}, Training Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, rocauc: {average_roc_auc:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Isomorphism Network "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GINConv, global_add_pool, Sequential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GIN(torch.nn.Module):\n",
    "    def __init__(self, num_features, num_classes):\n",
    "        super(GIN, self).__init__()\n",
    "\n",
    "        # Define the first GINConv layer with a neural network module nn1\n",
    "        nn1 = Sequential('x', [\n",
    "            (torch.nn.Linear(num_features, 32), 'x -> x'),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(32, 32)\n",
    "        ])\n",
    "        self.conv1 = GINConv(nn1, train_eps=True)\n",
    "\n",
    "        # Define the second GINConv layer with a neural network module nn2\n",
    "        nn2 = Sequential('x', [\n",
    "            (torch.nn.Linear(32, 32), 'x -> x'),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(32, 32)\n",
    "        ])\n",
    "        self.conv2 = GINConv(nn2, train_eps=True)\n",
    "\n",
    "        # Define a linear layer for final class predictions\n",
    "        self.fc = torch.nn.Linear(32, num_classes)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "\n",
    "        # Pass through the first and second GINConv layers\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = self.conv2(x, edge_index)\n",
    "\n",
    "        # Aggregate node features using global add pooling\n",
    "        x = global_add_pool(x, batch)\n",
    "        \n",
    "        # Apply a linear layer to get the final class predictions\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_gin(model, device, loader, optimizer, criterion):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        out = model(batch)\n",
    "        loss = criterion(out, batch.y.float()).mean()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    return total_loss / len(loader)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_gin(model, device, loader, criterion):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            out = model(batch)\n",
    "            loss = criterion(out, batch.y.float()).mean()\n",
    "            total_loss += loss.item()\n",
    "    return total_loss / len(loader)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_predictions_and_labels_gin(model, device, loader):\n",
    "    model.eval()\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            preds = model(batch)\n",
    "            all_preds.append(preds.cpu().numpy())\n",
    "            all_labels.append(batch.y.cpu().numpy())\n",
    "    return np.concatenate(all_preds, axis=0), np.concatenate(all_labels, axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== Epoch 1\n",
      "Epoch: 1, Training Loss: 0.7183, Validation Loss: 0.2396, Average ROC AUC: 0.3360\n",
      "==== Epoch 2\n",
      "Epoch: 2, Training Loss: 0.1890, Validation Loss: 0.2436, Average ROC AUC: 0.3589\n",
      "==== Epoch 3\n",
      "Epoch: 3, Training Loss: 0.1791, Validation Loss: 0.2252, Average ROC AUC: 0.3795\n",
      "==== Epoch 4\n",
      "Epoch: 4, Training Loss: 0.1733, Validation Loss: 0.1846, Average ROC AUC: 0.4085\n",
      "==== Epoch 5\n",
      "Epoch: 5, Training Loss: 0.1627, Validation Loss: 0.2161, Average ROC AUC: 0.4269\n",
      "==== Epoch 6\n",
      "Epoch: 6, Training Loss: 0.1595, Validation Loss: 0.1827, Average ROC AUC: 0.4672\n",
      "==== Epoch 7\n",
      "Epoch: 7, Training Loss: 0.1556, Validation Loss: 0.1874, Average ROC AUC: 0.4476\n",
      "==== Epoch 8\n",
      "Epoch: 8, Training Loss: 0.1555, Validation Loss: 0.1827, Average ROC AUC: 0.4782\n",
      "==== Epoch 9\n",
      "Epoch: 9, Training Loss: 0.1540, Validation Loss: 0.1887, Average ROC AUC: 0.4522\n"
     ]
    }
   ],
   "source": [
    "# Initialize the GIN model\n",
    "model_gin = GIN(num_features, num_classes).to(device)\n",
    "\n",
    "# Choose an optimizer and loss function\n",
    "optimizer = torch.optim.Adam(model_gin.parameters(), lr=0.001)\n",
    "criterion = torch.nn.BCEWithLogitsLoss(reduction=\"none\")\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(1, 10):\n",
    "    print(f\"==== Epoch {epoch}\")\n",
    "    train_loss = train_gin(model_gin, device, train_loader, optimizer, criterion)\n",
    "    val_loss = evaluate_gin(model_gin, device, val_loader, criterion)\n",
    "\n",
    "    preds, labels = collect_predictions_and_labels_gin(model_gin, device, val_loader)\n",
    "    roc_auc_scores = []\n",
    "    for i in range(labels.shape[1]):\n",
    "        if len(np.unique(labels[:, i])) == 2:\n",
    "            roc_auc = roc_auc_score(labels[:, i], preds[:, i])\n",
    "            roc_auc_scores.append(roc_auc)\n",
    "\n",
    "    average_roc_auc = np.mean(roc_auc_scores) if roc_auc_scores else 0\n",
    "    print(f\"Epoch: {epoch}, Training Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, Average ROC AUC: {average_roc_auc:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combining Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv, GINConv, global_mean_pool, Sequential\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Linear\n",
    "\n",
    "class CombinedGNN(torch.nn.Module):\n",
    "    def __init__(self, num_features, num_classes):\n",
    "        super(CombinedGNN, self).__init__()\n",
    "        # GCN layers\n",
    "        self.gcn_conv1 = GCNConv(num_features, 32)\n",
    "        self.gcn_conv2 = GCNConv(32, 32)\n",
    "\n",
    "        # GIN layers\n",
    "        nn1 = Sequential('x', [\n",
    "            (torch.nn.Linear(num_features, 32), 'x -> x'),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(32, 32)\n",
    "        ])\n",
    "        self.gin_conv1 = GINConv(nn1)\n",
    "        \n",
    "        nn2 = Sequential('x', [\n",
    "            (torch.nn.Linear(32, 32), 'x -> x'),  # Note the input size is 32 here\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(32, 32)\n",
    "        ])\n",
    "        self.gin_conv2 = GINConv(nn2)\n",
    "\n",
    "        # Linear layer to combine features\n",
    "        self.combined_lin = torch.nn.Linear(64, num_classes)  \n",
    "\n",
    "\n",
    "\n",
    "    def forward(self, data):\n",
    "        \n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        x = x.float()\n",
    "\n",
    "        # GCN part\n",
    "        x_gcn = F.relu(self.gcn_conv1(x, edge_index))\n",
    "        x_gcn = F.relu(self.gcn_conv2(x_gcn, edge_index))\n",
    "        x_gcn = global_mean_pool(x_gcn, batch)\n",
    "\n",
    "        # GIN part\n",
    "        x_gin = F.relu(self.gin_conv1(x, edge_index))\n",
    "        x_gin = F.relu(self.gin_conv2(x_gin, edge_index))\n",
    "        x_gin = global_mean_pool(x_gin, batch)\n",
    "\n",
    "        # Combine GCN and GIN outputs\n",
    "        x_combined = torch.cat([x_gcn, x_gin], dim=1)\n",
    "        \n",
    "        x_out = self.combined_lin(x_combined)\n",
    "        return x_out\n",
    "\n",
    "       \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model on the dataset in batches\n",
    "def train_combined(model, device, loader, optimizer, criterion):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "       \n",
    "        optimizer.zero_grad()\n",
    "        out = model(batch)\n",
    "        loss = criterion(out, batch.y.float()).mean()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    return total_loss / len(loader)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the model's performance on a dataset in batches\n",
    "def evaluate_combined(model, device, loader, criterion):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    \n",
    "    # Disable gradient computation for efficiency\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            \n",
    "            # Ensure node features and labels are of type float\n",
    "            batch.x = batch.x.float() \n",
    "            batch.y = batch.y.float()\n",
    "            \n",
    "            # Compute model predictions and loss\n",
    "            out = model(batch)\n",
    "            loss = criterion(out, batch.y.float()).mean()\n",
    "            total_loss += loss.item()\n",
    "    \n",
    "    # Return the average loss across all batches\n",
    "    return total_loss / len(loader)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect model predictions and true labels from a dataset in batches\n",
    "def collect_predictions_and_labels_combined(model, device, loader):\n",
    "    model.eval()\n",
    "    all_preds = []  # Predictions\n",
    "    all_labels = []  # True labels\n",
    "    \n",
    "    # Disable gradient computation for efficiency\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "\n",
    "            # Ensure node features and labels are of type float\n",
    "            batch.x = batch.x.float() \n",
    "            batch.y = batch.y.float()\n",
    "            \n",
    "            # Generate model predictions and collect them\n",
    "            preds = model(batch)\n",
    "            all_preds.append(preds.cpu().numpy())\n",
    "            \n",
    "            # Collect true labels\n",
    "            all_labels.append(batch.y.cpu().numpy())\n",
    "    \n",
    "    # Concatenate predictions and labels along the specified axis\n",
    "    return np.concatenate(all_preds, axis=0), np.concatenate(all_labels, axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch loaded\n",
      "Batch has data\n"
     ]
    }
   ],
   "source": [
    "for batch in train_loader:\n",
    "    print(\"Batch loaded\")\n",
    "    # Check if the batch contains data\n",
    "    if batch.x is not None:\n",
    "        print(\"Batch has data\")\n",
    "    break  # Just check the first batch for now\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== Epoch 1\n",
      "Epoch: 1, Training Loss: 0.2252, Validation Loss: 0.1347, Average ROC AUC: 0.4277\n",
      "==== Epoch 2\n",
      "Epoch: 2, Training Loss: 0.1247, Validation Loss: 0.1349, Average ROC AUC: 0.4441\n",
      "==== Epoch 3\n",
      "Epoch: 3, Training Loss: 0.1236, Validation Loss: 0.1353, Average ROC AUC: 0.4642\n",
      "==== Epoch 4\n",
      "Epoch: 4, Training Loss: 0.1210, Validation Loss: 0.1336, Average ROC AUC: 0.4803\n"
     ]
    }
   ],
   "source": [
    "# Initialize the combined GNN model\n",
    "model_combined = CombinedGNN(num_features, num_classes).to(device)\n",
    "\n",
    "# Choose an optimizer and loss function\n",
    "optimizer = torch.optim.Adam(model_combined.parameters(), lr=0.001)\n",
    "criterion = torch.nn.BCEWithLogitsLoss(reduction=\"none\")\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(1, 5):\n",
    "    print(f\"==== Epoch {epoch}\")\n",
    "    train_loss = train_combined(model_combined, device, train_loader, optimizer, criterion)\n",
    "    val_loss = evaluate_combined(model_combined, device, val_loader, criterion)\n",
    "\n",
    "    preds, labels = collect_predictions_and_labels_combined(model_combined, device, val_loader)\n",
    "    roc_auc_scores = []\n",
    "    for i in range(labels.shape[1]):\n",
    "        if len(np.unique(labels[:, i])) == 2:\n",
    "            roc_auc = roc_auc_score(labels[:, i], preds[:, i])\n",
    "            roc_auc_scores.append(roc_auc)\n",
    "\n",
    "    average_roc_auc = np.mean(roc_auc_scores) if roc_auc_scores else 0\n",
    "    print(f\"Epoch: {epoch}, Training Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, Average ROC AUC: {average_roc_auc:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GraphSAGE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import SAGEConv\n",
    "\n",
    "class GraphSAGE(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels, num_node_features, num_classes):\n",
    "        super(GraphSAGE, self).__init__()\n",
    "        self.encoder = AtomEncoder(hidden_channels)\n",
    "        self.conv1 = SAGEConv(hidden_channels, hidden_channels)\n",
    "        self.conv2 = SAGEConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = SAGEConv(hidden_channels, hidden_channels)\n",
    "        self.lin = torch.nn.Linear(hidden_channels, num_classes)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "\n",
    "        x = self.encoder(x)\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        x = F.relu(self.conv2(x, edge_index))\n",
    "        x = F.relu(self.conv3(x, edge_index))\n",
    "        x = global_mean_pool(x, batch)  # Graph-level classification\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.lin(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, device, loader, optimizer, criterion):\n",
    "    model.train()  # Set the model to training mode\n",
    "\n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "        optimizer.zero_grad()  # Clear gradients\n",
    "\n",
    "        out = model(batch)  # Forward pass\n",
    "        # print(\"Model output shape:\", out.shape)  # Print the shape of the model output\n",
    "        # print(\"Target shape:\", batch.y.shape)  # Print the shape of the target labels\n",
    "\n",
    "        loss = criterion(out, batch.y)  # Compute loss\n",
    "        #print(\"Loss:\", loss)  # Print the loss value\n",
    "        \n",
    "        loss.backward()  # Backward pass (compute gradients)\n",
    "        optimizer.step()  # Update model parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_with_auc(model, device, loader):\n",
    "    model.eval()\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            out = model(batch)\n",
    "            \n",
    "            preds = torch.sigmoid(out)  # Apply sigmoid to get probabilities\n",
    "            all_preds.append(preds.cpu())\n",
    "            all_labels.append(batch.y.cpu())\n",
    "\n",
    "    all_preds = torch.cat(all_preds, dim=0).numpy()\n",
    "    all_labels = torch.cat(all_labels, dim=0).numpy()\n",
    "\n",
    "    # Calculate ROC AUC for each class and average\n",
    "    roc_auc = roc_auc_score(all_labels, all_preds, average='macro')\n",
    "    return roc_auc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model setup\n",
    "hidden_channels = 32  \n",
    "num_node_features = 9  \n",
    "num_classes = 12     \n",
    "model = GraphSAGE(hidden_channels, num_node_features, num_classes).to(device)\n",
    "\n",
    "# Optimizer setup\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Criterion (Loss function) setup for multi-label classification\n",
    "criterion = torch.nn.BCEWithLogitsLoss(reduction='mean').to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== Epoch 1\n",
      "Epoch: 1, Val ROC AUC: 0.5455\n",
      "==== Epoch 2\n",
      "Epoch: 2, Val ROC AUC: 0.6127\n",
      "==== Epoch 3\n",
      "Epoch: 3, Val ROC AUC: 0.6342\n",
      "==== Epoch 4\n",
      "Epoch: 4, Val ROC AUC: 0.6435\n",
      "==== Epoch 5\n",
      "Epoch: 5, Val ROC AUC: 0.6750\n",
      "==== Epoch 6\n",
      "Epoch: 6, Val ROC AUC: 0.6628\n",
      "==== Epoch 7\n",
      "Epoch: 7, Val ROC AUC: 0.6890\n",
      "==== Epoch 8\n",
      "Epoch: 8, Val ROC AUC: 0.7084\n",
      "==== Epoch 9\n",
      "Epoch: 9, Val ROC AUC: 0.7083\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(1, 10):\n",
    "    print(\"==== Epoch \" + str(epoch))\n",
    "\n",
    "    train(model, device, train_loader, optimizer, criterion)  # Train the model\n",
    "    val_roc_auc = evaluate_with_auc(model, device, val_loader)  # Evaluate the model\n",
    "    print(f'Epoch: {epoch}, Val ROC AUC: {val_roc_auc:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}