make read_properties more flexible, don't recalculate properties

Author: sfluegel05

chebai_graph/preprocessing/datasets/chebi.py

@@ -15,6 +15,7 @@
 )
 from lightning_utilities.core.rank_zero import rank_zero_info
 from torch_geometric.data.data import Data as GeomData
+from rdkit import Chem
 
 from chebai_graph.preprocessing.properties import (
     AllNodeTypeProperty,
@@ -185,20 +186,23 @@ def _after_setup(self, **kwargs) -> None:
         super()._after_setup(**kwargs)
 
     def _preprocess_smiles_for_pred(
-        self, idx, smiles: str, model_hparams: Optional[dict] = None
+        self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None
     ) -> dict:
         """Preprocess prediction data."""
         # Add dummy labels because the collate function requires them.
         # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`,
         # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty.
         result = self.reader.to_data(
-            {"id": f"smiles_{idx}", "features": smiles, "labels": [1, 2]}
+            {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]}
         )
+        # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object
+        if isinstance(result["features"], tuple):
+            result["features"], raw_data = result["features"][0]
         if result is None or result["features"] is None:
             return None
         for property in self.properties:
             property.encoder.eval = True
-            property_value = self.reader.read_property(smiles, property)
+            property_value = self.reader.read_property(raw_data, property)
             if property_value is None or len(property_value) == 0:
                 encoded_value = None
             else:

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -59,16 +59,16 @@ def name(cls) -> str:
         """
         return f"{cls.__name__}".lower()
 
-    def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
+    def _read_data(self, raw_data: str | Chem.Mol) -> tuple[GeomData, dict] | None:
         """
         Reads and augments molecular data from a SMILES string.
 
         Args:
             raw_data (str | Chem.Mol): SMILES string or RDKit molecule object representing the molecule.
 
         Returns:
-            GeomData | None: A PyTorch Geometric Data object with augmented nodes and edges,
-            or None if parsing or augmentation fails.
+            tuple[GeomData, dict] | None: A tuple containing a PyTorch Geometric Data object with augmented nodes and edges,
+            and a dictionary of augmented molecule data, or None if parsing or augmentation fails.
 
         Raises:
             RuntimeError: If an unexpected error occurs during graph augmentation.
@@ -124,12 +124,12 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         NUM_ATOM_NODES = augmented_molecule["nodes"]["atom_nodes"].GetNumAtoms()
         is_atom_mask[:NUM_ATOM_NODES] = True
 
-        return GeomData(
+        return (GeomData(
             x=x,
             edge_index=edge_index,
             edge_attr=edge_attr,
             is_atom_node=is_atom_mask,
-        )
+        ), augmented_molecule)
 
     def _smiles_to_mol(self, smiles: str) -> Chem.Mol | None:
         """
@@ -285,32 +285,35 @@ def on_finish(self) -> None:
         )
         self.mol_object_buffer = {}
 
-    def read_property(self, data: str | Chem.Mol, property: MolecularProperty) -> list | None:
+    def read_property(self, raw_data: str | Chem.Mol | dict, property: MolecularProperty) -> list | None:
         """
         Reads a specific property from a molecule represented by a SMILES string.
 
         Args:
-            data (str | Chem.Mol): SMILES string or RDKit molecule object representing the molecule.
+            raw_data (str | Chem.Mol | dict): SMILES string, RDKit molecule object, or dictionary representation of a molecule.
             property (MolecularProperty): Molecular property object for which the value needs to be extracted.
 
         Returns:
             list | None: Property values if molecule parsing is successful, else None.
         """
-        if isinstance(data, Chem.Mol):
-            mol = data
+        if isinstance(raw_data, dict):
+            augmented_mol = raw_data
         else:
-            smiles = data
-            if smiles in self.mol_object_buffer:
-                return property.get_property_value(self.mol_object_buffer[smiles])
-            mol = self._smiles_to_mol(smiles)
-        if mol is None:
-            return None
+            if isinstance(raw_data, Chem.Mol):
+                mol = raw_data
+            else:
+                smiles = raw_data
+                if smiles in self.mol_object_buffer:
+                    return property.get_property_value(self.mol_object_buffer[smiles])
+                mol = self._smiles_to_mol(smiles)
+            if mol is None:
+                return None
 
-        returned_result = self._create_augmented_graph(mol)
-        if returned_result is None:
-            return None
+            returned_result = self._create_augmented_graph(mol)
+            if returned_result is None:
+                return None
 
-        _, augmented_mol = returned_result
+            _, augmented_mol = returned_result
         return property.get_property_value(augmented_mol)
 
 
@@ -680,14 +683,14 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         Returns:
             Data | None: Geometric data object with is_graph_node annotation.
         """
-        geom_data = super()._read_data(raw_data)
+        geom_data, augmented_mol = super()._read_data(raw_data)
         if geom_data is None:
             return None
         NUM_NODES = geom_data.x.shape[0]
         is_graph_node = torch.zeros(NUM_NODES, dtype=torch.bool)
         is_graph_node[-1] = True
         geom_data.is_graph_node = is_graph_node
-        return geom_data
+        return (geom_data, augmented_mol)
 
     def _add_graph_node_and_edges_to_nodes(
         self,

chebai_graph/preprocessing/reader/reader.py

@@ -68,7 +68,7 @@ def _smiles_to_mol(self, smiles: str) -> Chem.rdchem.Mol | None:
         self.mol_object_buffer[smiles] = mol
         return mol
 
-    def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
+    def _read_data(self, raw_data: str | Chem.Mol) -> tuple[GeomData, Chem.Mol] | None:
         """
         Convert raw SMILES string data into a PyTorch Geometric Data object.
 
@@ -95,7 +95,7 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         # edge_index.shape == [2, num_edges]; edge_attr.shape == [num_edges, num_edge_features]
         edge_attr = torch.zeros((edge_index.size(1), 0))
 
-        return GeomData(x=x, edge_index=edge_index, edge_attr=edge_attr)
+        return (GeomData(x=x, edge_index=edge_index, edge_attr=edge_attr), mol)
 
     def on_finish(self) -> None:
         """
@@ -104,18 +104,18 @@ def on_finish(self) -> None:
         print(f"Failed to read {self.failed_counter} SMILES in total")
         self.mol_object_buffer = {}
 
-    def read_property(self, smiles: str, property: MolecularProperty) -> list | None:
+    def read_property(self, raw_data: str | Chem.Mol, property: MolecularProperty) -> list | None:
         """
         Read a molecular property for a given SMILES string.
 
         Args:
-            smiles (str): SMILES string of the molecule.
+            raw_data (str | Chem.Mol): SMILES string or RDKit molecule object of the molecule.
             property (MolecularProperty): Property extractor to apply.
 
         Returns:
             list | None: Property values or None if molecule parsing failed.
         """
-        mol = self._smiles_to_mol(smiles)
+        mol = self._smiles_to_mol(raw_data) if isinstance(raw_data, str) else raw_data
         if mol is None:
             return None
         return property.get_property_value(mol)