From 60aa36702f0a47da3e270257cfd40bde10a87ee7 Mon Sep 17 00:00:00 2001
From: Mori Ludovico <lumori@uaf-2.t2.ucsd.edu>
Date: Fri, 25 Apr 2025 14:22:20 -0700
Subject: [PATCH] improve roc plotting, add exception printout when loading
 data

---
 src/boostedhh/plotting.py | 37 +++++++++++++++++++++++++------------
 src/boostedhh/utils.py    |  3 ++-
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/boostedhh/plotting.py b/src/boostedhh/plotting.py
index 97227ab..00f7ebe 100644
--- a/src/boostedhh/plotting.py
+++ b/src/boostedhh/plotting.py
@@ -846,14 +846,14 @@ def multiROCCurveGrey(
 
 
 th_colours = [
-    # "#36213E",
+    "#36213E",
     # "#9381FF",
     "#1f78b4",
-    # "#a6cee3",
+    "#a6cee3",
     # "#32965D",
     "#7CB518",
     "#EDB458",
-    # "#ff7f00",
+    "#ff7f00",
     "#a70000",
 ]
 
@@ -870,6 +870,7 @@ def multiROCCurve(
     plot_dir="",
     name="",
     prelim=True,
+    lumi=None,
     show=False,
 ):
     if ylim is None:
@@ -877,7 +878,11 @@ def multiROCCurve(
     if xlim is None:
         xlim = [0, 1]
     if thresholds is None:
-        thresholds = [[0.9, 0.98, 0.995, 0.9965, 0.998], [0.99, 0.997, 0.998, 0.999, 0.9997]]
+        thresholds = [
+            0.9,
+            0.99,
+            0.998,
+        ]  # [0.9, 0.98, 0.995, 0.9965, 0.998] # [[0.99, 0.997, 0.998, 0.999, 0.9997]]
 
     plt.rcParams.update({"font.size": 32})
 
@@ -897,13 +902,12 @@ def multiROCCurve(
                 color=COLOURS[ROC_COLOURS[i * len(roc_sigs) + j]],
                 linestyle=LINESTYLES[i * len(roc_sigs) + j],
             )
-
             pths = {th: [[], []] for th in pthresholds}
             for th in pthresholds:
                 idx = _find_nearest(roc["thresholds"], th)
                 pths[th][0].append(roc["tpr"][idx])
                 pths[th][1].append(roc["fpr"][idx])
-                print(roc["tpr"][idx])
+                # print(roc["tpr"][idx])
 
             for k, th in enumerate(pthresholds):
                 ax.scatter(
@@ -935,7 +939,15 @@ def multiROCCurve(
                     alpha=0.5,
                 )
 
-    add_cms_label(ax, year, data=False, label="Preliminary" if prelim else None, loc=1, lumi=False)
+                hep.cms.label(
+                    ax=ax,
+                    label="Preliminary" if prelim else "",
+                    data=True,
+                    year=year,
+                    com="13.6",
+                    fontsize=20,
+                    lumi=lumi,
+                )
 
     if log:
         plt.yscale("log")
@@ -944,16 +956,16 @@ def multiROCCurve(
     ax.set_ylabel("Background efficiency")
     ax.set_xlim(*xlim)
     ax.set_ylim(*ylim)
-    ax.legend(loc="lower right", fontsize=24)
+    ax.legend(loc="lower right", fontsize=21)
 
     if title:
         ax.text(
-            0.05,
-            0.83,
+            0.02,
+            0.93,
             title,
             transform=ax.transAxes,
-            fontsize=24,
-            fontproperties="Tex Gyre Heros:bold",
+            fontsize=20,
+            # fontproperties="Tex Gyre Heros:bold",
         )
 
     if kin_label:
@@ -968,6 +980,7 @@ def multiROCCurve(
 
     if len(name):
         plt.savefig(f"{plot_dir}/{name}.pdf", bbox_inches="tight")
+        plt.savefig(f"{plot_dir}/{name}.png", bbox_inches="tight")
 
     if show:
         plt.show()
diff --git a/src/boostedhh/utils.py b/src/boostedhh/utils.py
index f3ea959..f7f31f1 100644
--- a/src/boostedhh/utils.py
+++ b/src/boostedhh/utils.py
@@ -450,11 +450,12 @@ def load_sample(
         logger.debug(f"Loading {load_sample}")
         try:
             events = pd.read_parquet(parquet_path, filters=filters, columns=load_columns)
-        except Exception:
+        except Exception as e:
             warnings.warn(
                 f"Can't read file with requested columns/filters for {load_sample}!",
                 stacklevel=1,
             )
+            print(e)
             continue
 
         # no events?