adapt aug props for ablation data classes

Author: aditya0by0

chebai_graph/preprocessing/properties/augmented_properties.py

@@ -1,3 +1,4 @@
+import sys
 from abc import ABC
 from typing import Dict, List, Optional
 
@@ -13,6 +14,15 @@
 from . import properties as pr
 from .base import AtomProperty, BondProperty, FrozenPropertyAlias
 
+# 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
+assert sys.version_info >= (
+    3,
+    7,
+), "This code requires Python 3.7 or higher."
+# Order preservation is necessary to to create `prop_list`
+
 
 # --------------------- Atom Properties -----------------------------
 class AugmentedAtomProperty(AtomProperty, ABC):
@@ -24,9 +34,7 @@ def get_property_value(self, augmented_mol: Dict):
                 f"Key `{self.MAIN_KEY}` should be present in augmented molecule dict"
             )
 
-        missing_keys = {"atom_nodes", "fg_nodes", "graph_node"} - augmented_mol[
-            self.MAIN_KEY
-        ].keys()
+        missing_keys = {"atom_nodes"} - augmented_mol[self.MAIN_KEY].keys()
         if missing_keys:
             raise KeyError(f"Missing keys {missing_keys} in augmented molecule nodes")
 
@@ -35,26 +43,29 @@ def get_property_value(self, augmented_mol: Dict):
             raise TypeError(
                 f'augmented_mol["{self.MAIN_KEY}"]["atom_nodes"] must be an instance of rdkit.Chem.Mol'
             )
-
         prop_list = [self.get_atom_value(atom) for atom in atom_molecule.GetAtoms()]
 
-        fg_nodes = augmented_mol[self.MAIN_KEY]["fg_nodes"]
-        graph_node = augmented_mol[self.MAIN_KEY]["graph_node"]
-        if not isinstance(fg_nodes, dict) or not isinstance(graph_node, dict):
-            raise TypeError(
-                f'augmented_mol["{self.MAIN_KEY}"](["fg_nodes"]/["graph_node"]) must be an instance of dict '
-                f"containing its properties"
-            )
+        if "fg_nodes" in augmented_mol[self.MAIN_KEY]:
+            fg_nodes = augmented_mol[self.MAIN_KEY]["fg_nodes"]
+            if not isinstance(fg_nodes, dict):
+                raise TypeError(
+                    f'augmented_mol["{self.MAIN_KEY}"](["fg_nodes"]) must be an instance of dict '
+                    f"containing its properties"
+                )
+            prop_list.extend([self.get_atom_value(atom) for atom in fg_nodes.values()])
+
+        if "graph_node" in augmented_mol[self.MAIN_KEY]:
+            graph_node = augmented_mol[self.MAIN_KEY]["graph_node"]
+            if not isinstance(graph_node, dict):
+                raise TypeError(
+                    f'augmented_mol["{self.MAIN_KEY}"](["graph_node"]) must be an instance of dict '
+                    f"containing its properties"
+                )
+            prop_list.append(self.get_atom_value(graph_node))
 
-        # 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
-        prop_list.extend([self.get_atom_value(atom) for atom in fg_nodes.values()])
-        prop_list.append(self.get_atom_value(graph_node))
         assert (
             len(prop_list) == augmented_mol[self.MAIN_KEY]["num_nodes"]
         ), "Number of property values should be equal to number of nodes"
-
         return prop_list
 
     def _check_modify_atom_prop_value(self, atom: Chem.rdchem.Atom | Dict, prop: str):
@@ -228,7 +239,7 @@ def get_property_value(self, augmented_mol: Dict) -> List:
                 f"Key `{self.MAIN_KEY}` should be present in augmented molecule dict"
             )
 
-        missing_keys = k.EDGE_LEVELS - augmented_mol[self.MAIN_KEY].keys()
+        missing_keys = {k.WITHIN_ATOMS_EDGE} - augmented_mol[self.MAIN_KEY].keys()
         if missing_keys:
             raise KeyError(f"Missing keys {missing_keys} in augmented molecule nodes")
 
@@ -237,31 +248,38 @@ def get_property_value(self, augmented_mol: Dict) -> List:
             raise TypeError(
                 f'augmented_mol["{self.MAIN_KEY}"]["{k.WITHIN_ATOMS_EDGE}"] must be an instance of rdkit.Chem.Mol'
             )
-
         prop_list = [self.get_bond_value(bond) for bond in atom_molecule.GetBonds()]
 
-        fg_atom_edges = augmented_mol[self.MAIN_KEY][k.ATOM_FG_EDGE]
-        fg_edges = augmented_mol[self.MAIN_KEY][k.WITHIN_FG_EDGE]
-        fg_graph_node_edges = augmented_mol[self.MAIN_KEY][k.FG_GRAPHNODE_EDGE]
-
-        if (
-            not isinstance(fg_atom_edges, dict)
-            or not isinstance(fg_edges, dict)
-            or not isinstance(fg_graph_node_edges, dict)
-        ):
-            raise TypeError(
-                f'augmented_mol["{self.MAIN_KEY}"](["{k.ATOM_FG_EDGE}"]/["{k.WITHIN_FG_EDGE}"]/["{k.FG_GRAPHNODE_EDGE}"]) '
-                f"must be an instance of dict containing its properties"
+        if k.ATOM_FG_EDGE in augmented_mol[self.MAIN_KEY]:
+            fg_atom_edges = augmented_mol[self.MAIN_KEY][k.ATOM_FG_EDGE]
+            if not isinstance(fg_atom_edges, dict):
+                raise TypeError(
+                    f"augmented_mol['{self.MAIN_KEY}'](['{k.ATOM_FG_EDGE}'])"
+                    f"must be an instance of dict containing its properties"
+                )
+            prop_list.extend(
+                [self.get_bond_value(bond) for bond in fg_atom_edges.values()]
             )
 
-        # 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
-        prop_list.extend([self.get_bond_value(bond) for bond in fg_atom_edges.values()])
-        prop_list.extend([self.get_bond_value(bond) for bond in fg_edges.values()])
-        prop_list.extend(
-            [self.get_bond_value(bond) for bond in fg_graph_node_edges.values()]
-        )
+        if k.WITHIN_FG_EDGE in augmented_mol[self.MAIN_KEY]:
+            fg_edges = augmented_mol[self.MAIN_KEY][k.WITHIN_FG_EDGE]
+            if not isinstance(fg_edges, dict):
+                raise TypeError(
+                    f"augmented_mol['{self.MAIN_KEY}'](['{k.WITHIN_FG_EDGE}'])"
+                    f"must be an instance of dict containing its properties"
+                )
+            prop_list.extend([self.get_bond_value(bond) for bond in fg_edges.values()])
+
+        if k.TO_GRAPHNODE_EDGE in augmented_mol[self.MAIN_KEY]:
+            fg_graph_node_edges = augmented_mol[self.MAIN_KEY][k.TO_GRAPHNODE_EDGE]
+            if not isinstance(fg_graph_node_edges, dict):
+                raise TypeError(
+                    f"augmented_mol['{self.MAIN_KEY}'](['{k.TO_GRAPHNODE_EDGE}'])"
+                    f"must be an instance of dict containing its properties"
+                )
+            prop_list.extend(
+                [self.get_bond_value(bond) for bond in fg_graph_node_edges.values()]
+            )
 
         num_directed_edges = augmented_mol[self.MAIN_KEY][k.NUM_EDGES] // 2
         assert (

chebai_graph/preprocessing/properties/constants.py

@@ -8,6 +8,6 @@
 WITHIN_ATOMS_EDGE = "within_atoms_lvl"
 WITHIN_FG_EDGE = "within_fg_lvl"
 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}
+TO_GRAPHNODE_EDGE = "fg_graphNode_lvl"
+EDGE_LEVELS = {WITHIN_ATOMS_EDGE, WITHIN_FG_EDGE, ATOM_FG_EDGE, TO_GRAPHNODE_EDGE}
 NUM_EDGES = "num_undirected_edges"

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -598,7 +598,7 @@ def _add_graph_node_and_edges_to_nodes(
             self._construct_nodes_to_graph_node_structure(nodes_ids)
         )
 
-        augmented_struct["edge_info"][k.FG_GRAPHNODE_EDGE] = nodes_to_graph_edges
+        augmented_struct["edge_info"][k.TO_GRAPHNODE_EDGE] = nodes_to_graph_edges
         augmented_struct["edge_info"][k.NUM_EDGES] += len(nodes_to_graph_edges)
         assert (
             self._idx_of_edge == augmented_struct["edge_info"][k.NUM_EDGES]
@@ -643,7 +643,7 @@ def _construct_nodes_to_graph_node_structure(
             graph_edge_index[0].append(self._idx_of_node)
             graph_edge_index[1].append(fg_id)
             graph_to_nodes_edges[f"{self._idx_of_node}_{fg_id}"] = {
-                k.EDGE_LEVEL: k.FG_GRAPHNODE_EDGE
+                k.EDGE_LEVEL: k.TO_GRAPHNODE_EDGE
             }
             self._idx_of_edge += 1
         self._idx_of_node += 1

chebai_graph/preprocessing/utils/visualize_augmented_molecule.py

@@ -25,7 +25,7 @@
     k.WITHIN_ATOMS_EDGE: "#1f77b4",
     k.ATOM_FG_EDGE: "#9467bd",
     k.WITHIN_FG_EDGE: "#ff7f0e",
-    k.FG_GRAPHNODE_EDGE: "#2ca02c",
+    k.TO_GRAPHNODE_EDGE: "#2ca02c",
 }
 
 NODE_COLOR_MAP = {
@@ -119,8 +119,8 @@ def _create_graph(
         else set()
     )
     fg_graph_edges = (
-        set(augmented_graph_edges[k.FG_GRAPHNODE_EDGE])
-        if k.FG_GRAPHNODE_EDGE in augmented_graph_edges
+        set(augmented_graph_edges[k.TO_GRAPHNODE_EDGE])
+        if k.TO_GRAPHNODE_EDGE in augmented_graph_edges
         else set()
     )
 
@@ -133,7 +133,7 @@ def _create_graph(
         elif undirected_edge_set & within_fg_edges:
             edge_type = k.WITHIN_FG_EDGE
         elif undirected_edge_set & fg_graph_edges:
-            edge_type = k.FG_GRAPHNODE_EDGE
+            edge_type = k.TO_GRAPHNODE_EDGE
         else:
             raise ValueError("Unexpected edge type")
         G.add_edge(src, tgt, edge_type=edge_type, edge_color=EDGE_COLOR_MAP[edge_type])
@@ -318,7 +318,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None:
         k.WITHIN_ATOMS_EDGE: [],
         k.ATOM_FG_EDGE: [],
         k.WITHIN_FG_EDGE: [],
-        k.FG_GRAPHNODE_EDGE: [],
+        k.TO_GRAPHNODE_EDGE: [],
     }
     for src, tgt, data in G.edges(data=True):
         edge_type_to_edges[data["edge_type"]].append((src, tgt))