add aug rdkit norm prop

Author: aditya0by0

chebai_graph/preprocessing/datasets/chebi.py

@@ -1,6 +1,7 @@
 import os
 from abc import ABC
 from collections.abc import Callable
+from pprint import pformat
 
 import pandas as pd
 import torch
@@ -376,7 +377,7 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
         )
         rank_zero_info(
             f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n"
-            f"Properties Categories {props_categories}\n"
+            f"Properties Categories:\n{pformat(props_categories)}"
             f"n_atom_node_properties: {n_atom_node_properties}, "
             f"n_fg_node_properties: {n_fg_node_properties}, "
             f"n_bond_properties: {n_bond_properties}, "

chebai_graph/preprocessing/properties/__init__.py

@@ -43,6 +43,7 @@
     AugBondAromaticity,
     AugBondType,
     AugBondInRing,
+    AugRDKit2DNormalized,
 )
 
 # isort: on
@@ -84,4 +85,5 @@
     "AugBondAromaticity",
     "AugBondType",
     "AugBondInRing",
+    "AugRDKit2DNormalized",
 ]

chebai_graph/preprocessing/properties/augmented_properties.py

@@ -14,6 +14,7 @@
     AllNodeTypeProperty,
     AtomNodeTypeProperty,
     AugmentedBondProperty,
+    AugmentedMoleculeProperty,
     FGNodeTypeProperty,
     FrozenPropertyAlias,
 )
@@ -402,3 +403,9 @@ class AugBondInRing(AugBondValueDefaulter, pr.BondInRing):
     """
 
     ...
+
+
+# --------------------- Molecule Properties ------------------------------
+
+
+class AugRDKit2DNormalized(AugmentedMoleculeProperty, pr.RDKit2DNormalized): ...

chebai_graph/preprocessing/properties/base.py

@@ -439,3 +439,17 @@ def _get_bond_prop_value(bond: Chem.rdchem.Bond | dict, prop: str) -> str:
             return bond[prop]
         else:
             raise TypeError("Bond/Edge should be of type `Chem.rdchem.Bond` or `dict`.")
+
+
+class AugmentedMoleculeProperty(MoleculeProperty, ABC):
+    def get_property_value(self, augmented_mol: dict) -> list:
+        """
+        Get molecular property values from augmented molecule dict.
+        Args:
+            augmented_mol (dict): Augmented molecule dict.
+        Returns:
+            list: Property values of molecule.
+        """
+        mol: Chem.Mol = augmented_mol[AugmentedAtomProperty.MAIN_KEY]["atom_nodes"]
+        assert isinstance(mol, Chem.Mol), "Molecule should be instance of `Chem.Mol`"
+        return super().get_property_value(mol)

configs/data/chebi50_aug_prop_as_per_node.yml

@@ -16,7 +16,7 @@ init_args:
       - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG
       - chebai_graph.preprocessing.properties.IsFGAlkyl
       # Graph Node type properties
-      - chebai_graph.preprocessing.properties.RDKit2DNormalized
+      - chebai_graph.preprocessing.properties.AugRDKit2DNormalized
       # Bond properties
       - chebai_graph.preprocessing.properties.BondLevel
       - chebai_graph.preprocessing.properties.AugBondAromaticity