reformat w/ ruff

Author: sfluegel05

chebai_graph/models/dynamic_gni.py

@@ -85,9 +85,9 @@ def __init__(self, config: dict[str, Any], **kwargs: Any):
         print("Using complete randomness: ", self.complete_randomness)
 
         if not self.complete_randomness:
-            assert (
-                "pad_node_features" in config or "pad_edge_features" in config
-            ), "Missing 'pad_node_features' or 'pad_edge_features' in config when complete_randomness is False"
+            assert "pad_node_features" in config or "pad_edge_features" in config, (
+                "Missing 'pad_node_features' or 'pad_edge_features' in config when complete_randomness is False"
+            )
             self.pad_node_features = (
                 int(config["pad_node_features"])
                 if config.get("pad_node_features") is not None
@@ -112,9 +112,9 @@ def __init__(self, config: dict[str, Any], **kwargs: Any):
                     f"in each forward pass."
                 )
 
-            assert (
-                self.pad_node_features > 0 or self.pad_edge_features > 0
-            ), "'pad_node_features' or 'pad_edge_features' must be positive integers"
+            assert self.pad_node_features > 0 or self.pad_edge_features > 0, (
+                "'pad_node_features' or 'pad_edge_features' must be positive integers"
+            )
 
         self.resgated: BasicGNN = ResGatedModel(
             in_channels=self.in_channels,
@@ -182,9 +182,9 @@ def forward(self, batch: dict[str, Any]) -> Tensor:
                 )
                 new_edge_attr = torch.cat((graph_data.edge_attr, pad_edge), dim=1)
 
-        assert (
-            new_x is not None and new_edge_attr is not None
-        ), "Feature initialization failed"
+        assert new_x is not None and new_edge_attr is not None, (
+            "Feature initialization failed"
+        )
         out = self.resgated(
             x=new_x.float(),
             edge_index=graph_data.edge_index.long(),

chebai_graph/preprocessing/datasets/chebi.py

@@ -610,7 +610,14 @@ def _merge_props_into_base(
             enc_len = property_values.shape[1]
             # -------------- Node properties ---------------
             if isinstance(property, AllNodeTypeProperty):
-                x[:, atom_offset : atom_offset + enc_len] = property_values
+                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}"
+                    )
                 atom_offset += enc_len
                 fg_offset += enc_len
                 graph_offset += enc_len

chebai_graph/preprocessing/properties/base.py

@@ -262,9 +262,9 @@ def get_property_value(self, augmented_mol: dict) -> list:
                 )
             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"
+        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(
@@ -390,9 +390,9 @@ def get_property_value(self, augmented_mol: dict) -> list:
             )
 
         num_directed_edges = augmented_mol[self.MAIN_KEY][k.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} "
+        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
 

chebai_graph/preprocessing/property_encoder.py

@@ -261,9 +261,9 @@ def encode(self, token: float | int | None) -> torch.Tensor:
         """
         if token is None:
             return torch.zeros(1, self.get_encoding_length())
-        assert (
-            len(token) == self.get_encoding_length()
-        ), "Length of token should be equal to encoding length"
+        assert len(token) == self.get_encoding_length(), (
+            "Length of token should be equal to encoding length"
+        )
         # return torch.tensor([token]) # token is an ndarray, no need to create list of ndarray due to below warning
         # UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow.
         # Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor.

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -23,6 +23,7 @@
 # https://mail.python.org/pipermail/python-dev/2017-December/151283.html
 # Order preservation is necessary to to create `is_atom_node` mask
 
+
 class _AugmentorReader(DataReader, ABC):
     """
     Abstract base class for augmentor readers that extend ChemDataReader.
@@ -85,13 +86,17 @@ def _read_data(self, raw_data: str | Chem.Mol) -> tuple[GeomData, dict] | None:
         try:
             returned_result = self._create_augmented_graph(mol)
         except Exception as e:
-            print(f"Failed to construct augmented graph for smiles {smiles}, Error: {e}")
+            print(
+                f"Failed to construct augmented graph for smiles {smiles}, Error: {e}"
+            )
             self.f_cnt_for_aug_graph += 1
             return None
 
         # If the returned result is None, it indicates that the graph augmentation failed
         if returned_result is None:
-            print(f"Failed to construct augmented graph for smiles {smiles} (returned None)")
+            print(
+                f"Failed to construct augmented graph for smiles {smiles} (returned None)"
+            )
             self.f_cnt_for_aug_graph += 1
             return None
 
@@ -100,36 +105,39 @@ def _read_data(self, raw_data: str | Chem.Mol) -> tuple[GeomData, dict] | None:
 
         # Empty features initialized; node and edge features can be added later
         NUM_NODES = augmented_molecule["nodes"]["num_nodes"]
-        assert (
-            NUM_NODES is not None and NUM_NODES > 1
-        ), "Num of nodes in augmented graph should be more than 1"
+        assert NUM_NODES is not None and NUM_NODES > 1, (
+            "Num of nodes in augmented graph should be more than 1"
+        )
 
         x = torch.zeros((NUM_NODES, 0))
         edge_attr = torch.zeros((augmented_molecule["edges"][k.NUM_EDGES], 0))
 
-        assert (
-            edge_index.shape[0] == 2
-        ), f"Expected edge_index to have shape [2, num_edges], but got shape {edge_index.shape}"
+        assert edge_index.shape[0] == 2, (
+            f"Expected edge_index to have shape [2, num_edges], but got shape {edge_index.shape}"
+        )
 
-        assert (
-            edge_index.shape[1] == edge_attr.shape[0]
-        ), f"Mismatch between number of edges in edge_index ({edge_index.shape[1]}) and edge_attr ({edge_attr.shape[0]})"
+        assert edge_index.shape[1] == edge_attr.shape[0], (
+            f"Mismatch between number of edges in edge_index ({edge_index.shape[1]}) and edge_attr ({edge_attr.shape[0]})"
+        )
 
-        assert (
-            len(set(edge_index[0].tolist())) == x.shape[0]
-        ), f"Number of unique source nodes in edge_index ({len(set(edge_index[0].tolist()))}) does not match number of nodes in x ({x.shape[0]})"
+        assert len(set(edge_index[0].tolist())) == x.shape[0], (
+            f"Number of unique source nodes in edge_index ({len(set(edge_index[0].tolist()))}) does not match number of nodes in x ({x.shape[0]})"
+        )
 
         # Create a boolean mask: True for atom, False for augmented
         is_atom_mask = torch.zeros(NUM_NODES, dtype=torch.bool)
         NUM_ATOM_NODES = augmented_molecule["nodes"]["atom_nodes"].GetNumAtoms()
         is_atom_mask[:NUM_ATOM_NODES] = True
 
-        return (GeomData(
-            x=x,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            is_atom_node=is_atom_mask,
-        ), augmented_molecule)
+        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:
         """
@@ -219,19 +227,19 @@ def _augment_graph_structure(self, mol: Chem.Mol) -> dict:
         atom_edge_index = self._generate_atom_level_edge_index(mol)
 
         total_atoms = mol.GetNumAtoms()
-        assert (
-            self._idx_of_node == total_atoms
-        ), f"Mismatch in number of nodes: expected {total_atoms}, got {self._idx_of_node}"
+        assert self._idx_of_node == total_atoms, (
+            f"Mismatch in number of nodes: expected {total_atoms}, got {self._idx_of_node}"
+        )
 
         node_info = {
             "atom_nodes": mol,
             "num_nodes": self._idx_of_node,
         }
 
         total_edges = mol.GetNumBonds()
-        assert (
-            self._idx_of_edge == total_edges
-        ), f"Mismatch in number of edges: expected {total_edges}, got {self._idx_of_edge}"
+        assert self._idx_of_edge == total_edges, (
+            f"Mismatch in number of edges: expected {total_edges}, got {self._idx_of_edge}"
+        )
         edge_info = {
             k.WITHIN_ATOMS_EDGE: mol,
             k.NUM_EDGES: self._idx_of_edge,
@@ -285,7 +293,9 @@ def on_finish(self) -> None:
         )
         self.mol_object_buffer = {}
 
-    def read_property(self, raw_data: str | Chem.Mol | dict, 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.
 
@@ -362,16 +372,16 @@ def _augment_graph_structure(
         augmented_mol["directed_edge_index"] = directed_edge_index
 
         total_atoms = sum([mol.GetNumAtoms(), len(fg_nodes)])
-        assert (
-            self._idx_of_node == total_atoms
-        ), f"Mismatch in number of nodes: expected {total_atoms}, got {self._idx_of_node}"
+        assert self._idx_of_node == total_atoms, (
+            f"Mismatch in number of nodes: expected {total_atoms}, got {self._idx_of_node}"
+        )
         augmented_mol["node_info"]["fg_nodes"] = fg_nodes
         augmented_mol["node_info"]["num_nodes"] = self._idx_of_node
 
         total_edges = sum([mol.GetNumBonds(), len(atom_fg_edges)])
-        assert (
-            self._idx_of_edge == total_edges
-        ), f"Mismatch in number of edges: expected {total_edges}, got {self._idx_of_edge}"
+        assert self._idx_of_edge == total_edges, (
+            f"Mismatch in number of edges: expected {total_edges}, got {self._idx_of_edge}"
+        )
         augmented_mol["edge_info"][k.ATOM_FG_EDGE] = atom_fg_edges
         augmented_mol["edge_info"][k.NUM_EDGES] = self._idx_of_edge
 
@@ -594,12 +604,12 @@ def _augment_graph_structure(
         augmented_struct["edge_info"][k.WITHIN_FG_EDGE] = internal_fg_edges
         augmented_struct["edge_info"][k.NUM_EDGES] += len(internal_fg_edges)
 
-        assert (
-            self._idx_of_edge == augmented_struct["edge_info"][k.NUM_EDGES]
-        ), f"Mismatch in number of edges: expected {self._idx_of_edge}, got {augmented_struct['edge_info'][k.NUM_EDGES]}"
-        assert (
-            self._idx_of_node == augmented_struct["node_info"]["num_nodes"]
-        ), f"Mismatch in number of nodes: expected {self._idx_of_node}, got {augmented_struct['node_info']['num_nodes']}"
+        assert self._idx_of_edge == augmented_struct["edge_info"][k.NUM_EDGES], (
+            f"Mismatch in number of edges: expected {self._idx_of_edge}, got {augmented_struct['edge_info'][k.NUM_EDGES]}"
+        )
+        assert self._idx_of_node == augmented_struct["node_info"]["num_nodes"], (
+            f"Mismatch in number of nodes: expected {self._idx_of_node}, got {augmented_struct['node_info']['num_nodes']}"
+        )
 
         augmented_struct["directed_edge_index"] = torch.cat(
             [
@@ -629,9 +639,9 @@ def _construct_fg_level_structure(
         internal_edge_index = [[], []]
 
         def add_fg_internal_edge(source_fg: int, target_fg: int) -> None:
-            assert (
-                source_fg is not None and target_fg is not None
-            ), "Each bond should have a fg node on both end"
+            assert source_fg is not None and target_fg is not None, (
+                "Each bond should have a fg node on both end"
+            )
             assert source_fg != target_fg, "Source and Target FG should be different"
 
             edge_key = tuple(sorted((source_fg, target_fg)))
@@ -718,15 +728,15 @@ def _add_graph_node_and_edges_to_nodes(
 
         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]
-        ), f"Mismatch in number of edges: expected {self._idx_of_edge}, got {augmented_struct['edge_info'][k.NUM_EDGES]}"
+        assert self._idx_of_edge == augmented_struct["edge_info"][k.NUM_EDGES], (
+            f"Mismatch in number of edges: expected {self._idx_of_edge}, got {augmented_struct['edge_info'][k.NUM_EDGES]}"
+        )
 
         augmented_struct["node_info"]["graph_node"] = graph_node
         augmented_struct["node_info"]["num_nodes"] += 1
-        assert (
-            self._idx_of_node == augmented_struct["node_info"]["num_nodes"]
-        ), f"Mismatch in number of nodes: expected {self._idx_of_node}, got {augmented_struct['node_info']['num_nodes']}"
+        assert self._idx_of_node == augmented_struct["node_info"]["num_nodes"], (
+            f"Mismatch in number of nodes: expected {self._idx_of_node}, got {augmented_struct['node_info']['num_nodes']}"
+        )
 
         augmented_struct["directed_edge_index"] = torch.cat(
             [
@@ -957,4 +967,3 @@ def _augment_graph_structure(
         return self._add_graph_node_and_edges_to_nodes(
             augmented_struct, atom_ids | fg_to_atoms_map.keys()
         )
-    

chebai_graph/preprocessing/reader/reader.py

@@ -104,7 +104,9 @@ def on_finish(self) -> None:
         print(f"Failed to read {self.failed_counter} SMILES in total")
         self.mol_object_buffer = {}
 
-    def read_property(self, raw_data: str | Chem.Mol, 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.
 
@@ -160,7 +162,9 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         """
         # raw_data is a SMILES string
         try:
-            mol = self._smiles_to_mol(raw_data) if isinstance(raw_data, str) else raw_data
+            mol = (
+                self._smiles_to_mol(raw_data) if isinstance(raw_data, str) else raw_data
+            )
         except ValueError:
             return None
         assert isinstance(mol, nx.Graph)
@@ -193,7 +197,7 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         nx.set_edge_attributes(mol, de, "edge_attr")
         data = from_networkx(mol)
         return data
-    
+
     def _smiles_to_mol(self, smiles: str) -> Chem.rdchem.Mol | None:
         """
         Load SMILES string into an RDKit molecule object.