add usual prop for augmented graph with default value

Author: aditya0by0

chebai_graph/preprocessing/properties/__init__.py

@@ -21,10 +21,22 @@
 )
 
 from .augmented_properties import (
-    AtomNodeLevel,
-    AtomFunctionalGroup,
-    AtomRingSize,
-    BondLevel,
+    AugAtomNodeLevel,
+    AugAtomFunctionalGroup,
+    AugAtomRingSize,
+    AugBondLevel,
+    AugAtomType,
+    AugNumAtomBonds,
+    AugAtomCharge,
+    AugAtomChirality,
+    AugAtomHybridization,
+    AugAtomNumHs,
+    AugAtomAromaticity,
+    AugBondAromaticity,
+    AugBondType,
+    AugBondInRing,
+    AugMoleculeNumRings,
+    AugRDKit2DNormalized,
 )
 
 # isort: on
@@ -46,8 +58,20 @@
     "MoleculeNumRings",
     "RDKit2DNormalized",
     # -------- Augmented Molecular Properties --------
-    "AtomNodeLevel",
-    "AtomFunctionalGroup",
-    "AtomRingSize",
-    "BondLevel",
+    "AugAtomNodeLevel",
+    "AugAtomFunctionalGroup",
+    "AugAtomRingSize",
+    "AugBondLevel",
+    "AugAtomType",
+    "AugNumAtomBonds",
+    "AugAtomCharge",
+    "AugAtomChirality",
+    "AugAtomHybridization",
+    "AugAtomNumHs",
+    "AugAtomAromaticity",
+    "AugBondAromaticity",
+    "AugBondType",
+    "AugBondInRing",
+    "AugMoleculeNumRings",
+    "AugRDKit2DNormalized",
 ]

chebai_graph/preprocessing/properties/augmented_properties.py

@@ -5,83 +5,12 @@
 
 from chebai_graph.preprocessing.property_encoder import OneHotEncoder, PropertyEncoder
 
+from . import properties as pr
 from .constants import *
-from .properties import AtomProperty, BondProperty
 
 
-class AugmentedBondProperty(BondProperty, ABC):
-    MAIN_KEY = "edges"
-
-    def get_property_value(self, augmented_mol: Dict) -> List:
-        if self.MAIN_KEY not in augmented_mol:
-            raise KeyError(
-                f"Key `{self.MAIN_KEY}` should be present in augmented molecule dict"
-            )
-
-        missing_keys = EDGE_LEVELS - augmented_mol[self.MAIN_KEY].keys()
-        if missing_keys:
-            raise KeyError(f"Missing keys {missing_keys} in augmented molecule nodes")
-
-        atom_molecule: Chem.Mol = augmented_mol[self.MAIN_KEY][WITHIN_ATOMS_EDGE]
-        if not isinstance(atom_molecule, Chem.Mol):
-            raise TypeError(
-                f'augmented_mol["{self.MAIN_KEY}"]["{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][ATOM_FG_EDGE]
-        fg_edges = augmented_mol[self.MAIN_KEY][WITHIN_FG_EDGE]
-        fg_graph_node_edges = augmented_mol[self.MAIN_KEY][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}"](["{ATOM_FG_EDGE}"]/["{WITHIN_FG_EDGE}"]/["{FG_GRAPHNODE_EDGE}"]) '
-                f"must be an instance of dict containing its properties"
-            )
-
-        # 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()]
-        )
-
-        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} "
-
-        return prop_list
-
-    @abstractmethod
-    def get_bond_value(self, bond: Chem.rdchem.Bond | Dict):
-        pass
-
-    def _check_modify_bond_prop_value(self, bond: Chem.rdchem.Bond | Dict, prop: str):
-        value = self._get_bond_prop_value(bond, prop)
-        if not value:
-            # Every atom/node should have given value
-            raise ValueError(f"'{prop}' is set but empty.")
-        return value
-
-    @staticmethod
-    def _get_bond_prop_value(bond: Chem.rdchem.Bond | Dict, prop: str):
-        if isinstance(bond, Chem.rdchem.Bond):
-            return bond.GetProp(prop)
-        elif isinstance(bond, dict):
-            return bond[prop]
-        else:
-            raise TypeError("Bond/Edge should be of type `Chem.rdchem.Bond` or `dict`.")
-
-
-class AugmentedAtomProperty(AtomProperty, ABC):
+# --------------------- Atom Properties -----------------------------
+class AugmentedAtomProperty(pr.AtomProperty, ABC):
     MAIN_KEY = "nodes"
 
     def get_property_value(self, augmented_mol: Dict):
@@ -145,31 +74,23 @@ def _get_atom_prop_value(self, atom: Chem.rdchem.Atom | Dict, prop: str):
             )
 
 
-class AtomNodeLevel(AugmentedAtomProperty):
+class AugAtomNodeLevel(AugmentedAtomProperty):
     def __init__(self, encoder: Optional[PropertyEncoder] = None):
         super().__init__(encoder or OneHotEncoder(self))
 
     def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
         return self._check_modify_atom_prop_value(atom, NODE_LEVEL)
 
 
-class AtomFunctionalGroup(AugmentedAtomProperty):
+class AugAtomFunctionalGroup(AugmentedAtomProperty):
     def __init__(self, encoder: Optional[PropertyEncoder] = None):
         super().__init__(encoder or OneHotEncoder(self))
 
     def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
         return self._check_modify_atom_prop_value(atom, "FG")
 
-    def _get_atom_prop_value(self, atom: Chem.rdchem.Atom | Dict, prop: str):
-        if isinstance(atom, Chem.rdchem.Atom):
-            return atom.GetProp(prop)
-        elif isinstance(atom, dict):
-            return atom[prop]
-        else:
-            raise TypeError("Atom/Node should be of type `Chem.rdchem.Atom` or `dict`.")
-
 
-class AtomRingSize(AugmentedAtomProperty):
+class AugAtomRingSize(AugmentedAtomProperty):
     def __init__(self, encoder: Optional[PropertyEncoder] = None):
         super().__init__(encoder or OneHotEncoder(self))
 
@@ -186,9 +107,150 @@ def _check_modify_atom_prop_value(self, atom: Chem.rdchem.Atom | Dict, prop: str
             return 0
 
 
-class BondLevel(AugmentedBondProperty):
+class AugNodeValueDefaulter(AugmentedAtomProperty, ABC):
+    def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
+        if isinstance(atom, Chem.rdchem.Atom):
+            # Delegate to superclass method for atom
+            return super().get_atom_value(atom)
+        elif isinstance(atom, dict):
+            return 0
+        else:
+            raise TypeError(
+                f"Expected Chem.rdchem.Atom or dict, got {type(atom).__name__}"
+            )
+
+
+class AugAtomType(AugNodeValueDefaulter, pr.AtomType): ...
+
+
+class AugNumAtomBonds(AugNodeValueDefaulter, pr.NumAtomBonds): ...
+
+
+class AugAtomCharge(AugNodeValueDefaulter, pr.AtomCharge): ...
+
+
+class AugAtomChirality(AugNodeValueDefaulter, pr.AtomChirality): ...
+
+
+class AugAtomHybridization(AugNodeValueDefaulter, pr.AtomHybridization): ...
+
+
+class AugAtomNumHs(AugNodeValueDefaulter, pr.AtomNumHs): ...
+
+
+class AugAtomAromaticity(AugNodeValueDefaulter, pr.AtomAromaticity): ...
+
+
+# --------------------- Bond Properties ------------------------------
+class AugmentedBondProperty(pr.BondProperty, ABC):
+    MAIN_KEY = "edges"
+
+    def get_property_value(self, augmented_mol: Dict) -> List:
+        if self.MAIN_KEY not in augmented_mol:
+            raise KeyError(
+                f"Key `{self.MAIN_KEY}` should be present in augmented molecule dict"
+            )
+
+        missing_keys = EDGE_LEVELS - augmented_mol[self.MAIN_KEY].keys()
+        if missing_keys:
+            raise KeyError(f"Missing keys {missing_keys} in augmented molecule nodes")
+
+        atom_molecule: Chem.Mol = augmented_mol[self.MAIN_KEY][WITHIN_ATOMS_EDGE]
+        if not isinstance(atom_molecule, Chem.Mol):
+            raise TypeError(
+                f'augmented_mol["{self.MAIN_KEY}"]["{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][ATOM_FG_EDGE]
+        fg_edges = augmented_mol[self.MAIN_KEY][WITHIN_FG_EDGE]
+        fg_graph_node_edges = augmented_mol[self.MAIN_KEY][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}"](["{ATOM_FG_EDGE}"]/["{WITHIN_FG_EDGE}"]/["{FG_GRAPHNODE_EDGE}"]) '
+                f"must be an instance of dict containing its properties"
+            )
+
+        # 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()]
+        )
+
+        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} "
+
+        return prop_list
+
+    @abstractmethod
+    def get_bond_value(self, bond: Chem.rdchem.Bond | Dict):
+        pass
+
+    def _check_modify_bond_prop_value(self, bond: Chem.rdchem.Bond | Dict, prop: str):
+        value = self._get_bond_prop_value(bond, prop)
+        if not value:
+            # Every atom/node should have given value
+            raise ValueError(f"'{prop}' is set but empty.")
+        return value
+
+    @staticmethod
+    def _get_bond_prop_value(bond: Chem.rdchem.Bond | Dict, prop: str):
+        if isinstance(bond, Chem.rdchem.Bond):
+            return bond.GetProp(prop)
+        elif isinstance(bond, dict):
+            return bond[prop]
+        else:
+            raise TypeError("Bond/Edge should be of type `Chem.rdchem.Bond` or `dict`.")
+
+
+class AugBondLevel(AugmentedBondProperty):
     def __init__(self, encoder: Optional[PropertyEncoder] = None):
         super().__init__(encoder or OneHotEncoder(self))
 
     def get_bond_value(self, bond: Chem.rdchem.Bond | Dict):
         return self._check_modify_bond_prop_value(bond, EDGE_LEVEL)
+
+
+class AugBondValueDefaulter(AugmentedBondProperty, ABC):
+    def get_bond_value(self, bond: Chem.rdchem.Bond | Dict):
+        if isinstance(bond, Chem.rdchem.Bond):
+            # Delegate to superclass method for bond
+            return super().get_bond_value(bond)
+        elif isinstance(bond, dict):
+            return 0
+        else:
+            raise TypeError("Bond/Edge should be of type `Chem.rdchem.Bond` or `dict`.")
+
+
+class AugBondAromaticity(AugBondValueDefaulter, pr.BondAromaticity): ...
+
+
+class AugBondType(AugBondValueDefaulter, pr.BondType): ...
+
+
+class AugBondInRing(AugBondValueDefaulter, pr.BondInRing): ...
+
+
+# --------------------- Molecular Properties ------------------------------
+class AugmentedMolecularProperty(pr.MolecularProperty, ABC):
+    def get_property_value(self, augmented_mol: Dict) -> list:
+        mol: Chem.Mol = augmented_mol[self.MAIN_KEY]["atom_nodes"]
+        assert isinstance(mol, Chem.Mol), "Molecule should be instance of `Chem.Mol`"
+        return super().get_property_value(mol)
+
+
+class AugMoleculeNumRings(AugmentedMolecularProperty, pr.MoleculeNumRings): ...
+
+
+class AugRDKit2DNormalized(AugmentedMolecularProperty, pr.RDKit2DNormalized): ...

chebai_graph/preprocessing/properties/properties.py

@@ -32,8 +32,8 @@ def on_finish(self):
     def __str__(self):
         return self.name
 
-    def get_property_value(self, mol: Chem.rdchem.Mol | Dict):
-        raise NotImplementedError
+    @abstractmethod
+    def get_property_value(self, mol: Chem.rdchem.Mol | Dict): ...
 
 
 class AtomProperty(MolecularProperty, ABC):

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -396,7 +396,6 @@ def _set_ring_fg_prop(self, connected_atoms, fg_nodes):
         ring_size = len(connected_atoms)
         fg_nodes[self._num_of_nodes] = {
             NODE_LEVEL: FG_NODE_LEVEL,
-            # E.g.,  Fused Ring has size "5-6", indicating size of each connected ring in fused ring
             "FG": f"RING_{ring_size}",
             "RING": ring_size,
         }
@@ -405,6 +404,8 @@ def _set_ring_fg_prop(self, connected_atoms, fg_nodes):
             ring_prop = atom.GetProp("RING")
             if not ring_prop:
                 raise ValueError("Atom does not have a ring size set")
+            # TODO: discuss the case, should it be max or average
+            # An atom belonging to multiple rings in fused Ring has size "5-6", indicating size of each ring
             max_ring_size = max(list(map(int, ring_prop.split("-"))))
             atom.SetProp("FG", f"RING_{max_ring_size}")
 

configs/data/chebi50_augmented_baseline.yml

@@ -0,0 +1,14 @@
+class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphFGAugmentorReader
+init_args:
+  properties:
+      - chebai_graph.preprocessing.properties.AugAtomType
+      - chebai_graph.preprocessing.properties.AugNumAtomBonds
+      - chebai_graph.preprocessing.properties.AugAtomCharge
+      - chebai_graph.preprocessing.properties.AugAtomAromaticity
+      - chebai_graph.preprocessing.properties.AugAtomHybridization
+      - chebai_graph.preprocessing.properties.AugAtomNumHs
+      - chebai_graph.preprocessing.properties.AugBondType
+      - chebai_graph.preprocessing.properties.AugBondInRing
+      - chebai_graph.preprocessing.properties.AugBondAromaticity
+      #- chebai_graph.preprocessing.properties.AugMoleculeNumRings
+      - chebai_graph.preprocessing.properties.AugRDKit2DNormalized