property failure handling

Author: sfluegel05

chebai_graph/preprocessing/datasets/chebi.py

@@ -42,7 +42,7 @@
     RandomFeatureInitializationReader,
 )
 
-from .utils import resolve_property
+from chebai_graph.preprocessing.datasets.utils import resolve_property
 
 
 class ChEBI50GraphData(ChEBIOver50):
@@ -142,9 +142,8 @@ def enc_if_not_none(encode, value):
                         r = None
                     returned_results.append(r)
                 mols = [
-                    augmented_mol[1]
+                    augmented_mol[1] if augmented_mol is not None else None
                     for augmented_mol in returned_results
-                    if augmented_mol is not None
                 ]
             else:
                 mols = features
@@ -155,7 +154,7 @@ def enc_if_not_none(encode, value):
                     # read all property values first, then encode
                     rank_zero_info(f"\tReading property values of {property.name}...")
                     property_values = [
-                        self.reader.read_property(mol, property)
+                        self.reader.read_property(mol, property) if mol is not None else None
                         for mol in tqdm.tqdm(mols)
                     ]
                     rank_zero_info(f"\tEncoding property values of {property.name}...")
@@ -610,14 +609,7 @@ def _merge_props_into_base(
             enc_len = property_values.shape[1]
             # -------------- Node properties ---------------
             if isinstance(property, AllNodeTypeProperty):
-                try:
-                    x[:, atom_offset : atom_offset + enc_len] = property_values
-                except Exception as e:
-                    raise ValueError(
-                        f"Error assigning property '{property.name}' values to node features: {e}\n"
-                        f"Property values shape: {property_values.shape}, expected (num_nodes, {enc_len})\n"
-                        f"Node feature matrix shape: {x.shape}"
-                    )
+                x[:, atom_offset : atom_offset + enc_len] = property_values
                 atom_offset += enc_len
                 fg_offset += enc_len
                 graph_offset += enc_len
@@ -855,3 +847,9 @@ class ChEBI25_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOverX):
     READER = AtomFGReader_WithFGEdges_WithGraphNode
 
     THRESHOLD = 25
+
+
+if __name__ == "__main__":
+    dataset = ChEBI25_WFGE_WGN_AsPerNodeType(chebi_version=248, subset="3_STAR")
+    dataset.prepare_data()
+    dataset.setup()

chebai_graph/preprocessing/properties/properties.py

@@ -292,8 +292,13 @@ def get_property_value(self, mol: Chem.rdchem.Mol) -> list[np.ndarray]:
         Returns:
             list[np.ndarray]: List containing the descriptor numpy array (excluding first element).
         """
+        if mol is None:
+            raise ValueError("Input molecule is None.")
+        smiles = Chem.MolToSmiles(mol)
+        if smiles is None:
+            raise ValueError("Could not convert molecule to SMILES.")
         features_normalized = self.generator_normalized.processMol(
-            mol, Chem.MolToSmiles(mol)
+            mol, smiles
         )
         features_normalized = np.nan_to_num(features_normalized)
         return [features_normalized[1:]]