add graph node mask

Author: aditya0by0

chebai_graph/preprocessing/datasets/chebi.py

@@ -182,12 +182,16 @@ def _merge_props_into_base(self, row):
         is_atom_node = (
             geom_data.is_atom_node if hasattr(geom_data, "is_atom_node") else None
         )
+        is_graph_node = (
+            geom_data.is_graph_node if hasattr(geom_data, "is_graph_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,
+            is_graph_node=is_graph_node,
         )
 
     def load_processed_data_from_file(self, filename):

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -205,9 +205,15 @@ def _read_data(self, smiles: str) -> GeomData | None:
         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
+        is_graph_node = torch.zeros(NUM_NODES, dtype=torch.bool)
+        is_graph_node[-1] = True
 
         return GeomData(
-            x=x, edge_index=edge_index, edge_attr=edge_attr, is_atom_node=is_atom_mask
+            x=x,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            is_atom_node=is_atom_mask,
+            is_graph_node=is_graph_node,
         )
 
     def _create_augmented_graph(self, mol: Chem.Mol) -> Tuple[torch.Tensor, dict]: