Merge pull request #11 from LPC-HH/main

Update

Author: ludomori99

.github/workflows/ci.yml

@@ -50,6 +50,6 @@ jobs:
           --durations=20
 
       - name: Upload coverage report
-        uses: codecov/codecov-action@v5.3.1
+        uses: codecov/codecov-action@v5.4.2
         with:
           token: ${{ secrets.CODECOV_TOKEN }}

src/boostedhh/plotting.py

@@ -550,7 +550,7 @@ def ratioHistPlot(
     if ylim is not None:
         ax.set_ylim([y_lowlim, ylim])
     else:
-        ax.set_ylim(y_lowlim)
+        ax.set_ylim(y_lowlim, ax.get_ylim()[1] * 2)
 
     ax.margins(x=0)
 
@@ -787,7 +787,7 @@ def _find_nearest(array, value):
     return idx
 
 
-def multiROCCurveGrey(
+def multiROCCurveGreyOld(
     rocs: dict,
     sig_effs: list[float],
     plot_dir: Path,
@@ -858,6 +858,162 @@ def multiROCCurveGrey(
 ]
 
 
+def multiROCCurveGrey(
+    rocs: dict,
+    sig_effs: list[float] = None,
+    bkg_effs: list[float] = None,
+    xlim=None,
+    ylim=None,
+    plot_dir: Path = None,
+    name: str = "",
+    show: bool = False,
+    add_cms_label=False,
+    legtitle: str = None,
+    title: str = None,
+    plot_thresholds: dict = None,  # plot signal and bkg efficiency for a given discriminator threshold
+    find_from_sigeff: dict = None,  # find discriminator threshold that matches signal efficiency
+):
+    """Plot multiple ROC curves (e.g. train and test) + multiple signals"""
+    if ylim is None:
+        ylim = [1e-06, 1]
+    if xlim is None:
+        xlim = [0, 1]
+    line_style = {"colors": "lightgrey", "linestyles": "dashed"}
+    th_colours = ["cornflowerblue", "deepskyblue", "mediumblue", "cyan", "cadetblue"]
+    eff_colours = ["lime", "aquamarine", "greenyellow"]
+
+    fig = plt.figure(figsize=(12, 12))
+    ax = fig.gca()
+    for roc_sigs in rocs.values():
+
+        # plots roc curves for each type of signal
+        for roc in roc_sigs.values():
+
+            plt.plot(
+                roc["tpr"],
+                roc["fpr"],
+                label=roc["label"],
+                color=roc["color"],
+                linewidth=2,
+            )
+
+            # determines the point on the ROC curve that corresponds to the signal efficiency
+            # plots a vertical and horizontal line to the point
+            if sig_effs is not None:
+                for sig_eff in sig_effs:
+                    y = roc["fpr"][np.searchsorted(roc["tpr"], sig_eff)]
+                    plt.hlines(y=y, xmin=0, xmax=sig_eff, **line_style)
+                    plt.vlines(x=sig_eff, ymin=0, ymax=y, **line_style)
+
+            # determines the point on the ROC curve that corresponds to the background efficiency
+            # plots a vertical and horizontal line to the point
+            if bkg_effs is not None:
+                for bkg_eff in bkg_effs:
+                    x = roc["tpr"][np.searchsorted(roc["fpr"], bkg_eff)]
+                    plt.vlines(x=x, ymin=0, ymax=bkg_eff, **line_style)
+                    plt.hlines(y=bkg_eff, xmin=0, xmax=x, **line_style)
+
+    # plots points and lines on plot corresponding to classifier thresholds
+    for roc_sigs in rocs.values():
+        if plot_thresholds is None:
+            break
+        i_sigeff = 0
+        i_th = 0
+        for rockey, roc in roc_sigs.items():
+            if rockey in plot_thresholds:
+                pths = {th: [[], []] for th in plot_thresholds[rockey]}
+                for th in plot_thresholds[rockey]:
+                    idx = _find_nearest(roc["thresholds"], th)
+                    pths[th][0].append(roc["tpr"][idx])
+                    pths[th][1].append(roc["fpr"][idx])
+                for th in plot_thresholds[rockey]:
+                    plt.scatter(
+                        *pths[th],
+                        marker="o",
+                        s=40,
+                        label=rf"{rockey} > {th:.2f}",
+                        zorder=100,
+                        color=th_colours[i_th],
+                    )
+                    plt.vlines(
+                        x=pths[th][0],
+                        ymin=0,
+                        ymax=pths[th][1],
+                        color=th_colours[i_th],
+                        linestyles="dashed",
+                        alpha=0.5,
+                    )
+                    plt.hlines(
+                        y=pths[th][1],
+                        xmin=0,
+                        xmax=pths[th][0],
+                        color=th_colours[i_th],
+                        linestyles="dashed",
+                        alpha=0.5,
+                    )
+                    i_th += 1
+
+            if find_from_sigeff is not None and rockey in find_from_sigeff:
+                pths = {sig_eff: [[], []] for sig_eff in find_from_sigeff[rockey]}
+                thrs = {}
+                for sig_eff in find_from_sigeff[rockey]:
+                    idx = _find_nearest(roc["tpr"], sig_eff)
+                    thrs[sig_eff] = roc["thresholds"][idx]
+                    pths[sig_eff][0].append(roc["tpr"][idx])
+                    pths[sig_eff][1].append(roc["fpr"][idx])
+                for sig_eff in find_from_sigeff[rockey]:
+                    plt.scatter(
+                        *pths[sig_eff],
+                        marker="o",
+                        s=40,
+                        label=rf"{rockey} > {thrs[sig_eff]:.2f}",
+                        zorder=100,
+                        color=eff_colours[i_sigeff],
+                    )
+                    plt.vlines(
+                        x=pths[sig_eff][0],
+                        ymin=0,
+                        ymax=pths[sig_eff][1],
+                        color=eff_colours[i_sigeff],
+                        linestyles="dashed",
+                        alpha=0.5,
+                    )
+                    plt.hlines(
+                        y=pths[sig_eff][1],
+                        xmin=0,
+                        xmax=pths[sig_eff][0],
+                        color=eff_colours[i_sigeff],
+                        linestyles="dashed",
+                        alpha=0.5,
+                    )
+                    i_sigeff += 1
+
+    if add_cms_label:
+        hep.cms.label(data=False, rlabel="")
+    if title:
+        plt.title(title)
+    plt.yscale("log")
+    plt.xlabel("Signal efficiency")
+    plt.ylabel("Background efficiency")
+    plt.xlim(*xlim)
+    plt.ylim(*ylim)
+    ax.xaxis.grid(True, which="major")
+    ax.yaxis.grid(True, which="major")
+    if legtitle:
+        plt.legend(title=legtitle, loc="center left", bbox_to_anchor=(1, 0.5))
+    else:
+        plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+
+    if len(name):
+        plt.savefig(plot_dir / f"{name}.png", bbox_inches="tight")
+        plt.savefig(plot_dir / f"{name}.pdf", bbox_inches="tight")
+
+    if show:
+        plt.show()
+    else:
+        plt.close()
+
+
 def multiROCCurve(
     rocs: dict,
     thresholds=None,

src/boostedhh/run_utils.py

@@ -109,7 +109,8 @@ def parse_common_hh_args(parser):
         "--year",
         help="year",
         type=str,
-        default="2022",
+        nargs="+",
+        required=True,
         choices=["2018", "2022", "2022EE", "2023", "2023BPix"],
     )
 

src/boostedhh/submit_utils.py

@@ -133,6 +133,13 @@ def check_branch(
 def init_args(args):
     # check that branch exists
     check_branch(args.analysis, args.git_branch, args.git_user, args.allow_diff_local_repo)
+
+    if isinstance(args.year, list):
+        if len(args.year) == 1:
+            args.year = args.year[0]
+        else:
+            raise ValueError("Submitting multiple years without --yaml option is not supported yet")
+
     username = os.environ["USER"]
 
     if args.site == "lpc":
@@ -189,8 +196,8 @@ def submit(
 
     # submit jobs
     nsubmit = 0
-    for sample in fileset:
-        for subsample, tot_files in fileset[sample].items():
+    for sample, sfiles in fileset.items():
+        for subsample, tot_files in sfiles.items():
             if args.submit:
                 print("Submitting " + subsample)
 
@@ -239,7 +246,7 @@ def submit(
                     Path(f"{localcondor}.log").unlink()
 
                 if args.submit:
-                    os.system("condor_submit %s" % localcondor)
+                    os.system(f"condor_submit {localcondor}")
                 else:
                     print("To submit ", localcondor)
                 nsubmit = nsubmit + 1
@@ -250,14 +257,14 @@ def submit(
 def replace_batch_size(file_path: Path, new_batch_size: int):
     """Replacing batch size in given file"""
     import re
-    
+
     # Read the file
-    with file_path.open('r') as file:
+    with file_path.open("r") as file:
         content = file.read()
-    
+
     # Replace using regex
-    updated_content = re.sub(r'--batch-size \d+', f'--batch-size {new_batch_size}', content)
-    
+    updated_content = re.sub(r"--batch-size \d+", f"--batch-size {new_batch_size}", content)
+
     # Write back to the file
-    with open(file_path, 'w') as file:
-        file.write(updated_content)
+    with file_path.open("w") as file:
+        file.write(updated_content)