seperate pooling for atom and augmented nodes

One vector: average of atom embeddings
One vector: average of augmented node embeddings
#2 (comment)

Author: aditya0by0

chebai_graph/models/__init__.py

@@ -1,3 +1,4 @@
 from ._gat import GATModelWrapper
+from .graph import ResGatedAugmentedGraphPred
 
-__all__ = ["GATModelWrapper"]
+__all__ = ["GATModelWrapper", "ResGatedAugmentedGraphPred"]

chebai_graph/models/graph.py

@@ -188,6 +188,68 @@ def forward(self, batch):
         return a
 
 
+class ResGatedAugmentedGraphPred(GraphBaseNet):
+    """GNN for graph-level prediction for augmented graphs"""
+
+    NAME = "ResGatedAugmentedGraphPred"
+
+    def __init__(
+        self,
+        config: typing.Dict,
+        n_linear_layers=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.gnn = ResGatedGraphConvNetBase(config, **kwargs)
+        self.linear_layers = torch.nn.ModuleList(
+            [
+                torch.nn.Linear(
+                    self.gnn.hidden_length
+                    + (i == 0) * self.gnn.n_molecule_properties
+                    + (i == 0) * self.gnn.hidden_length,
+                    self.gnn.hidden_length,
+                )
+                for i in range(n_linear_layers - 1)
+            ]
+        )
+        self.final_layer = nn.Linear(self.gnn.hidden_length, self.out_dim)
+
+    def forward(self, batch):
+        graph_data = batch["features"][0]
+        assert isinstance(graph_data, GraphData)
+        is_atom_node = graph_data.is_atom_node.bool()  # Boolean mask: shape [num_nodes]
+        is_augmented_node = ~is_atom_node
+
+        node_embeddings = self.gnn(batch)
+
+        atom_embeddings = node_embeddings[is_atom_node]
+        atom_batch = graph_data.batch[is_atom_node]
+
+        augmented_node_embeddings = node_embeddings[is_augmented_node]
+        augmented_node_batch = graph_data.batch[is_augmented_node]
+
+        # Scatter add separately
+        graph_vec_atoms = scatter_add(atom_embeddings, atom_batch, dim=0)
+        graph_vec_augmented_nodes = scatter_add(
+            augmented_node_embeddings, augmented_node_batch, dim=0
+        )
+
+        # Concatenate both
+        graph_vector = torch.cat(
+            [
+                graph_vec_atoms,
+                graph_data.molecule_attr,
+                graph_vec_augmented_nodes,
+            ],
+            dim=1,
+        )
+
+        for lin in self.linear_layers:
+            a = self.gnn.activation(lin(graph_vector))
+        a = self.final_layer(a)
+        return a
+
+
 class ResGatedGraphConvNetPretrain(GraphBaseNet):
     """For pretraining. BaseNet with an additional output layer for predicting atom properties"""
 

chebai_graph/preprocessing/datasets/chebi.py

@@ -178,11 +178,16 @@ def _merge_props_into_base(self, row):
                 )
             else:
                 molecule_attr = torch.cat([molecule_attr, property_values], dim=1)
+
+        is_atom_node = (
+            geom_data.is_atom_node if hasattr(geom_data, "is_atom_node") else None
+        )
         return GeomData(
             x=x,
             edge_index=geom_data.edge_index,
             edge_attr=edge_attr,
             molecule_attr=molecule_attr,
+            is_atom_node=is_atom_node,
         )
 
     def load_processed_data_from_file(self, filename):

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -1,3 +1,4 @@
+import sys
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional, Tuple
 
@@ -179,7 +180,12 @@ def _read_data(self, smiles: str) -> GeomData | None:
         self.mol_object_buffer[smiles] = augmented_molecule
 
         # Empty features initialized; node and edge features can be added later
-        x = torch.zeros((augmented_molecule["nodes"]["num_nodes"], 0))
+        NUM_NODES = augmented_molecule["nodes"]["num_nodes"]
+        assert (
+            NUM_NODES is not None and NUM_NODES > 1
+        ), "Num of nodes in augmented graph should be more than 1"
+
+        x = torch.zeros((NUM_NODES, 0))
         edge_attr = torch.zeros((augmented_molecule["edges"][k.NUM_EDGES], 0))
 
         assert (
@@ -194,7 +200,14 @@ def _read_data(self, smiles: str) -> GeomData | None:
             len(set(edge_index[0].tolist())) == x.shape[0]
         ), f"Number of unique source nodes in edge_index ({len(set(edge_index[0].tolist()))}) does not match number of nodes in x ({x.shape[0]})"
 
-        return GeomData(x=x, edge_index=edge_index, edge_attr=edge_attr)
+        # Create a boolean mask: True for atom, False for augmented
+        is_atom_mask = torch.zeros(NUM_NODES, dtype=torch.bool)
+        NUM_ATOM_NODES = augmented_molecule["nodes"]["atom_nodes"].GetNumAtoms()
+        is_atom_mask[:NUM_ATOM_NODES] = True
+
+        return GeomData(
+            x=x, edge_index=edge_index, edge_attr=edge_attr, is_atom_node=is_atom_mask
+        )
 
     def _create_augmented_graph(self, mol: Chem.Mol) -> Tuple[torch.Tensor, dict]:
         """
@@ -267,6 +280,14 @@ def _augment_graph_structure(
         assert (
             self._num_of_nodes == total_atoms
         ), f"Mismatch in number of nodes: expected {total_atoms}, got {self._num_of_nodes}"
+        assert sys.version_info >= (
+            3,
+            7,
+        ), "This code requires Python 3.7 or higher."
+        # For python 3.7+, the standard dict type preserves insertion order, and is iterated over in same order
+        # https://docs.python.org/3/whatsnew/3.7.html#summary-release-highlights
+        # https://mail.python.org/pipermail/python-dev/2017-December/151283.html
+        # Order preservation is necessary to to create `is_atom_node` mask
         node_info = {
             "atom_nodes": mol,
             "fg_nodes": fg_nodes,

configs/model/gnn_resgated_aug.yml

@@ -0,0 +1,13 @@
+class_path: chebai_graph.models.ResGatedAugmentedGraphPred
+init_args:
+  optimizer_kwargs:
+    lr: 1e-3
+  config:
+    in_length: 256
+    hidden_length: 512
+    dropout_rate: 0.1
+    n_conv_layers: 3
+    n_linear_layers: 3
+    n_atom_properties: 158
+    n_bond_properties: 7
+    n_molecule_properties: 200