Add ROC-AUC score for each feature

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
+import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
 
+from sklearn import metrics
 from datasets import load_dataset
 
 import histos
@@ -16,9 +18,22 @@ def get_plot(features, n_bins):
     plotting_df = dataset_df.copy()
     if len(features) == 1:
         fig, ax = plt.subplots()
+        pos_samples = plotting_df[plotting_df["target"] == "spin-ON"][features[0]]
+        neg_samples = plotting_df[plotting_df["target"] == "spin-OFF"][features[0]]
+        y_score = np.concatenate([pos_samples, neg_samples], axis=0)
+        if pos_samples.mean() >= neg_samples.mean():
+            y_true = np.concatenate(
+                [np.ones_like(pos_samples), np.zeros_like(neg_samples)], axis=0
+            )
+            roc_auc_score = metrics.roc_auc_score(y_true, y_score)
+        else:
+            y_true = np.concatenate(
+                [np.zeros_like(pos_samples), np.ones_like(neg_samples)], axis=0
+            )
+            roc_auc_score = metrics.roc_auc_score(y_true, y_score)
         values = [
-            plotting_df[plotting_df["target"] == "spin-ON"][features[0]],
-            plotting_df[plotting_df["target"] == "spin-OFF"][features[0]],
+            pos_samples,
+            neg_samples,
         ]
         labels = ["spin-ON", "spin-OFF"]
         fig = histos.ratio_hist(
@@ -27,7 +42,7 @@ def get_plot(features, n_bins):
             reference_label=labels[1],
             n_bins=n_bins,
             hist_range=None,
-            title=features[0],
+            title=f"{features[0]} (ROC AUC: {roc_auc_score:.3f})",
         )
         return fig
     if len(features) == 2:

requirements.txt

@@ -1,2 +1,3 @@
 matplotlib==3.7.1
+scikit-learn==1.2.2
 seaborn==0.12.2