Merge pull request #2 from ChEB-AI/thesis_augmented_gnn

Thesis: Integration Chemical Knowledge into GNN

Author: sfluegel05

.gitignore

@@ -169,3 +169,8 @@ cython_debug/
 electra_pretrained.ckpt
 .isort.cfg
 /.vscode
+
+*.err
+*.out
+*.sh
+*.ckpt

chebai_graph/models/__init__.py

@@ -0,0 +1,19 @@
+from .augmented import (
+    GATAugNodePoolGraphPred,
+    GATGraphNodeFGNodePoolGraphPred,
+    ResGatedAugNodePoolGraphPred,
+    ResGatedGraphNodeFGNodePoolGraphPred,
+)
+from .dynamic_gni import ResGatedDynamicGNIGraphPred
+from .gat import GATGraphPred
+from .resgated import ResGatedGraphPred
+
+__all__ = [
+    "ResGatedGraphPred",
+    "ResGatedAugNodePoolGraphPred",
+    "ResGatedGraphNodeFGNodePoolGraphPred",
+    "GATGraphPred",
+    "GATAugNodePoolGraphPred",
+    "GATGraphNodeFGNodePoolGraphPred",
+    "ResGatedDynamicGNIGraphPred",
+]

chebai_graph/models/augmented.py

@@ -0,0 +1,45 @@
+from .base import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet
+from .gat import GATGraphPred
+from .resgated import ResGatedGraphPred
+
+
+class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred):
+    """
+    Combines:
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes).
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred):
+    """
+    Combines:
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes).
+    - GATGraphPred: Graph attention network for final graph prediction.
+    """
+
+    ...
+
+
+class ResGatedGraphNodeFGNodePoolGraphPred(
+    GraphNodeFGNodePoolingNet, ResGatedGraphPred
+):
+    """
+    Combines:
+    - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes).
+    - ResGatedGraphPred: Residual gated network for final graph prediction.
+    """
+
+    ...
+
+
+class GATGraphNodeFGNodePoolGraphPred(GraphNodeFGNodePoolingNet, GATGraphPred):
+    """
+    Combines:
+    - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes).
+    - GATGraphPred: Graph attention network for final graph prediction.
+    """
+
+    ...