Implement safe ROC at inference

alan-turing-institute · Oct 19, 2023 · e7c9981 · e7c9981
1 parent 76c64c7
commit e7c9981
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/grace/evaluation/inference.py b/grace/evaluation/inference.py
@@ -12,13 +12,13 @@
 from sklearn.metrics import (
     accuracy_score,
     precision_recall_fscore_support,
-    roc_auc_score,
     average_precision_score,
 )
 
 from grace.styling import LOGGER
 from grace.base import GraphAttrs, Annotation, Prediction
 from grace.models.datasets import dataset_from_graph
+from grace.evaluation.metrics_classifier import safe_roc_auc_score
 from grace.visualisation.plotting import (
     plot_confusion_matrix_tiles,
     plot_areas_under_curves,
@@ -359,11 +359,11 @@ def calculate_numerical_results_on_entire_batch(
         inference_batch_metrics["Batch F1-score (edges)"] = prf1_edges[2]
 
         # AUC scores:
-        inference_batch_metrics["Batch AUROC (nodes)"] = roc_auc_score(
+        inference_batch_metrics["Batch AUROC (nodes)"] = safe_roc_auc_score(
             y_true=predictions_data["n_true"],
             y_score=predictions_data["n_prob"],
         )
-        inference_batch_metrics["Batch AUROC (edges)"] = roc_auc_score(
+        inference_batch_metrics["Batch AUROC (edges)"] = safe_roc_auc_score(
             y_true=predictions_data["e_true"],
             y_score=predictions_data["e_prob"],
         )

diff --git a/grace/visualisation/plotting.py b/grace/visualisation/plotting.py
@@ -1,5 +1,6 @@
-from grace.base import GraphAttrs, Annotation
 from grace.styling import COLORMAPS
+from grace.base import GraphAttrs, Annotation
+from grace.evaluation.metrics_classifier import safe_roc_auc_score
 
 import matplotlib.pyplot as plt
 import networkx as nx
@@ -10,11 +11,10 @@
 
 from skimage.util import montage
 from sklearn.metrics import (
-    ConfusionMatrixDisplay,
-    roc_auc_score,
-    RocCurveDisplay,
     average_precision_score,
+    ConfusionMatrixDisplay,
     PrecisionRecallDisplay,
+    RocCurveDisplay,
 )
 
 
@@ -179,7 +179,7 @@ def plot_areas_under_curves(
     fig, axes = plt.subplots(nrows=1, ncols=2, figsize=figsize)
 
     # Area under ROC:
-    roc_score_nodes = roc_auc_score(y_true=node_true, y_score=node_pred)
+    roc_score_nodes = safe_roc_auc_score(y_true=node_true, y_score=node_pred)
     RocCurveDisplay.from_predictions(
         y_true=node_true,
         y_pred=node_pred,
@@ -189,7 +189,7 @@ def plot_areas_under_curves(
         ax=axes[0],
     )
 
-    roc_score_edges = roc_auc_score(y_true=edge_true, y_score=edge_pred)
+    roc_score_edges = safe_roc_auc_score(y_true=edge_true, y_score=edge_pred)
     RocCurveDisplay.from_predictions(
         y_true=edge_true,
         y_pred=edge_pred,