GATv2 amg pool

Author: aditya0by0

chebai_graph/models/__init__.py

@@ -1,29 +1,17 @@
 from .augmented import (
     GATAugNodePoolGraphPred,
+    GATGraphNodeFGNodePoolGraphPred,
     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",
+    "GATGraphPred",
     "GATAugNodePoolGraphPred",
+    "GATGraphNodeFGNodePoolGraphPred",
 ]

chebai_graph/models/augmented.py

@@ -1,22 +1,12 @@
-from .base import (
-    AugmentedNodePoolingNet,
-    AugmentedOnlyPoolingNet,
-    FGNodePoolingNet,
-    FGNodePoolingNoGraphNodeNet,
-    FGOnlyPoolingNet,
-    GraphNodeFGNodePoolingNet,
-    GraphNodeNoFGNodePoolingNet,
-    GraphNodeOnlyPoolingNet,
-    GraphNodePoolingNet,
-)
+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 with molecule attributes.
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes).
     - ResGatedGraphPred: Residual gated network for final graph prediction.
     """
 
@@ -26,94 +16,30 @@ class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred):
 class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred):
     """
     Combines:
-    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings with molecule attributes.
+    - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally 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.
+    - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes).
     - ResGatedGraphPred: Residual gated network for final graph prediction.
     """
 
     ...
 
 
-class ResGatedGraphNodeOnlyPoolGraphPred(GraphNodeOnlyPoolingNet, ResGatedGraphPred):
+class GATGraphNodeFGNodePoolGraphPred(GraphNodeFGNodePoolingNet, GATGraphPred):
     """
     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.
+    - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes).
+    - GATGraphPred: Graph attention network for final graph prediction.
     """
 
     ...

chebai_graph/models/base.py

@@ -200,7 +200,7 @@ def _get_lin_seq_input_dim(
 
         Includes:
         - Atom embeddings
-        - Molecular attributes
+        - Molecular attributes (if any)
         - Augmented node embeddings
 
         Args:

chebai_graph/models/gat.py

@@ -40,7 +40,7 @@ def __init__(self, config: dict, **kwargs):
             heads=self.heads,
             v2=self.v2,
             act=self.activation,
-            share_weights=self.share_weights
+            share_weights=self.share_weights,
         )
 
     def forward(self, batch: dict) -> torch.Tensor: