chebai-graph 1.0.0__tar.gz

chebai_graph-1.0.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: chebai-graph
+Version: 1.0.0
+Summary: GNNs for ChEB-AI
+Author-email: Martin Glauer <martin.glauer@ovgu.de>
+Requires-Python: >=3.8
+Requires-Dist: chebai
+Requires-Dist: descriptastorus
+Requires-Dist: tox ; extra == "dev"
+Requires-Dist: isort ; extra == "linters"
+Requires-Dist: pre-commit ; extra == "linters"
+Requires-Dist: black ; extra == "linters"
+Provides-Extra: dev
+Provides-Extra: linters

chebai_graph-1.0.0/chebai_graph/loss/pretraining.py

@@ -0,0 +1,31 @@
+import torch
+
+
+class MaskPretrainingLoss(torch.nn.Module):
+    # Mask atoms and edges, try to predict them (see Hu et al., 2020: Strategies for Pre-training Graph Neural Networks)
+    def __init__(self):
+        super().__init__()
+        self.ce = torch.nn.functional.binary_cross_entropy_with_logits
+
+    def forward(self, input, target, **loss_kwargs):
+        if isinstance(input, tuple):
+            atom_preds, bond_preds = input
+            atom_targets, bond_targets = target
+            try:
+                bond_loss = self.ce(bond_preds, bond_targets)
+            except RuntimeError as e:
+                print(f"Failed to compute bond loss: {e}")
+                print(f"Input: preds: {bond_preds.shape}, labels: {bond_targets.shape}")
+                bond_loss = 0
+        else:
+            atom_preds = input
+            atom_targets = target
+            bond_loss = 0
+        try:
+            atom_loss = self.ce(atom_preds, atom_targets)
+        except RuntimeError as e:
+            print(f"Failed to compute atom loss: {e}")
+            print(f"Input: preds: {atom_preds.shape}, labels: {atom_targets.shape}")
+            atom_loss = 0
+
+        return atom_loss + bond_loss

chebai_graph-1.0.0/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-1.0.0/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.
+    """
+
+    ...