2_Graph_Datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install --use-pep517 \"graphein[extras]\" lightning torch torch-geometric tensorboard nbformat \"jsonargparse[signatures]\" ipywidgets tabulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import graphein\n",
    "graphein.verbose(enabled=False)\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "from graphein.ml import GraphFormatConvertor\n",
    "import torch\n",
    "import lightning\n",
    "from torch.utils.data import random_split\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch_geometric.data import Dataset\n",
    "from pathlib import Path\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from torch_geometric.utils import to_networkx\n",
    "from graphein.protein.visualisation import plotly_protein_structure_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Don't forget to use your `load_graph` function from the previous notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphein.protein.config import ProteinGraphConfig\n",
    "from graphein.protein.graphs import construct_graph\n",
    "from graphein.protein.features.nodes import amino_acid as graphein_nodes\n",
    "from graphein.protein import edges as graphein_edges\n",
    "from graphein.protein.subgraphs import extract_subgraph\n",
    "from functools import partial\n",
    "\n",
    "def load_graph(pdb_id, chain):\n",
    "    graph_config = ProteinGraphConfig(\n",
    "        node_metadata_functions = [graphein_nodes.amino_acid_one_hot, graphein_nodes.meiler_embedding],\n",
    "        edge_construction_functions = [graphein_edges.add_peptide_bonds, \n",
    "                                       partial(graphein_edges.add_distance_threshold, \n",
    "                                               threshold=8.,\n",
    "                                               long_interaction_threshold=2)],\n",
    "    )\n",
    "    graph = construct_graph(pdb_code=pdb_id, config=graph_config, verbose=False)\n",
    "    interface_residues = set()\n",
    "    for source, target, kind in graph.edges(data=True):\n",
    "        c1, c2 = source.split(\":\")[0], target.split(\":\")[0]\n",
    "        if 'distance_threshold' in kind['kind']:\n",
    "            if c1 == chain and c2 != chain:\n",
    "                interface_residues.add(source)\n",
    "            elif c2 == chain and c1 != chain:\n",
    "                interface_residues.add(target)\n",
    "    graph = extract_subgraph(graph, chains=chain)\n",
    "    for node, data in graph.nodes(data=True):\n",
    "        if node in interface_residues:\n",
    "            data['interface_label'] = 1\n",
    "        else:\n",
    "            data['interface_label'] = 0\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting graphs to deep learning datasets\n",
    "\n",
    "Deep learning libraries like PyTorch, and by extension PyTorch-Geometric, have some standardized ways of handling data and datasets, in order to optimize the operations they perform on the various numeric features involved. In this notebook, we will see how to convert a graph into a torch `Data` object, which is the standard way of representing a graph in PyTorch-Geometric. Then we'll go from a single graph to a `Dataset` of graphs, which is the standard way of representing a dataset in PyTorch. And finally, we'll see how to wrap this `Dataset` into a Lightning `DataModule`, which is the standard way of handling data-related operations in PyTorch-Lightning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data \n",
    "\n",
    "We first need to make a torch `Data` object from our graphs. This is easily done with graphein's conversion functions, specifically the `GraphFormatConvertor`, where you can specify which features of the NetworkX graph you'd like to retain in the `Data` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = [\n",
    "            \"chain_id\",\n",
    "            \"coords\",\n",
    "            \"edge_index\",\n",
    "            \"node_id\",\n",
    "            \"residue_number\",\n",
    "            \"amino_acid_one_hot\",\n",
    "            \"meiler\",\n",
    "            \"interface_label\"\n",
    "]\n",
    "convertor = GraphFormatConvertor(src_format=\"nx\", # From NetworkX \n",
    "                                 dst_format=\"pyg\", # To PyTorch Geometric\n",
    "                                 columns=columns, # The columns to be used\n",
    "                                 verbose=None)\n",
    "graphein_graph = load_graph(\"1A0G\", \"A\")\n",
    "torch_data = convertor(graphein_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(zip(torch_data.node_id[:5], torch_data.interface_label[:5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_data.edge_index.T[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, here are the amino acid types across interface and non-interface residues:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extracted_amino_acids = [s.split(\":\")[1] for s in torch_data.node_id]\n",
    "\n",
    "data_with_aa = {\"amino acid\": extracted_amino_acids, \"interface labels\": torch_data.interface_label}\n",
    "\n",
    "plt.figure(figsize=(14, 8))\n",
    "sns.countplot(x=\"amino acid\", hue=\"interface labels\", data=data_with_aa, palette=\"Set2\")\n",
    "plt.title(\"Distribution of amino acids types across interface and non-interface residues\")\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Amino acid\")\n",
    "plt.legend(title=\"Interface label\", loc='upper right')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the `to_networkx` function to convert the `Data` object back to a NetworkX graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "graphein_graph_again = to_networkx(torch_data, \n",
    "                                   node_attrs=[\"chain_id\",\n",
    "                                               \"coords\",\n",
    "                                               \"node_id\",\n",
    "                                               \"residue_number\",\n",
    "                                               \"amino_acid_one_hot\",\n",
    "                                               \"meiler\",\n",
    "                                               \"interface_label\"],\n",
    "                                    to_undirected=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "p = plotly_protein_structure_graph(\n",
    "    graphein_graph_again,\n",
    "    colour_edges_by=None,\n",
    "    colour_nodes_by='interface_label',\n",
    "    label_node_ids=False,\n",
    "    plot_title=\"Peptide backbone graph with distance connections. Nodes coloured by interface labels.\",\n",
    "    node_size_multiplier=1\n",
    "    )\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `torch_geometric.data.Dataset` class is a standard way of representing a graph dataset in PyTorch. It is an abstract class that you can subclass to create your own dataset. The functions that need to be included are:\n",
    "\n",
    "- `download()`: this downloads the dataset (in our case from `dataset.txt`) and saves each data point (in our case as a pickle file containing the graphein graph that our `load_graph` function returns) in `self.raw_dir`.\n",
    "- `process()`: this processes the data from `self.raw_dir` to torch-geometric `Data` objects (as we did above), and saves them as `.pt` files in `self.processed_dir`.\n",
    "- property functions: `raw_file_names`, `processed_file_names` return the names of the raw pickle and processed pt files for each data point.\n",
    "- `len()`: this returns the number of graphs in the dataset\n",
    "- `get()`: this returns the `Data` object for a given index\n",
    "\n",
    "See the [documentation](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.data.Dataset.html#torch_geometric.data.Dataset) for more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProteinDataset(Dataset):\n",
    "    \"\"\"\n",
    "    torch-geometric Dataset class for loading protein files as graphs.\n",
    "    \"\"\"\n",
    "    def __init__(self, root,\n",
    "                 protein_names: list):\n",
    "        columns = [\n",
    "            \"chain_id\",\n",
    "            \"coords\",\n",
    "            \"edge_index\",\n",
    "            \"kind\",\n",
    "            \"node_id\",\n",
    "            \"residue_number\",\n",
    "            \"amino_acid_one_hot\",\n",
    "            \"meiler\",\n",
    "            \"interface_label\",\n",
    "        ]\n",
    "        self.convertor = GraphFormatConvertor(src_format=\"nx\", dst_format=\"pyg\", columns=columns, verbose=None)\n",
    "        self.protein_names = protein_names\n",
    "        super(ProteinDataset, self).__init__(root)\n",
    "\n",
    "    def download(self):\n",
    "        for protein_name in self.protein_names:\n",
    "            output = Path(self.raw_dir) / f'{protein_name}.pkl'\n",
    "            if not output.exists():\n",
    "                pdb_id, chain = protein_name.split(\"_\")\n",
    "                graphein_graph = load_graph(pdb_id, chain)\n",
    "                with open(output, \"wb\") as f:\n",
    "                    pickle.dump(graphein_graph, f)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return [Path(self.raw_dir) / f\"{protein_name}.pkl\" for protein_name in self.protein_names if (Path(self.raw_dir) / f\"{protein_name}.pt\").exists()]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return [Path(self.processed_dir) / f\"{protein_name}.pt\" for protein_name in self.protein_names if (Path(self.processed_dir) / f\"{protein_name}.pt\").exists()]\n",
    "\n",
    "    def process(self):\n",
    "        for protein_name in self.protein_names:\n",
    "            output = Path(self.processed_dir) / f'{protein_name}.pt'\n",
    "            if not output.exists():\n",
    "                with open(Path(self.raw_dir) / f\"{protein_name}.pkl\", \"rb\") as f:\n",
    "                    graphein_graph = pickle.load(f)\n",
    "                torch_data = self.convertor(graphein_graph)\n",
    "                torch.save(torch_data, output)\n",
    "\n",
    "    def len(self):\n",
    "        return len(self.processed_file_names)\n",
    "\n",
    "    def get(self, idx):\n",
    "        data = torch.load(self.processed_file_names[idx])\n",
    "        return data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make our dataset! We run it for the first 20 proteins to save time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('dataset.txt') as f:\n",
    "    protein_names = [line.strip() for line in f]\n",
    "\n",
    "dataset = ProteinDataset(root='./test_data', protein_names=protein_names[:20])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_amino_acids = []\n",
    "all_interface_labels = []\n",
    "\n",
    "for torch_graph in dataset:\n",
    "    extracted_amino_acids = [s.split(\":\")[1] for s in torch_graph.node_id]\n",
    "    all_amino_acids.extend(extracted_amino_acids)\n",
    "    all_interface_labels.extend(torch_graph.interface_label.tolist())\n",
    "\n",
    "data_with_aa = {\n",
    "    \"amino acid\": all_amino_acids,\n",
    "    \"interface labels\": all_interface_labels\n",
    "}\n",
    "\n",
    "plt.figure(figsize=(14, 8))\n",
    "sns.countplot(x=\"amino acid\", hue=\"interface labels\", data=data_with_aa, palette=\"Set2\")\n",
    "plt.title(\"Distribution of amino acids types across interface and non-interface residues\")\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Amino acid\")\n",
    "plt.legend(title=\"Interface label\", loc='upper right')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Graphein also has a built-in `ProteinGraphDataset` class that combines these steps. It also has some nice features like \n",
    "- the ability to load a dataset of proteins from both the PDB or AlphaFold Database directory of PDB files\n",
    "- the ability to apply custom transformations from your bioinformatics tools of choice to the PDB files (with the `pdb_transform` argument)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataModule\n",
    "\n",
    "Now that we have our `Dataset` ready, we need to specificy how the `Data` objects within the created Dataset are split into training, validation and test sets. This is where PyTorch Lightning's DataModule comes in ([documentation](https://lightning.ai/docs/pytorch/stable/data/datamodule.html)). The `DataModule` is a class that encapsulates the logic for loading, batching and splitting the data. It's a way of separating the logic for data loading and batching separate from both the data ingestion and the model and training logic, which makes the code more modular and easier to maintain. It also makes it easier to switch between different datasets and data loading strategies.\n",
    "\n",
    "To define a `DataModule` the following methods are necessary:\n",
    "- `prepare_data()` - this defines the downloading and IO operations that are generally slower and need to only be run once. In our case it just runs the Dataset function once to download and process all the pickle and pt files. This is called once in the beginning of training so all the future calls of Dataset in setup (which is called on every node/process) just loads the data from the saved files.\n",
    "- `setup()` - this defines how to split the dataset. It also takes a `stage` argument (one of `fit,validate,test,predict`).\n",
    "- `train_dataloader()` - this returns the `DataLoader` for the training data\n",
    "\n",
    "and the following are optional:\n",
    "- `val_dataloader()` - this returns the `DataLoader` for the validation data\n",
    "- `test_dataloader()` - this returns the `DataLoader` for the test data\n",
    "- `predict_dataloader()` - this returns the `DataLoader` for the inference data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ProteinGraphDataModule(lightning.LightningDataModule):\n",
    "    def __init__(self, root, dataset_file, batch_size=8):\n",
    "        super().__init__()\n",
    "        self.root = root\n",
    "        self.dataset_file = dataset_file\n",
    "        with open(dataset_file) as f:\n",
    "            self.protein_names = [line.strip() for line in f]\n",
    "        self.protein_names = self.protein_names[:20] # SMALL DATASET FOR TESTING\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "    def prepare_data(self):\n",
    "        ProteinDataset(root=self.root, protein_names=self.protein_names)\n",
    "    \n",
    "    def setup(self, stage):\n",
    "        dataset = ProteinDataset(root=self.root, protein_names=self.protein_names)\n",
    "        # Here we just do a random split of 80/10/10 for train/val/test\n",
    "        train_idx, val_idx, test_idx = random_split(range(len(dataset)), [0.8, 0.1, 0.1])\n",
    "        self.train, self.val, self.test = dataset[list(train_idx)], dataset[list(val_idx)], dataset[list(test_idx)]\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.train, batch_size=self.batch_size)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.val, batch_size=self.batch_size)\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(self.test, batch_size=self.batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `DataModule` is now ready, give it a try and loop through the dataloader to see how they work!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "datamodule = ProteinGraphDataModule(\"./test_data\", \"dataset.txt\")\n",
    "datamodule.prepare_data()\n",
    "datamodule.setup(\"fit\")\n",
    "\n",
    "train_loader = datamodule.train_dataloader()\n",
    "example_train_protein = datamodule.train[0]\n",
    "example_train_batch = next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have an example train data point (`example_train_protein`) but training is almost always done on batches of data points controlled by the `batch_size`. This batch size defines the number of input graphs looked at in each iteration of the training process (one forward and backward pass). It has a trade-off between the speed of the training and the generalizability of the model.\n",
    "\n",
    "In the graph neural network setting, a batch essentially combines all the graphs of the individual proteins into a bigger graph, with an additional batch attribute that specifies which protein each node belongs to. Since there are no edges between the different proteins, training on this batch graph is equivalent to training on the individual graphs separately, since no information flows between the different proteins. This is what is returned by the `train_dataloader` of the `DataModule`, in `example_train_batch`. \n",
    "\n",
    "Let's check what each of variables contain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_train_protein"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_train_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Move the `load_graph`, `ProteinDataset` and `ProteinGraphDataModule` functions and classes to `src/dataloader.py` so that we can use them in the next notebooks."
   ]
  }
 ],
 "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}