Thesis: Integration Chemical Knowledge into GNN

.github/workflows/lint.yml

@@ -0,0 +1,26 @@
+name: Lint
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'  # or any version your project uses
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install black==25.1.0 ruff==0.12.2
+
+      - name: Run Black
+        run: black --check .
+
+      - name: Run Ruff (no formatting)
+        run: ruff check . --no-fix

.pre-commit-config.yaml

@@ -1,9 +1,31 @@
 repos:
-#-   repo: https://github.com/PyCQA/isort
-#    rev: "5.12.0"
-#    hooks:
-#    -   id: isort
 -   repo: https://github.com/psf/black
-    rev: "22.10.0"
+    rev: "25.1.0"
     hooks:
-    -   id: black
+    -   id: black
+    -   id: black-jupyter # for formatting jupyter-notebook
+
+-   repo: https://github.com/pycqa/isort
+    rev: 5.13.2
+    hooks:
+      - id: isort
+        name: isort (python)
+        args: ["--profile=black"]
+
+-   repo: https://github.com/asottile/seed-isort-config
+    rev: v2.2.0
+    hooks:
+    -   id: seed-isort-config
+
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.2
+    hooks:
+    -   id: ruff
+        args: [--fix]

README.md

@@ -7,7 +7,7 @@ Graph-based models for molecular property prediction and ontology classification
 
 ## Installation
 
-To install this repository, download it and run 
+To install this repository, download it and run
 
 ```bash
 pip install .
@@ -25,7 +25,7 @@ Replace:
 - `${TORCH}` with your installed PyTorch version (e.g., `2.6.0`)
 - `${CUDA}` with e.g. `cpu`, `cu118`, or `cu121` depending on your system and CUDA version
 
-If you already have `torch` installed, make sure that `torch_scatter` and `torch_geometric` are compatible with your 
+If you already have `torch` installed, make sure that `torch_scatter` and `torch_geometric` are compatible with your
 PyTorch version and are installed with the same CUDA version.
 
 For a full list of currently available PyTorch versions and CUDA compatibility, please refer to libraries' official documentation:
@@ -63,11 +63,10 @@ my_projects/
 ### Ontology Prediction
 
 
-This example command trains a Residual Gated Graph Convolutional Network on the ChEBI50 dataset (see [wiki](https://github.com/ChEB-AI/python-chebai/wiki/Data-Management)). 
-The dataset has a customizable list of properties for atoms, bonds and molecules that are added to the graph. 
+This example command trains a Residual Gated Graph Convolutional Network on the ChEBI50 dataset (see [wiki](https://github.com/ChEB-AI/python-chebai/wiki/Data-Management)).
+The dataset has a customizable list of properties for atoms, bonds and molecules that are added to the graph.
 The list can be found in the `configs/data/chebi50_graph_properties.yml` file.
 
 ```bash
 python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/csv_logger.yml --model=../python-chebai-graph/configs/model/gnn_res_gated.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_graph_properties.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce.yml
 ```
-

chebai_graph/models/__init__.py

@@ -0,0 +1,29 @@
+from .augmented import (
+    GATAugNodePoolGraphPred,
+    ResGatedAugNodePoolGraphPred,
+    ResGatedAugOnlyPoolGraphPred,
+    ResGatedFGNodeNoGraphNodeGraphPred,
+    ResGatedFGNodePoolGraphPred,
+    ResGatedFGOnlyPoolGraphPred,
+    ResGatedGraphNodeFGNodePoolGraphPred,
+    ResGatedGraphNodeNoFGNodeGraphPred,
+    ResGatedGraphNodeOnlyPoolGraphPred,
+    ResGatedGraphNodePoolGraphPred,
+)
+from .gat import GATGraphPred
+from .resgated import ResGatedGraphPred
+
+__all__ = [
+    "GATGraphPred",
+    "ResGatedGraphPred",
+    "ResGatedFGNodeNoGraphNodeGraphPred",
+    "ResGatedAugNodePoolGraphPred",
+    "ResGatedGraphNodeFGNodePoolGraphPred",
+    "ResGatedGraphNodePoolGraphPred",
+    "ResGatedGraphNodeNoFGNodeGraphPred",
+    "ResGatedFGNodePoolGraphPred",
+    "ResGatedAugOnlyPoolGraphPred",
+    "ResGatedGraphNodeOnlyPoolGraphPred",
+    "ResGatedFGOnlyPoolGraphPred",
+    "GATAugNodePoolGraphPred",
+]

chebai_graph/models/augmented.py

@@ -0,0 +1,119 @@
+from .base import (
+    AugmentedNodePoolingNet,
+    AugmentedOnlyPoolingNet,
+    FGNodePoolingNet,
+    FGNodePoolingNoGraphNodeNet,
+    FGOnlyPoolingNet,
+    GraphNodeFGNodePoolingNet,
+    GraphNodeNoFGNodePoolingNet,
+    GraphNodeOnlyPoolingNet,
+    GraphNodePoolingNet,
+)
+from .gat import GATGraphPred
+from .resgated import ResGatedGraphPred
+
+
+class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred):
+    """
+    Combines:
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings with molecule attributes.
+    - GATGraphPred: Graph attention network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedGraphNodePoolGraphPred(GraphNodePoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - GraphNodePoolingNet: Pools atom and graph node embeddings with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedFGNodePoolGraphPred(FGNodePoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - FGNodePoolingNet: Pools functional group nodes and other nodes with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedGraphNodeFGNodePoolGraphPred(
+    GraphNodeFGNodePoolingNet, ResGatedGraphPred
+):
+    """
+    Combines:
+    - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedGraphNodeNoFGNodeGraphPred(
+    GraphNodeNoFGNodePoolingNet, ResGatedGraphPred
+):
+    """
+    Combines:
+    - GraphNodeNoFGNodePoolingNet: Pools atom and graph nodes, excluding functional groups.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedFGNodeNoGraphNodeGraphPred(
+    FGNodePoolingNoGraphNodeNet, ResGatedGraphPred
+):
+    """
+    Combines:
+    - FGNodePoolingNoGraphNodeNet: Pools atom and functional group nodes, excluding graph nodes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedAugOnlyPoolGraphPred(AugmentedOnlyPoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - AugmentedOnlyPoolingNet: Pools only augmented nodes with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedGraphNodeOnlyPoolGraphPred(GraphNodeOnlyPoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - GraphNodeOnlyPoolingNet: Pools only graph nodes with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedFGOnlyPoolGraphPred(FGOnlyPoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - FGOnlyPoolingNet: Pools only functional group nodes with molecule attributes.
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...