num_undir_edges key fix

Author: aditya0by0

chebai_graph/preprocessing/properties/augmented_properties.py

@@ -53,7 +53,7 @@ def get_property_value(self, augmented_mol: Dict) -> List:
             [self.get_bond_value(bond) for bond in fg_graph_node_edges.values()]
         )
 
-        num_directed_edges = augmented_mol[self.MAIN_KEY]["num_undirected_edges"] // 2
+        num_directed_edges = augmented_mol[self.MAIN_KEY][NUM_EDGES] // 2
         assert (
             len(prop_list) == num_directed_edges
         ), f"Number of property values ({len(prop_list)}) should be equal to number of half the number of undirected edges i.e. must be equal to {num_directed_edges} "

chebai_graph/preprocessing/properties/constants.py

@@ -10,3 +10,4 @@
 ATOM_FG_EDGE = "atom_fg_lvl"
 FG_GRAPHNODE_EDGE = "fg_graphNode_lvl"
 EDGE_LEVELS = {WITHIN_ATOMS_EDGE, WITHIN_FG_EDGE, ATOM_FG_EDGE, FG_GRAPHNODE_EDGE}
+NUM_EDGES = "num_undirected_edges"

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -181,7 +181,7 @@ def _read_data(self, smiles: str) -> GeomData | None:
 
         # Empty features initialized; node and edge features can be added later
         x = torch.zeros((augmented_molecule["nodes"]["num_nodes"], 0))
-        edge_attr = torch.zeros((augmented_molecule["edges"]["num_edges"], 0))
+        edge_attr = torch.zeros((augmented_molecule["edges"][NUM_EDGES], 0))
 
         assert (
             edge_index.shape[0] == 2
@@ -297,7 +297,7 @@ def _augment_graph_structure(
             ATOM_FG_EDGE: atom_fg_edges,
             WITHIN_FG_EDGE: internal_fg_edges,
             FG_GRAPHNODE_EDGE: fg_to_graph_edges,
-            "num_undirected_edges": self._num_of_edges * 2,  # Undirected edges
+            NUM_EDGES: self._num_of_edges * 2,  # Undirected edges
         }
         return undirected_edge_index, node_info, edge_info
 
@@ -449,12 +449,12 @@ def _construct_fg_level_structure(
                 source_fg is not None and target_fg is not None
             ), "Each bond should have a fg node on both end"
 
-            internal_edge_index[0].append(source_fg)
-            internal_edge_index[1].append(target_fg)
             edge_str = f"{source_fg}_{target_fg}"
             if edge_str not in internal_fg_edges:
                 # If two atoms of a FG points to atom(s) belonging to another FG. In this case, only one edge is counted.
                 # Eg. In CHEBI:52723, atom idx 13 and 16 of a FG points to atom idx 18 of another FG
+                internal_edge_index[0].append(source_fg)
+                internal_edge_index[1].append(target_fg)
                 internal_fg_edges[edge_str] = {EDGE_LEVEL: WITHIN_FG_EDGE}
                 self._num_of_edges += 1