fix: save ensemble results correctly, skip undefined f1 scores

Author: sfluegel05

chebifier/ensemble/base_ensemble.py

@@ -39,15 +39,15 @@ def gather_predictions(self, smiles_list):
                         predicted_classes.add(cls)
         print(f"Sorting predictions...")
         predicted_classes = sorted(list(predicted_classes))
-        predicted_classes = {cls: i for i, cls in enumerate(predicted_classes)}
+        predicted_classes_dict = {cls: i for i, cls in enumerate(predicted_classes)}
         ordered_logits = torch.zeros(len(smiles_list), len(predicted_classes), len(self.models)) * torch.nan
         for i, model_prediction in enumerate(model_predictions):
             for j, logits_for_smiles in tqdm.tqdm(enumerate(model_prediction),
                                                  total=len(model_prediction),
                                                  desc=f"Sorting predictions for {self.models[i].model_name}"):
                 if logits_for_smiles is not None:
                     for cls in logits_for_smiles:
-                        ordered_logits[j, predicted_classes[cls], i] = logits_for_smiles[cls]
+                        ordered_logits[j, predicted_classes_dict[cls], i] = logits_for_smiles[cls]
 
         return ordered_logits, predicted_classes
 
@@ -114,7 +114,7 @@ def predict_smiles_list(self, smiles_list, load_preds_if_possible=True) -> list:
         preds_file = f"predictions_by_model_{'_'.join(model.model_name for model in self.models)}.pt"
         predicted_classes_file = f"predicted_classes_{'_'.join(model.model_name for model in self.models)}.txt"
         if not load_preds_if_possible or not os.path.isfile(preds_file):
-            ordered_predictions = predicted_classes = self.gather_predictions(smiles_list)
+            ordered_predictions, predicted_classes = self.gather_predictions(smiles_list)
             # save predictions
             torch.save(ordered_predictions, preds_file)
             with open(predicted_classes_file, "w") as f:

chebifier/ensemble/weighted_majority_ensemble.py

@@ -44,8 +44,9 @@ def calculate_classwise_weights(self, predicted_classes):
             if model.classwise_weights is None:
                 continue
             for cls, weights in model.classwise_weights.items():
-                f1 = 2 * weights["TP"] / (2 * weights["TP"] + weights["FP"] + weights["FN"])
-                weights_by_cls[predicted_classes[cls], j] *= f1
+                if (2 * weights["TP"] + weights["FP"] + weights["FN"]) > 0:
+                    f1 = 2 * weights["TP"] / (2 * weights["TP"] + weights["FP"] + weights["FN"])
+                    weights_by_cls[predicted_classes[cls], j] *= f1
 
         print(f"Calculated model weightings. The average weights are:")
         for i, model in enumerate(self.models):