demos/main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ToGePi: Topology and Geometry informed Positional Information\n",
    "\n",
    "**Authors**: Veljko Kovac, Gerard Planella, Adam Valin and Luca Pantea \\\n",
    "**Course**: Deep Learning 2, University of Amsterdam \\\n",
    "**Course Year**: 2023 \\\n",
    "**Course Website**: https://uvadl2c.github.io/ \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is meant to showcase the main experiments that our group worked on in the Deep Learning 2 course timeline. Our goal is to develop a generic method that also combines geometric and topological information by improving upon the established LSPE framework. By combining these distinct approaches, we seek to leverage the complementary nature of geometric and topological information in capturing complex graph relationships and enhancing the discriminative capabilities of GNN models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W&B offline. Running your script from this directory will only write metadata locally. Use wandb disabled to completely turn off W&B.\n"
     ]
    }
   ],
   "source": [
    "# Standard imports\n",
    "import os\n",
    "import torch\n",
    "\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "script_dir = os.path.abspath('')\n",
    "\n",
    "# Set this to false to see wandb output\n",
    "os.environ[\"WANDB_SILENT\"] = \"true\"\n",
    "\n",
    "# Set this to \"online\" for loggin metrics to cloud\n",
    "!wandb offline "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## About this notebook\n",
    "\n",
    "This notebook showcases the main experiments that we base our paper on. \\\n",
    "By integrating the **geometrical features** of the graph (node distances in the case of the QM9 dataset) with **topological features** given by PEs, we seek to achieve more expressive node attributes.\n",
    "\n",
    "We thus divide our approach into **two main directions**:\n",
    "\n",
    "**i)** **ToGePi-MPNN**: A method to **combine the LSPE method** while also making use of the **geometrical information found in EGNNs**, by taking relative absolute distances between nodes into account in the message function, and \n",
    "\n",
    "**ii)** **A study of MPNN Architectures**: For both _standard_ and _isotropic_ MPNNs, we conduct experiments with Baseline*, Geometry only, PE only, Geometry and PE, LSPE only, LSPE and Geometry. Formulas found in the table.\n",
    "\n",
    "*Baseline refers to the barebone implementation of MPNN models, i.e. no added topological, geometrical or learnt features. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We acknowledge that not anybody might have access to the required computational resources to train each of the models we tested, and thus we provide the saved model weights in the HuggingFace repository [here](https://huggingface.co/datasets/lucapantea/egnn-lspe/tree/main). We thus load each of the model weights saved during training in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model weights already initialized.\n"
     ]
    }
   ],
   "source": [
    "# Create saved_models dir\n",
    "saved_models_dir = os.path.join(os.path.dirname(script_dir), 'saved_models')\n",
    "if not os.path.exists(saved_models_dir):\n",
    "    os.makedirs(saved_models_dir)\n",
    "\n",
    "# Load model weights from hugging face\n",
    "saved_models_dir_git = r'\"{}\"'.format(saved_models_dir)\n",
    "if os.path.exists(saved_models_dir) and len(os.listdir(saved_models_dir)) == 0:\n",
    "    !git clone https://huggingface.co/datasets/lucapantea/egnn-lspe {saved_models_dir_git}\n",
    "else:\n",
    "    print('Model weights already initialized.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results and Analysis\n",
    "\n",
    "This section of the notebook corresponds to its homonymous counterpart in the blogpost/report, we will first examine how infusing the models with implicit topological information in the shape of Random Walk PEs (RWPE) affects their performance on the QM9 dataset in a fully connected (FC) and non-fully connected (NFC) setting. Moreover, we will demonstrate how geometrical information, the absolute distance between nodes, can be utilized effectively to learn\n",
    "better node embeddings.\n",
    "\n",
    "We make the run configurations, training and test metric,alongside with visualisations accessible via a WandB report [here](https://api.wandb.ai/links/dl2-gnn-lspe/uotynqoo). It is additionally displayed below via an embedded HTML element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe src=\"https://wandb.ai/dl2-gnn-lspe/dl2-modularized-exp/reports/EGNN-LSPE-Experiments--Vmlldzo0NDAyMjQ0\" style=\"border:none;height:1024px;width:100%\">"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb_visualizations_code = r'<iframe src=\"https://wandb.ai/dl2-gnn-lspe/dl2-modularized-exp/reports/EGNN-LSPE-Experiments--Vmlldzo0NDAyMjQ0\" style=\"border:none;height:1024px;width:100%\">'\n",
    "display(HTML(wandb_visualizations_code))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In case one would want to run evaluations for each of the individual experiments, please follow the run argument pattern outlined below, by specifying the model path under the \"evaluate\" argument. We showcase a two of examples below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Namespace(evaluate='mpnn_qm9_rw24_yes-lspe_no-dist_no-reduced_epochs-1000_num_layers-7_in_c-11_h_c-128_o_c-1_bs-96_lr-0.0005.pt', model='mpnn', dataset='qm9', pe='rw', pe_dim=24, lspe=True, seed=42, epochs=1000, batch_size=96, learning_rate=0.0005, weight_decay=1e-16, in_channels=11, hidden_channels=128, num_layers=7, out_channels=1, include_dist=False, reduced=False, update_with_pe=False)\n",
      "\n",
      "MPS available. Setting device to MPS.\n",
      "Number of parameters: 1670273\n",
      "\n",
      "Loading model with weights stored at /Users/luca/Documents/Masters/Deep Learning 2/LSPE-EGNN/demos/../saved_models/mpnn_qm9_rw24_yes-lspe_no-dist_no-reduced_epochs-1000_num_layers-7_in_c-11_h_c-128_o_c-1_bs-96_lr-0.0005.pt...\n",
      "\n",
      "Beginning evaluation...\n",
      "100%|█████████████████████████████████████████| 105/105 [00:15<00:00,  6.80it/s]\n",
      "\n",
      "Test MAE: 0.214\n",
      "Evaluation finished. Exiting...\n"
     ]
    }
   ],
   "source": [
    "# Standard MPNN, using PE and LSPE, yet without conditioning on distance.\n",
    "!python -W ignore ../main.py --evaluate \"mpnn_qm9_rw24_yes-lspe_no-dist_no-reduced_epochs-1000_num_layers-7_in_c-11_h_c-128_o_c-1_bs-96_lr-0.0005.pt\" \\\n",
    "                            --dataset \"qm9\" \\\n",
    "                            --pe \"rw\" \\\n",
    "                            --pe_dim 24 \\\n",
    "                            --lspe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Namespace(evaluate='mpnn_qm9_fc_nope_no-lspe_yes-dist_no-reduced_no-update_with_pe_epochs-1000_num_layers-7_in_c-11_h_c-128_o_c-1_bs-96_lr-0.0005.pt', model='mpnn', dataset='qm9_fc', pe='nope', pe_dim=24, lspe=False, seed=42, epochs=1000, batch_size=96, learning_rate=0.0005, weight_decay=1e-16, in_channels=11, hidden_channels=128, num_layers=7, out_channels=1, include_dist=True, reduced=False, update_with_pe=False)\n",
      "\n",
      "MPS available. Setting device to MPS.\n",
      "Number of parameters: 743809\n",
      "\n",
      "Loading model with weights stored at /Users/luca/Documents/Masters/Deep Learning 2/LSPE-EGNN/demos/../saved_models/mpnn_qm9_fc_nope_no-lspe_yes-dist_no-reduced_no-update_with_pe_epochs-1000_num_layers-7_in_c-11_h_c-128_o_c-1_bs-96_lr-0.0005.pt...\n",
      "\n",
      "Beginning evaluation...\n",
      "100%|█████████████████████████████████████████| 105/105 [00:23<00:00,  4.55it/s]\n",
      "\n",
      "Test MAE: 0.113\n",
      "Evaluation finished. Exiting...\n"
     ]
    }
   ],
   "source": [
    "# Standard MPNN operating on a fully connected dataset, without PE or LSPE, now with conditioning on distance.\n",
    "!python -W ignore ../main.py --evaluate \"mpnn_qm9_fc_nope_no-lspe_yes-dist_no-reduced_no-update_with_pe_epochs-1000_num_layers-7_in_c-11_h_c-128_o_c-1_bs-96_lr-0.0005.pt\" \\\n",
    "                            --dataset \"qm9_fc\" \\\n",
    "                            --pe \"nope\" \\\n",
    "                            --include_dist\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, to run a full experiment (train **and** evaluate), one can simply specify the parameters for each argument to the main function. Here is an overview of each of the possible arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "usage: Model runner. [-h] [--config S] [--write_config_to S] [--evaluate S]\r\n",
      "                     [--model S] [--dataset S] [--pe S] [--pe_dim N] [--lspe]\r\n",
      "                     [--seed N] [--epochs N] [--batch_size N]\r\n",
      "                     [--learning_rate N] [--weight_decay N] [--in_channels N]\r\n",
      "                     [--hidden_channels N] [--num_layers N] [--out_channels N]\r\n",
      "                     [--include_dist] [--reduced] [--update_with_pe]\r\n",
      "\r\n",
      "options:\r\n",
      "  -h, --help           show this help message and exit\r\n",
      "  --config S           Config file for parsing arguments. Command line\r\n",
      "                       arguments will be overriden.\r\n",
      "  --write_config_to S  Writes the current arguments as a json file for config\r\n",
      "                       with the specified filename.\r\n",
      "  --evaluate S         Directly evaluates the model with the model weightsof\r\n",
      "                       the path specified here. No need to specify the\r\n",
      "                       directory.\r\n",
      "  --model S            Available models: egnn | mpnn\r\n",
      "  --dataset S          Available datasets: qm9 | qm9_fc\r\n",
      "  --pe S               Available PEs: nope | rw | lap\r\n",
      "  --pe_dim N           PE dimension\r\n",
      "  --lspe               Whether or not to use LSPE framework. (default: False)\r\n",
      "  --seed N             Random seed\r\n",
      "  --epochs N           Number of epochs to train\r\n",
      "  --batch_size N       Batch size for training\r\n",
      "  --learning_rate N    Learning rate\r\n",
      "  --weight_decay N     clamp the output of the coords function if get too\r\n",
      "                       large\r\n",
      "  --in_channels N      Input dimension of features\r\n",
      "  --hidden_channels N  Hidden dimensions\r\n",
      "  --num_layers N       Number of model layers\r\n",
      "  --out_channels N     Output dimensions\r\n",
      "  --include_dist       Whether or not to include distance in the message\r\n",
      "                       state. (default: False)\r\n",
      "  --reduced            Whether or not to used the reduced version. (default:\r\n",
      "                       False)\r\n",
      "  --update_with_pe     Whether or not to include pe in the update network.\r\n",
      "                       (default: False)\r\n"
     ]
    }
   ],
   "source": [
    "!python ../main.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Individual Contributions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Team member      | Contributions |\n",
    "| :-               | :-            | \n",
    "| Veljko Kovac     | Throughout our group project, we initially attempted to divide the workload evenly among team members. However, we soon realized that achieving true equality in task distribution was challenging. Consequently, we decided to collaborate closely, working together on most aspects of the project. Initially, I assisted in shaping the code with PyTorch Geometric, using the tutorial code on EGNN as a starting point. While the performance did not match that of the original paper, I acquired valuable skills in utilizing this framework for faster implementations. To rule out potential issues with PyTorch Geometric functions, I also implemented the code using the standard PyTorch framework. After making the code of EGNN to work with my teammates, our focus was to experiment with various combinations of methods suggested by our TA. If I were to identify my biggest contribution in this project, that would be the mathematical derivations and their respective coding that we had to try. In addition to the approaches we pursued within the project's scope, I was keen to explore alternative directions, such as leveraging the differences between positional encodings of two nodes. Although these experiments lacked a solid theoretical foundation, the results were almost identical with the original EGNN, suggesting the potential value of further investigation. Furthermore, the writing of the project report was a collaborative effort, with all team members sitting together and contributing to its completion. Overall, this project provided me with valuable learning experiences and great enjoyment. We plan to continue working on it, with the potential goal of refining it for publication.\n",
    "| Gerard Planella  | I like that we all had the opportunity to work in everything. I started my work by adapting the code from the EGNN in PyG and PyTorch Lightning to work on the QM9 dataset, while also making a pipeline for easily configuring the different model parameters and having a modular code to easily expand to different datasets and models (some predicted quantities are categorical). I also worked on the integration with WandB. Sadly, the PyG model did not give a good performance, we all tried fixing it and making it have the same performance as Floor's implementation but we finally decided it was more efficient to leave PyG and PyTorch Lighting. Once we had the code working with a good baseline performance, I helped in coming up and coding the equations for the different models used. We all had to run many different tests, which can be observed from our WandB logs. After correcting many errors and aligning with our TA in the direction we wanted to go in, we were able to perform all the runs required for the project. While our GPUs were working on that, I continued writing the report with my teammates, analysing the obtained results, discussing the next steps to take and writing them in the report. I also worked on converting our report to a blogpost.\n",
    "| Adam Valin       | As everybody on the project, I was happy to work on a little bit of everything, from the implementation of LSPE to the EGNN architecture with Pytorch Geometric, to the state-of-the-art review and writing of the report. I especially spent a lot of time implementing and making the architecture work, first with Pytorch Geometric then later with standard Pytorch by coding the different parts of the MPNNs. This later part was really instructive as implementing is only possible when the equations are understood. On top of that, I obviously did use my Lisa account to run and evaluate our run models and participated actively in the writing of the report. Finally, I also worked on converting on the blogpost, which is a modified version of our report. \n",
    "| Luca Pantea      | During the course of the project, my primary focus was on making contributions that significantly improved our workflow and the overall effectiveness of our work. One notable contribution involved the development of an end-to-end method that greatly simplified the exploration of various model configurations. This streamlined approach not only expedited last-minute experiments but also reduced the time and effort required for designing, experimenting with, and evaluating new model setups. To ensure the reproducibility of our work, I took the initiative to create a HuggingFace repository where we made our trained model weights available to the public. Additionally, I contributed to the project's accessibility by creating a report on WandB, which allowed others to easily access our resources and explore our experimental runs. Beyond the project's primary objectives, I also delved into a detailed exploration of the inner workings of our chosen topological features, specifically the Random Walk Positional Encoding (RWPE) vectors. This led me to dedicate a dedicated notebook to extensively study the qm9 dataset and visualize the RWPE vectors assigned to each atom in a molecule. These efforts significantly deepened my understanding of its underlying mechanisms, and why it works so well in practice. Furthermore, I actively participated in the writing process alongside my colleagues, with a focus on developing sections such as the introduction, related works, and experiments. In addition to that, I worked on distilling our report into a poster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 ('egnn_lspe')",
   "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.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "a481cf1fb4b2476369d538cfa7febb4501776dad9ca5bbc0e6e01e506cc16df9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}