Merge pull request #8 from ChEB-AI/features-sfluegel

Minor fixes

Author: sfluegel05

.pre-commit-config.yaml

@@ -4,6 +4,6 @@ repos:
 #    hooks:
 #    -   id: isort
 -   repo: https://github.com/psf/black
-    rev: "22.10.0"
+    rev: "24.2.0"
     hooks:
     -   id: black

chebai/loss/__init__.py

@@ -87,18 +87,18 @@ def create_feature_vectors(self):
         for idx in range(self.no_of_atoms):
             sorted_path = self.directed_graphs[idx, :, :]
 
-            self.local_input_vector[
-                idx, idx, :length_of_atom_features
-            ] = self.get_atom_features(idx)
+            self.local_input_vector[idx, idx, :length_of_atom_features] = (
+                self.get_atom_features(idx)
+            )
 
             no_of_incoming_edges = {}
             for i in range(self.no_of_atoms - 1):
                 node1 = sorted_path[i, 0]
                 node2 = sorted_path[i, 1]
 
-                self.local_input_vector[
-                    idx, node1, :length_of_atom_features
-                ] = self.get_atom_features(node1)
+                self.local_input_vector[idx, node1, :length_of_atom_features] = (
+                    self.get_atom_features(node1)
+                )
 
                 if node2 in no_of_incoming_edges:
                     index = no_of_incoming_edges[node2]

chebai/molecule.py

@@ -9,11 +9,11 @@
     "JCI_500_COLUMNS_INT",
 ]
 
-import queue
 from abc import ABC
 from collections import OrderedDict
 import os
 import pickle
+import queue
 
 from iterstrat.ml_stratifiers import (
     MultilabelStratifiedKFold,
@@ -219,7 +219,7 @@ def setup_processed(self):
             self._setup_pruned_test_set()
 
     def get_test_split(self, df: pd.DataFrame):
-        print("Split dataset into train (including val) / test")
+        print("Get test data split")
 
         df_list = df.values.tolist()
         df_list = [row[3:] for row in df_list]
@@ -247,8 +247,8 @@ def get_train_val_splits_given_test(self, df: pd.DataFrame, test_df: pd.DataFram
         print(f"Split dataset into train / val with given test set")
 
         df_trainval = df
-        test_smiles = test_df["SMILES"].tolist()
-        mask = [smiles not in test_smiles for smiles in df_trainval["SMILES"]]
+        test_ids = test_df["id"].tolist()
+        mask = [trainval_id not in test_ids for trainval_id in df_trainval["id"]]
         df_trainval = df_trainval[mask]
         df_trainval_list = df_trainval.values.tolist()
         df_trainval_list = [row[3:] for row in df_trainval_list]
@@ -265,9 +265,9 @@ def get_train_val_splits_given_test(self, df: pd.DataFrame, test_df: pd.DataFram
                 df_validation = df_trainval.iloc[val_ids]
                 df_train = df_trainval.iloc[train_ids]
                 folds[self.raw_file_names_dict[f"fold_{fold}_train"]] = df_train
-                folds[
-                    self.raw_file_names_dict[f"fold_{fold}_validation"]
-                ] = df_validation
+                folds[self.raw_file_names_dict[f"fold_{fold}_validation"]] = (
+                    df_validation
+                )
 
             return folds
 
@@ -513,7 +513,7 @@ def extract_class_hierarchy(self, chebi_path):
         g.add_edges_from([(p, q["id"]) for q in elements for p in q["parents"]])
 
         g = nx.transitive_closure_dag(g)
-        g = g.subgraph(nx.descendants(g, self.top_class_id))
+        g = g.subgraph(list(nx.descendants(g, self.top_class_id)) + [self.top_class_id])
         print("Compute transitive closure")
         return g
 

chebai/preprocessing/datasets/chebi.py

@@ -60,7 +60,7 @@ def evaluate_model(
         collated.x = collated.to_x(model.device)
         collated.y = collated.to_y(model.device)
         processable_data = model._process_batch(collated, 0)
-        model_output = model(processable_data)
+        model_output = model(processable_data, **processable_data["model_kwargs"])
         preds, labels = model._get_prediction_and_labels(
             processable_data, processable_data["labels"], model_output
         )
@@ -166,6 +166,8 @@ def print_metrics(preds, labels, device, classes=None, top_k=10, markdown_output
 
     zeros = []
     for i, f1 in enumerate(classwise_f1):
-        if f1 == 0.0 and torch.sum(labels[:, i]):
+        if f1 == 0.0 and torch.sum(labels[:, i]) != 0:
             zeros.append(f"{classes[i] if classes is not None else i}")
-    print(f'Classes with F1-score == 0 (and non-zero labels): {", ".join(zeros)}')
+    print(
+        f'Found {len(zeros)} classes with F1-score == 0 (and non-zero labels): {", ".join(zeros)}'
+    )

chebai/result/classification.py

@@ -42,6 +42,9 @@
         "lightning",
         "jsonargparse[signatures]>=4.17.0",
         "omegaconf",
+        "seaborn",
+        "deepsmiles",
+        "iterative-stratification",
     ],
     extras_require={"dev": ["black", "isort", "pre-commit"]},
 )