diff --git a/BDT/BDT_training.py b/BDT/BDT_training.py new file mode 100644 index 0000000..d5d490e --- /dev/null +++ b/BDT/BDT_training.py @@ -0,0 +1,541 @@ +import glob +import re +import uproot +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import os +from pathlib import Path + +from sklearn.model_selection import train_test_split +from xgboost import XGBClassifier +from sklearn.metrics import roc_auc_score, roc_curve + +# --------------------------------------------------------------------------- +# Paths — flat tuples produced by fillTuplesScouting.py +# --------------------------------------------------------------------------- +tuples_dir = Path("/home/users/garciaja/fullRun3/CMSSW_15_0_2/src/run3_scouting/tuples") + +use_conditional = False +model = 'A' + +# --------------------------------------------------------------------------- +# Discover all signal files and parse (mpi, mA, ctau) from filenames +# --------------------------------------------------------------------------- +_SIG_RE = re.compile( + r"tuples_Signal_ScenarioA_Par_2024_mpi-(\w+)_mA-(\w+)_ctau-(\w+)mm_2024\.root" +) + +def _p2f(s): + return float(s.replace("p", ".")) + +def _flabel(f): + s = f"{f:g}" + return s.replace(".", "p") + +MASS_COLORS = ["#d62728", "#ff7f0e", "#2ca02c", "#1f77b4", "#e377c2"] +BKG_FACE = "#7fc7c4" +BKG_EDGE = "#2f5f5d" + +sig_file_params = [] +for fpath in sorted(glob.glob(str(tuples_dir / "tuples_Signal_ScenarioA_Par_2024_*.root"))): + m = _SIG_RE.search(os.path.basename(fpath)) + if m: + mpi_s, mA_s, ctau_s = m.groups() + sig_file_params.append((fpath, _p2f(mpi_s), _p2f(mA_s), _p2f(ctau_s))) + +ctau_values = sorted(set(p[3] for p in sig_file_params)) +param_grid = [(p[3], p[2], p[1]) for p in sig_file_params] # (ctau, mA, mpi) +print(f"Found {len(sig_file_params)} signal files, ctau values: {ctau_values}") + +# --------------------------------------------------------------------------- +# BDT input variables — comment out variables to exclude from training +# Mirrors the selection in plotsignalvsbkg.py +# --------------------------------------------------------------------------- +def _make_bdt_vars(): + sv_stems = [ + # "chi2", + "chi2Ndof", + "d3d_mumu_SV", + "dphi_mumu_SV", + "l3d", + "lxy", + # "maxd3d", + # "maxdx", + # "maxdxy", + # "maxdy", + # "maxdz", + # "minDistanceFromDet", + # "minDistanceFromDet_x", + # "minDistanceFromDet_y", + # "minDistanceFromDet_z", + # "mind3d", + # "mindx", + # "mindxy", + # "mindy", + # "mindz", + # "ndof", + # "onModule", + # "onModuleWithinUnc", + "prob", + "ptmm", + "x", + "xErr", + "y", + "yErr", + "z", + "zErr", + "dr_mumu", + "dphi_mumu", + "deta_mumu", + "deta_mumu_SV", + "sindphi_lxy", + # "closestDet_x", + # "closestDet_y", + # "closestDet_z", + "a3d_mumu" + ] + mu_stems = [ + "dxy", + # "dxyErr", + "dxysig", + "dxy_lxy", + "dz", + # "dze", + "dzsig", + "ecalIso", + "ecalRelIso", + "eta", + "hcalIso", + "hcalRelIso", + "isGlobal", + # "isStandAlone", + "isTracker", + "isvtx", + "maxdr", + # "mindetaJet", + # "mindphiJet", + "mindr", + # "mindrJet", + # "mindrPF0p3", + # "mindrPF0p4", + "muCSCDT", + "muChambs", + # "muExpMatchedStats", + "muHits", + # "muMatch", + # "muMatchedRPC", + # "muMatchedStats", + # "ncompatible", + # "ncompatibletotal", + # "nexpectedhits", + # "nexpectedhitsmultiple", + # "nexpectedhitsmultipletotal", + # "nexpectedhitstotal", + "nhitsbeforesv", + "normChi2", + # "phi", + "phiCorr", + "pixHits", + "pixLayers", + "pt", + # "saHits", + # "saMatchedStats", + "stripHits", + "trackIso", + "trackRelIso", + "trkLayers", + # "PFIsoChg0p3", + "PFIsoAll0p3", + # "PFIsoChg0p4", + # "PFIsoAll0p4", + # "PFRelIsoChg0p3", + "PFRelIsoAll0p3", + # "PFRelIsoChg0p4", + # "PFRelIsoAll0p4", + ] + vars_ = [] + for sv in ("SV1", "SV2"): + for s in sv_stems: + vars_.append(f"{sv}_{s}") + for mu in ("mu1", "mu2"): + for s in mu_stems: + vars_.append(f"{sv}_{mu}_{s}") + return vars_ + +BDT_VARIABLES = _make_bdt_vars() +_LOAD_BRANCHES = list(dict.fromkeys( + BDT_VARIABLES + ["SV1_lxy", "SV1_mass", "SV2_mass"] +)) + +BKG_FILES = [ + "tuples_QCD_Bin-PT-15to20_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-20to30_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-30to50_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-50to80_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-80to120_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-120to170_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-170to300_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-300to470_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-470to600_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-600to800_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-800to1000_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-1000_Fil-MuEnriched_2024_2024.root", +] + + +def read_flat(path): + with uproot.open(path) as f: + t = f['tuples'] + available = set(t.keys()) + branches = [b for b in _LOAD_BRANCHES if b in available] + return t.arrays(branches, library='pd') + + +def compute_class_weights(y): + n_sig = (y == 1).sum() + n_bkg = (y == 0).sum() + return np.where(y == 1, n_bkg / n_sig, 1.0) + + +# --------------------------------------------------------------------------- +# Load signal — all mass/mpi/ctau points +# --------------------------------------------------------------------------- +print('Making signal dataframes...') +sig_frames = [] +for fpath, mpi_val, mA_val, ctau_val in sig_file_params: + if not Path(fpath).exists(): + continue + df = read_flat(fpath) + df['param_ctau'] = float(ctau_val) + df['param_mA'] = float(mA_val) + df['param_mpi'] = float(mpi_val) + df['label'] = 1 + sig_frames.append(df) + print(f' sig mpi={mpi_val} mA={mA_val} ctau={ctau_val}: {len(df)} events') + +df_sig = pd.concat(sig_frames, ignore_index=True) + +# --------------------------------------------------------------------------- +# Load background +# --------------------------------------------------------------------------- +print('Making background dataframes...') +bkg_frames = [] +for fname in BKG_FILES: + fpath = tuples_dir / fname + if not fpath.exists(): + continue + df = read_flat(fpath) + df['label'] = 0 + bkg_frames.append(df) + print(f' bkg {fname}: {len(df)} events') + +df_bkg = pd.concat(bkg_frames, ignore_index=True) + + +def add_dxy_lxy(df): + for sv in ("SV1", "SV2"): + denom = df[f"{sv}_lxy"] * df[f"{sv}_mass"] / df[f"{sv}_ptmm"] + denom = np.where(denom > 1e-9, denom, 1e-9) + for mu in ("mu1", "mu2"): + df[f"{sv}_{mu}_dxy_lxy"] = np.abs(df[f"{sv}_{mu}_dxy"]) / denom + + +add_dxy_lxy(df_sig) +add_dxy_lxy(df_bkg) + +# --------------------------------------------------------------------------- +# Lxy binning +# --------------------------------------------------------------------------- +lxy_bins = [0.0, 0.2, 1.0, 2.4, 3.1, 7.0, 11.0, 16.0, 70.0] +lxy_labels = ["0p0to0p2", "0p2to1p0", "1p0to2p4", "2p4to3p1", + "3p1to7p0", "7p0to11p0", "11p0to16p0", "16p0to70p0"] + +df_sig['lxy_bin'] = pd.cut(df_sig['SV1_lxy'], bins=lxy_bins, labels=lxy_labels, include_lowest=True) +df_bkg['lxy_bin'] = pd.cut(df_bkg['SV1_lxy'], bins=lxy_bins, labels=lxy_labels, include_lowest=True) + +available = set(df_sig.columns) & set(df_bkg.columns) +input_vars = [v for v in BDT_VARIABLES if v in available] +missing = [v for v in BDT_VARIABLES if v not in available] +if missing: + print(f'\nWARNING: {len(missing)} BDT variables not in tuples (regenerate with fillTuplesScouting.py?):') + for v in missing: + print(f' {v}') +print(f'\nBDT input variables ({len(input_vars)}):') +for v in input_vars: + print(f' {v}') + +cond_vars = ['param_ctau', 'param_mA', 'param_mpi'] if use_conditional else [] + +cond_tag = 'conditional' if use_conditional else 'non_cond' +out_dir = f'BDT/curves_OR_{model}_{cond_tag}' +os.makedirs(f'{out_dir}/binned_lxy', exist_ok=True) +os.makedirs(f'{out_dir}/global', exist_ok=True) + +# --------------------------------------------------------------------------- +# Training loop — one BDT per lxy bin +# --------------------------------------------------------------------------- +print('Training...') +for bin_label in lxy_labels: + sig_bin = df_sig[df_sig['lxy_bin'] == bin_label].copy() + bkg_bin = df_bkg[df_bkg['lxy_bin'] == bin_label].copy() + + if len(bkg_bin) < 100 or len(sig_bin) < 10: + continue + + if use_conditional: + # Replicate background once per param point so conditional vars are set + all_combined = [] + for ctau_val, mA_val, mpi_val in param_grid: + sig_pt = sig_bin[ + (sig_bin['param_ctau'] == ctau_val) & + (sig_bin['param_mA'] == mA_val) & + (sig_bin['param_mpi'] == mpi_val) + ].copy() + bkg_cp = bkg_bin.copy() + bkg_cp['param_ctau'] = float(ctau_val) + bkg_cp['param_mA'] = float(mA_val) + bkg_cp['param_mpi'] = float(mpi_val) + all_combined.append(pd.concat([sig_pt, bkg_cp], ignore_index=True)) + df_combined = pd.concat(all_combined, ignore_index=True) + else: + # Non-conditional: all signal + background combined once + print(f' Combining datasets for lxy={bin_label}...') + df_combined = pd.concat([sig_bin, bkg_bin], ignore_index=True) + + X = df_combined[input_vars + cond_vars] + y = df_combined['label'] + w = compute_class_weights(y.values) + + print(f' Training lxy={bin_label} ({len(sig_bin)} sig, {len(bkg_bin)} bkg)...') + X_train, X_test, y_train, y_test, w_train, w_test = train_test_split( + X, y, w, test_size=0.3, random_state=42, stratify=y) + + bdt = XGBClassifier( + n_estimators=100, max_depth=3, learning_rate=0.1, + use_label_encoder=False, eval_metric='logloss', + tree_method='hist', n_jobs=4, + ) + bdt.fit(X_train, y_train, sample_weight=w_train) + + # Feature importance + print(' Feature importance...') + importances = bdt.feature_importances_ + mask = importances >= 0.001 + imp_filt = importances[mask] + col_filt = np.array(X.columns)[mask] + sorted_idx = np.argsort(imp_filt) + fig, ax = plt.subplots(figsize=(8, max(4, len(col_filt) * 0.3)), constrained_layout=True) + ax.barh(col_filt[sorted_idx], imp_filt[sorted_idx]) + ax.set_title(f'Feature Importance (Lxy={bin_label})') + ax.tick_params(axis='y', labelsize=8) + fig.savefig(f'{out_dir}/binned_lxy/FeatImp_lxy_{bin_label}.png', dpi=150, bbox_inches='tight') + plt.close(fig) + + # ROC curves — one per (mpi, mA), lines = ctau values + df_test_meta = df_combined.loc[X_test.index] + bkg_mask_roc = df_test_meta['label'] == 0 + + mpi_mA_groups = {} + for ctau_val, mA_val, mpi_val in param_grid: + mpi_mA_groups.setdefault((mpi_val, mA_val), []).append(ctau_val) + + for (mpi_val, mA_val), ctau_vals in sorted(mpi_mA_groups.items()): + plot_dir = f'{out_dir}/binned_lxy/mpi{_flabel(mpi_val)}/mA_{_flabel(mA_val)}' + os.makedirs(plot_dir, exist_ok=True) + fig, ax = plt.subplots(figsize=(6, 6), constrained_layout=True) + ax.plot([0, 1], [0, 1], 'k--', alpha=0.4, linewidth=1.0) + + for i, ctau_val in enumerate(sorted(ctau_vals)): + sig_mask = ( + (df_test_meta['param_ctau'] == float(ctau_val)) & + (df_test_meta['param_mA'] == float(mA_val)) & + (df_test_meta['param_mpi'] == float(mpi_val)) + ) + eval_mask = sig_mask | bkg_mask_roc if not use_conditional else sig_mask + X_eval = X_test[eval_mask] + y_eval = y_test[eval_mask] + if len(y_eval) < 10 or y_eval.nunique() < 2: + continue + y_pred = bdt.predict_proba(X_eval)[:, 1] + fpr, tpr, _ = roc_curve(y_eval, y_pred) + auc = roc_auc_score(y_eval, y_pred) + ax.plot(fpr, tpr, color=MASS_COLORS[i % len(MASS_COLORS)], linewidth=1.6, + label=rf'$c\tau={ctau_val:g}$ mm (AUC={auc:.3f})') + + ax.set_xlabel('False Positive Rate') + ax.set_ylabel('True Positive Rate') + ax.set_title(f'ROC mpi={mpi_val} mA={mA_val} lxy={bin_label}') + ax.legend(loc='lower right', fontsize=9, framealpha=0.9) + fig.savefig( + f'{plot_dir}/ROC_mpi{_flabel(mpi_val)}_mA{_flabel(mA_val)}_lxy_{bin_label}.png', + dpi=150, bbox_inches='tight', + ) + plt.close(fig) + + # Grouped discriminant plots — one per (mpi, mA), lines = ctau values + bkg_disc_scores = bdt.predict_proba(X_test[bkg_mask_roc])[:, 1] + + disc_bins = np.linspace(0, 1, 51) + widths = np.diff(disc_bins) + for (mpi_val, mA_val), ctau_vals in sorted(mpi_mA_groups.items()): + plot_dir = f'{out_dir}/binned_lxy/mpi{_flabel(mpi_val)}/mA_{_flabel(mA_val)}' + os.makedirs(plot_dir, exist_ok=True) + fig, ax = plt.subplots(figsize=(7.2, 5.6), constrained_layout=True) + + hb, _ = np.histogram(bkg_disc_scores, bins=disc_bins) + if hb.sum() > 0: hb = hb / hb.sum() + ax.bar(disc_bins[:-1], hb, width=widths, align='edge', + color=BKG_FACE, edgecolor=BKG_EDGE, linewidth=0.6, + label='Background', zorder=1) + + for i, ctau_val in enumerate(sorted(ctau_vals)): + sig_mask = ( + (df_test_meta['param_ctau'] == float(ctau_val)) & + (df_test_meta['param_mA'] == float(mA_val)) & + (df_test_meta['param_mpi'] == float(mpi_val)) + ) + if sig_mask.sum() < 5: + continue + hs, _ = np.histogram(bdt.predict_proba(X_test[sig_mask])[:, 1], bins=disc_bins) + if hs.sum() > 0: hs = hs / hs.sum() + ax.stairs(hs, disc_bins, color=MASS_COLORS[i % len(MASS_COLORS)], linewidth=1.6, + label=rf'$c\tau={ctau_val:g}$ mm', zorder=3 + i) + + ax.set_xlabel('BDT score') + ax.set_ylabel('a.u.') + ax.set_title(rf'Discriminant $m_\pi={mpi_val:g}$ GeV, $m_A={mA_val:g}$ GeV, lxy={bin_label}') + ax.legend(loc='upper center', fontsize=9, framealpha=0.9) + fig.savefig( + f'{plot_dir}/Disc_mpi{_flabel(mpi_val)}_mA{_flabel(mA_val)}_lxy_{bin_label}.png', + dpi=150, bbox_inches='tight', + ) + plt.close(fig) + +# --------------------------------------------------------------------------- +# Global BDT — trained on all lxy bins combined +# --------------------------------------------------------------------------- +print('Training global BDT...') +if use_conditional: + all_combined_global = [] + for ctau_val, mA_val, mpi_val in param_grid: + sig_pt = df_sig[ + (df_sig['param_ctau'] == ctau_val) & + (df_sig['param_mA'] == mA_val) & + (df_sig['param_mpi'] == mpi_val) + ].copy() + bkg_cp = df_bkg.copy() + bkg_cp['param_ctau'] = float(ctau_val) + bkg_cp['param_mA'] = float(mA_val) + bkg_cp['param_mpi'] = float(mpi_val) + all_combined_global.append(pd.concat([sig_pt, bkg_cp], ignore_index=True)) + df_global = pd.concat(all_combined_global, ignore_index=True) +else: + df_global = pd.concat([df_sig, df_bkg], ignore_index=True) + +X_g = df_global[input_vars + cond_vars] +y_g = df_global['label'] +w_g = compute_class_weights(y_g.values) + +X_train_g, X_test_g, y_train_g, y_test_g, w_train_g, w_test_g = train_test_split( + X_g, y_g, w_g, test_size=0.3, random_state=42, stratify=y_g) + +bdt_global = XGBClassifier( + n_estimators=100, max_depth=3, learning_rate=0.1, + use_label_encoder=False, eval_metric='logloss', + tree_method='hist', n_jobs=4, +) +bdt_global.fit(X_train_g, y_train_g, sample_weight=w_train_g) + +# Feature importance +importances_g = bdt_global.feature_importances_ +mask_g = importances_g >= 0.001 +imp_filt_g = importances_g[mask_g] +col_filt_g = np.array(X_g.columns)[mask_g] +sorted_idx_g = np.argsort(imp_filt_g) +fig, ax = plt.subplots(figsize=(8, max(4, len(col_filt_g) * 0.3)), constrained_layout=True) +ax.barh(col_filt_g[sorted_idx_g], imp_filt_g[sorted_idx_g]) +ax.set_title('Feature Importance (Global, all lxy)') +ax.tick_params(axis='y', labelsize=8) +fig.savefig(f'{out_dir}/global/FeatImp_global.png', dpi=150, bbox_inches='tight') +plt.close(fig) + +# ROC curves — global BDT, one per (mpi, mA), lines = ctau values +df_test_meta_g = df_global.loc[X_test_g.index] +bkg_mask_roc_g = df_test_meta_g['label'] == 0 + +mpi_mA_groups_g = {} +for ctau_val, mA_val, mpi_val in param_grid: + mpi_mA_groups_g.setdefault((mpi_val, mA_val), []).append(ctau_val) + +for (mpi_val, mA_val), ctau_vals in sorted(mpi_mA_groups_g.items()): + plot_dir_g = f'{out_dir}/global/mpi{_flabel(mpi_val)}/mA_{_flabel(mA_val)}' + os.makedirs(plot_dir_g, exist_ok=True) + fig, ax = plt.subplots(figsize=(6, 6), constrained_layout=True) + ax.plot([0, 1], [0, 1], 'k--', alpha=0.4, linewidth=1.0) + + for i, ctau_val in enumerate(sorted(ctau_vals)): + sig_mask = ( + (df_test_meta_g['param_ctau'] == float(ctau_val)) & + (df_test_meta_g['param_mA'] == float(mA_val)) & + (df_test_meta_g['param_mpi'] == float(mpi_val)) + ) + eval_mask = sig_mask | bkg_mask_roc_g if not use_conditional else sig_mask + X_eval_g = X_test_g[eval_mask] + y_eval_g = y_test_g[eval_mask] + if len(y_eval_g) < 10 or y_eval_g.nunique() < 2: + continue + y_pred_g = bdt_global.predict_proba(X_eval_g)[:, 1] + fpr_g, tpr_g, _ = roc_curve(y_eval_g, y_pred_g) + auc_g = roc_auc_score(y_eval_g, y_pred_g) + ax.plot(fpr_g, tpr_g, color=MASS_COLORS[i % len(MASS_COLORS)], linewidth=1.6, + label=rf'$c\tau={ctau_val:g}$ mm (AUC={auc_g:.3f})') + + ax.set_xlabel('False Positive Rate') + ax.set_ylabel('True Positive Rate') + ax.set_title(f'ROC Global mpi={mpi_val} mA={mA_val}') + ax.legend(loc='lower right', fontsize=9, framealpha=0.9) + fig.savefig( + f'{plot_dir_g}/ROC_global_mpi{_flabel(mpi_val)}_mA{_flabel(mA_val)}.png', + dpi=150, bbox_inches='tight', + ) + plt.close(fig) + +# Grouped discriminant plots — global BDT, one per (mpi, mA), lines = ctau values +bkg_disc_scores_g = bdt_global.predict_proba(X_test_g[bkg_mask_roc_g])[:, 1] + +disc_bins = np.linspace(0, 1, 51) +widths = np.diff(disc_bins) +for (mpi_val, mA_val), ctau_vals in sorted(mpi_mA_groups_g.items()): + plot_dir_g = f'{out_dir}/global/mpi{_flabel(mpi_val)}/mA_{_flabel(mA_val)}' + os.makedirs(plot_dir_g, exist_ok=True) + fig, ax = plt.subplots(figsize=(7.2, 5.6), constrained_layout=True) + + hb_g, _ = np.histogram(bkg_disc_scores_g, bins=disc_bins) + if hb_g.sum() > 0: hb_g = hb_g / hb_g.sum() + ax.bar(disc_bins[:-1], hb_g, width=widths, align='edge', + color=BKG_FACE, edgecolor=BKG_EDGE, linewidth=0.6, + label='Background', zorder=1) + + for i, ctau_val in enumerate(sorted(ctau_vals)): + sig_mask = ( + (df_test_meta_g['param_ctau'] == float(ctau_val)) & + (df_test_meta_g['param_mA'] == float(mA_val)) & + (df_test_meta_g['param_mpi'] == float(mpi_val)) + ) + if sig_mask.sum() < 5: + continue + hs_g, _ = np.histogram(bdt_global.predict_proba(X_test_g[sig_mask])[:, 1], bins=disc_bins) + if hs_g.sum() > 0: hs_g = hs_g / hs_g.sum() + ax.stairs(hs_g, disc_bins, color=MASS_COLORS[i % len(MASS_COLORS)], linewidth=1.6, + label=rf'$c\tau={ctau_val:g}$ mm', zorder=3 + i) + + ax.set_xlabel('BDT score') + ax.set_ylabel('a.u.') + ax.set_title(rf'Discriminant Global $m_\pi={mpi_val:g}$ GeV, $m_A={mA_val:g}$ GeV') + ax.legend(loc='upper center', fontsize=9, framealpha=0.9) + fig.savefig( + f'{plot_dir_g}/Disc_global_mpi{_flabel(mpi_val)}_mA{_flabel(mA_val)}.png', + dpi=150, bbox_inches='tight', + ) + plt.close(fig) diff --git a/BDT/cutncount.py b/BDT/cutncount.py new file mode 100644 index 0000000..076030b --- /dev/null +++ b/BDT/cutncount.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +"""Cut-and-count: signal efficiency vs background rejection in lxy bins. + +Ported from the old Vtx-collection ntuples to the new flat tuples format. +Key changes vs old version: + - Tree is 'tuples' (was 'tout') + - Branches are flat scalars, no jagged max_or_default needed + - Muon vars live on the SV object: SV1_mu1_pt, SV1_mu1_dxysig, etc. + - dphi: SV1_dphi_mumu_SV (was SV1_dphi_Vtx) + - chi2/ndof: SV1_chi2Ndof already computed (was SV1_chi2/SV1_ndof) + - dxysig: SV1_mu1_dxysig already computed + - SV1_3Dangle has no equivalent in new ntuples — cut dropped +""" +import glob +import numpy as np +import uproot +from pathlib import Path + +TUPLES_DIR = Path(__file__).resolve().parent.parent / "tuples" +TREE_NAME = "tuples" + +# One benchmark signal — edit as needed +SIG_PATTERN = str(TUPLES_DIR / "tuples_Signal_ScenarioA_Par_2024_mpi-2_mA-0p50_ctau-10mm_2024.root") +BKG_PATTERN = str(TUPLES_DIR / "tuples_QCD_*.root") + +SIG_FILES = sorted(glob.glob(SIG_PATTERN)) +BKG_FILES = sorted(glob.glob(BKG_PATTERN)) + +BRANCHES = [ + "SV1_lxy", "SV1_chi2Ndof", "SV1_prob", "SV1_mass", "SV1_ptmm", + "SV1_xErr", "SV1_yErr", "SV1_zErr", + "SV1_dphi_mumu_SV", + "SV1_mu1_dxysig", "SV1_mu1_dxy", "SV1_mu1_normChi2", "SV1_mu1_phi", "SV1_mu1_eta", + "SV1_mu1_nhitsbeforesv", + "SV1_mu2_dxysig", "SV1_mu2_dxy", "SV1_mu2_normChi2", "SV1_mu2_phi", "SV1_mu2_eta", + "SV1_mu2_nhitsbeforesv", + "SV1_minDistanceFromDet_x", "SV1_minDistanceFromDet_y", "SV1_minDistanceFromDet_z", + + "SV2_lxy", "SV2_chi2Ndof", "SV2_prob", "SV2_mass", "SV2_ptmm", + "SV2_xErr", "SV2_yErr", "SV2_zErr", + "SV2_dphi_mumu_SV", + "SV2_mu1_dxysig", "SV2_mu1_dxy", "SV2_mu1_normChi2", "SV2_mu1_phi", "SV2_mu1_eta", + "SV2_mu1_nhitsbeforesv", + "SV2_mu2_dxysig", "SV2_mu2_dxy", "SV2_mu2_normChi2", "SV2_mu2_phi", "SV2_mu2_eta", + "SV2_mu2_nhitsbeforesv", + "SV2_minDistanceFromDet_x", "SV2_minDistanceFromDet_y", "SV2_minDistanceFromDet_z", +] + +LXY_BINS = [0.0, 0.2, 1.0, 2.4, 3.1, 7.0, 11.0, 16.0, 70.0] +LXY_LABELS = ["0p0-0p2", "0p2-1p0", "1p0-2p4", "2p4-3p1", + "3p1-7p0", "7p0-11p0", "11p0-16p0", "16p0-70p0"] + + +def load_arrays(files, label): + arrays = {b: [] for b in BRANCHES} + for path in files: + with uproot.open(path) as f: + t = f[TREE_NAME] + for b in BRANCHES: + arrays[b].append(t[b].array(library="np")) + arr = {b: np.concatenate(arrays[b]) for b in BRANCHES} + arr["label"] = np.full(len(arr[BRANCHES[0]]), label, dtype=np.int8) + return arr + + +def sv_mask(arr, sv): + dphi = arr[f"SV{sv}_dphi_mumu_SV"] + xErr = arr[f"SV{sv}_xErr"] + yErr = arr[f"SV{sv}_yErr"] + zErr = arr[f"SV{sv}_zErr"] + lxy = arr[f"SV{sv}_lxy"] + chi2ndof = arr[f"SV{sv}_chi2Ndof"] + return ( + (xErr < 0.05) & + (yErr < 0.05) & + (zErr < 0.10) & + (dphi > 0 ) & + (lxy > 0 ) & + (lxy < 70 ) & + (chi2ndof < 3 ) + ) + + +def dimuon_mask(arr, sv): + dxysig1 = arr[f"SV{sv}_mu1_dxysig"] + dxysig2 = arr[f"SV{sv}_mu2_dxysig"] + dxy1 = arr[f"SV{sv}_mu1_dxy"] + dxy2 = arr[f"SV{sv}_mu2_dxy"] + chi2ndof1 = arr[f"SV{sv}_mu1_normChi2"] + chi2ndof2 = arr[f"SV{sv}_mu2_normChi2"] + phi1 = arr[f"SV{sv}_mu1_phi"] + phi2 = arr[f"SV{sv}_mu2_phi"] + eta1 = arr[f"SV{sv}_mu1_eta"] + eta2 = arr[f"SV{sv}_mu2_eta"] + lxy = arr[f"SV{sv}_lxy"] + mass = arr[f"SV{sv}_mass"] + ptmm = arr[f"SV{sv}_ptmm"] + + dphi = (phi1 - phi2 + np.pi) % (2 * np.pi) - np.pi + dphi = np.where(np.abs(dphi) > 1e-6, dphi, 1e-6) + deta = np.where(np.abs(eta1 - eta2) > 1e-6, eta1 - eta2, 1e-6) + log_ratio = np.log10(np.abs(deta) / np.abs(dphi)) + + denom = lxy * mass / ptmm + dxy_lxy1 = np.abs(dxy1) / np.where(denom > 1e-9, denom, 1e-9) + dxy_lxy2 = np.abs(dxy2) / np.where(denom > 1e-9, denom, 1e-9) + + return ( + (np.abs(dxysig1) > 2 ) & + (np.abs(dxysig2) > 2 ) & + (chi2ndof1 < 3 ) & + (chi2ndof2 < 3 ) & + (dphi < 2.8 ) & + (log_ratio < 1.25) & + (dxy_lxy1 > 0.1 ) & + (dxy_lxy2 > 0.1 ) + ) + + +def material_veto_mask(arr, sv): + dx = arr[f"SV{sv}_minDistanceFromDet_x"] + dy = arr[f"SV{sv}_minDistanceFromDet_y"] + dz = arr[f"SV{sv}_minDistanceFromDet_z"] + return (np.abs(dx) >= 0.81) | (np.abs(dy) >= 3.24) | (np.abs(dz) >= 0.0145) + + +def excess_hits_mask(arr, sv): + lxy = arr[f"SV{sv}_lxy"] + n_excess = arr[f"SV{sv}_mu1_nhitsbeforesv"] + arr[f"SV{sv}_mu2_nhitsbeforesv"] + max_hits = np.where(lxy < 11, 0, np.where(lxy < 16, 1, 2)) + return n_excess <= max_hits + + +def apply_cuts(arr): + pass1 = sv_mask(arr, 1) & dimuon_mask(arr, 1) & excess_hits_mask(arr, 1) & material_veto_mask(arr, 1) + pass2 = sv_mask(arr, 2) & dimuon_mask(arr, 2) & excess_hits_mask(arr, 2) & material_veto_mask(arr, 2) + mask = pass1 | pass2 + return {k: v[mask] for k, v in arr.items()} + + +print("Loading signal...") +sig = load_arrays(SIG_FILES, label=1) +print("Loading background...") +bkg = load_arrays(BKG_FILES, label=0) + +all_events = {k: np.concatenate([sig[k], bkg[k]]) for k in sig} + +lxy = all_events["SV1_lxy"] +labels = all_events["label"] + +print() +print(f"{'lxy bin':>14} | {'sig pass/total':>16} {'sig eff':>8} | {'bkg pass/total':>16} {'bkg rej':>8}") +print("-" * 75) + +for lo, hi, lbl in zip(LXY_BINS[:-1], LXY_BINS[1:], LXY_LABELS): + bin_mask = (lxy >= lo) & (lxy < hi) + arr_bin = {k: v[bin_mask] for k, v in all_events.items()} + arr_pass = apply_cuts(arr_bin) + + sig_total = int(np.sum(arr_bin["label"] == 1)) + bkg_total = int(np.sum(arr_bin["label"] == 0)) + sig_pass = int(np.sum(arr_pass["label"] == 1)) + bkg_pass = int(np.sum(arr_pass["label"] == 0)) + + sig_eff = sig_pass / sig_total if sig_total > 0 else float("nan") + bkg_rej = 1 - bkg_pass / bkg_total if bkg_total > 0 else float("nan") + + print( + f"{lbl:>14} | " + f"{sig_pass:>6}/{sig_total:<8} {sig_eff:>8.3f} | " + f"{bkg_pass:>6}/{bkg_total:<8} {bkg_rej:>8.3f}" + ) + +# --------------------------------------------------------------------------- +# Cutflow — sequential cut breakdown (all lxy bins combined) +# --------------------------------------------------------------------------- +CUT_STAGES = [ + ("SV quality", + lambda a: (sv_mask(a, 1)) + | (sv_mask(a, 2))), + ("+ Dimuon cuts", + lambda a: (sv_mask(a, 1) & dimuon_mask(a, 1)) + | (sv_mask(a, 2) & dimuon_mask(a, 2))), + ("+ Excess hits", + lambda a: (sv_mask(a, 1) & dimuon_mask(a, 1) & excess_hits_mask(a, 1)) + | (sv_mask(a, 2) & dimuon_mask(a, 2) & excess_hits_mask(a, 2))), + ("+ Material veto", + lambda a: (sv_mask(a, 1) & dimuon_mask(a, 1) & excess_hits_mask(a, 1) & material_veto_mask(a, 1)) + | (sv_mask(a, 2) & dimuon_mask(a, 2) & excess_hits_mask(a, 2) & material_veto_mask(a, 2))), +] + +sig_total_all = len(sig["label"]) +bkg_total_all = len(bkg["label"]) + +print() +print("Cutflow (all lxy bins)") +print(f"{'cut':<18} | {'sig pass':>8} {'sig eff':>8} {'sig drop':>9} | {'bkg pass':>9} {'bkg rej':>8} {'bkg drop':>9}") +print("-" * 82) +print(f"{'All events':<18} | {sig_total_all:>8} {'100.0%':>8} {'—':>9} | {bkg_total_all:>9} {'0.0%':>8} {'—':>9}") + +prev_sig = sig_total_all +prev_bkg = bkg_total_all +for name, cut_fn in CUT_STAGES: + sig_pass = int(cut_fn(sig).sum()) + bkg_pass = int(cut_fn(bkg).sum()) + sig_eff = sig_pass / sig_total_all if sig_total_all > 0 else float("nan") + bkg_rej = 1 - bkg_pass / bkg_total_all if bkg_total_all > 0 else float("nan") + sig_drop = (prev_sig - sig_pass) / prev_sig if prev_sig > 0 else float("nan") + bkg_drop = (prev_bkg - bkg_pass) / prev_bkg if prev_bkg > 0 else float("nan") + print(f"{name:<18} | {sig_pass:>8} {sig_eff:>8.1%} {sig_drop:>9.1%} | {bkg_pass:>9} {bkg_rej:>8.1%} {bkg_drop:>9.1%}") + prev_sig = sig_pass + prev_bkg = bkg_pass + +# --------------------------------------------------------------------------- +# Dimuon sub-cut breakdown on signal (SV1, all events) +# --------------------------------------------------------------------------- +print() +print("Dimuon sub-cut breakdown on signal (SV1, standalone — not cumulative)") +print(f"{'cut':<30} | {'sig pass':>8} {'sig eff':>8}") +print("-" * 46) + +def _dimuon_subcuts(arr, sv): + dxysig1 = arr[f"SV{sv}_mu1_dxysig"] + dxysig2 = arr[f"SV{sv}_mu2_dxysig"] + dxy1 = arr[f"SV{sv}_mu1_dxy"] + dxy2 = arr[f"SV{sv}_mu2_dxy"] + chi2ndof1 = arr[f"SV{sv}_mu1_normChi2"] + chi2ndof2 = arr[f"SV{sv}_mu2_normChi2"] + phi1 = arr[f"SV{sv}_mu1_phi"] + phi2 = arr[f"SV{sv}_mu2_phi"] + eta1 = arr[f"SV{sv}_mu1_eta"] + eta2 = arr[f"SV{sv}_mu2_eta"] + lxy = arr[f"SV{sv}_lxy"] + mass = arr[f"SV{sv}_mass"] + ptmm = arr[f"SV{sv}_ptmm"] + dphi = (phi1 - phi2 + np.pi) % (2 * np.pi) - np.pi + dphi = np.where(np.abs(dphi) > 1e-6, dphi, 1e-6) + deta = np.where(np.abs(eta1 - eta2) > 1e-6, eta1 - eta2, 1e-6) + log_ratio = np.log10(np.abs(deta) / np.abs(dphi)) + denom = lxy * mass / ptmm + dxy_lxy1 = np.abs(dxy1) / np.where(denom > 1e-9, denom, 1e-9) + dxy_lxy2 = np.abs(dxy2) / np.where(denom > 1e-9, denom, 1e-9) + return [ + ("|dxysig1| > 2", np.abs(dxysig1) > 2), + ("|dxysig2| > 2", np.abs(dxysig2) > 2), + ("chi2ndof1 < 3", chi2ndof1 < 3 ), + ("chi2ndof2 < 3", chi2ndof2 < 3 ), + ("dphi < 2.8", dphi < 2.8 ), + ("log_ratio < 1.25", log_ratio < 1.25), + ("dxy_lxy1 > 0.1", dxy_lxy1 > 0.1 ), + ("dxy_lxy2 > 0.1", dxy_lxy2 > 0.1 ), + ] + +n_sig = sig_total_all +for sv in (1, 2): + print(f"\n SV{sv}:") + for cut_name, mask in _dimuon_subcuts(sig, sv): + n = int(mask.sum()) + print(f" {cut_name:<28} | {n:>8} {n/n_sig:>8.1%}") diff --git a/BDT/plotsignalvsbkg.py b/BDT/plotsignalvsbkg.py new file mode 100644 index 0000000..2a61b8d --- /dev/null +++ b/BDT/plotsignalvsbkg.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +"""Signal vs background shape comparison for all BDT input variables. + +One set of plots per ctau lifetime. Each plot overlays: + - the combined QCD background (filled, with sqrt(N) uncertainty band) + - one step-line per (mpi, mA) mass point available at that ctau + +Ranges are auto-computed from combined sig+bkg data (1st-99th percentile). +Output: BDT/signal_vs_bkg_ctau{X}mm/{var}.png +""" +import os +import re +import glob +import numpy as np +import uproot +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import mplhep as hep +hep.style.use("CMS") + +# --- Aesthetic overrides on top of the CMS style --------------------------- +# The CMS mplhep style is sized for publication-scale plots; for these many +# small figures we want lighter text and a more compact look. +plt.rcParams.update({ + "font.size": 13, + "axes.labelsize": 13, + "axes.titlesize": 13, + "xtick.labelsize": 11, + "ytick.labelsize": 11, + "legend.fontsize": 9, + "legend.title_fontsize": 10, + "axes.linewidth": 1.0, + "xtick.major.size": 5, + "ytick.major.size": 5, + "xtick.minor.size": 3, + "ytick.minor.size": 3, + "xtick.major.width": 0.9, + "ytick.major.width": 0.9, +}) + +FIGSIZE = (8.5, 6.5) + +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +TUPLES_DIR = str(_HERE.parent / "tuples") +TREE_NAME = "tuples" + +N_BINS = 50 + +# Colors for mass-point overlays (ordered light -> heavy). +MASS_COLORS = ["#d62728", "#ff7f0e", "#2ca02c", "#1f77b4", "#e377c2"] + +# Background colours +BKG_FACE = "#7fc7c4" +BKG_EDGE = "#2f5f5d" + +# --------------------------------------------------------------------------- +# Discover signal files grouped by (mpi, mA) then ctau +# --------------------------------------------------------------------------- +_SIG_RE = re.compile( + r"tuples_Signal_ScenarioA_Par_2024_mpi-(\w+?)_mA-(\w+?)_ctau-(\w+?)mm_2024\.root" +) + +def _p2f(s): + return float(s.replace("p", ".")) + +def _f2lbl(x): + return f"{x:g}".replace(".", "p") + +# mpi_mA_files[(mpi, mA)][ctau] = [files] +mpi_mA_files = {} +for fpath in sorted(glob.glob(os.path.join(TUPLES_DIR, "tuples_Signal_ScenarioA_Par_2024_*.root"))): + m = _SIG_RE.search(os.path.basename(fpath)) + if not m: + continue + mpi = _p2f(m.group(1)) + mA = _p2f(m.group(2)) + ctau = _p2f(m.group(3)) + mpi_mA_files.setdefault((mpi, mA), {}).setdefault(ctau, []).append(fpath) + +mpi_values = sorted(set(k[0] for k in mpi_mA_files)) +print(f"(mpi, mA) points found ({len(mpi_mA_files)}):") +for (mpi, mA), ctau_dict in sorted(mpi_mA_files.items()): + print(f" mpi={mpi}, mA={mA}: ctau={sorted(ctau_dict)}") + +BKG_FILES = sorted(glob.glob(os.path.join(TUPLES_DIR, "tuples_QCD_*.root"))) + +# --------------------------------------------------------------------------- +# Axis label definitions (matplotlib LaTeX) +# --------------------------------------------------------------------------- + +_SV_STEM_LABELS = { + "ptmm": r"$p_{T}^{\mu\mu}$ [GeV]", + #"chi2": r"SV $\chi^{2}$", + "prob": r"SV $\chi^{2}$ probability", + "x": r"SV $x$ [cm]", + "y": r"SV $y$ [cm]", + "z": r"SV $z$ [cm]", + "lxy": r"$l_{xy}$ (from PV) [cm]", + "xErr": r"SV $x$ error [cm]", + "yErr": r"SV $y$ error [cm]", + "zErr": r"SV $z$ error [cm]", + "dphi_mumu_SV": r"$|\Delta\phi(\vec{\mu\mu},\,\vec{SV})|$ [rad]", + "d3d_mumu_SV": r"3D angle$(\vec{\mu\mu},\,\vec{SV})$ [rad]", + "a3d_mumu": r"3D angle$(\mu,\,\mu)$ [rad]", + "mass": r"$m_{\mu\mu}$ [GeV]", + #"ndof": r"SV ndof", + "chi2Ndof": r"SV $\chi^{2}$/ndof", + "l3d": r"$l_{3D}$ (from PV) [cm]", + #"mindx": r"min $D_{x}(SV_{i},SV_{j})$ [cm]", + #"mindy": r"min $D_{y}(SV_{i},SV_{j})$ [cm]", + #"mindz": r"min $D_{z}(SV_{i},SV_{j})$ [cm]", + #"mindxy": r"min $D_{xy}(SV_{i},SV_{j})$ [cm]", + #"mind3d": r"min $D_{3D}(SV_{i},SV_{j})$ [cm]", + #"maxdx": r"max $D_{x}(SV_{i},SV_{j})$ [cm]", + #"maxdy": r"max $D_{y}(SV_{i},SV_{j})$ [cm]", + #"maxdz": r"max $D_{z}(SV_{i},SV_{j})$ [cm]", + #"maxdxy": r"max $D_{xy}(SV_{i},SV_{j})$ [cm]", + #"maxd3d": r"max $D_{3D}(SV_{i},SV_{j})$ [cm]", + #"onModule": r"SV on module", + #"onModuleWithinUnc": r"SV on module (within unc.)", + #"minDistanceFromDet": r"Min. distance to module [cm]", + #"minDistanceFromDet_x": r"Min. distance to module, $x$ [cm]", + #"minDistanceFromDet_y": r"Min. distance to module, $y$ [cm]", + #"minDistanceFromDet_z": r"Min. distance to module, $z$ [cm]", + #"closestDet_x": r"Closest module $x$ [cm]", + #"closestDet_y": r"Closest module $y$ [cm]", + #"closestDet_z": r"Closest module $z$ [cm]", +} + +_MU_STEM_LABELS = { + "pt": r"Muon $p_{T}$ [GeV]", + "eta": r"Muon $\eta$", + #"phi": r"Muon $\phi$ [rad]", + "phiCorr": r"Muon $\phi$ corrected [rad]", + "isvtx": r"Muon is from vertex", + "normChi2": r"Muon $\chi^{2}$/ndof", + "dxy": r"Muon $|d_{xy}|$ [cm]", + #"dxyErr": r"Muon $d_{xy}$ error [cm]", + "dxysig": r"Muon $|d_{xy}|/\sigma_{xy}$", + "dz": r"Muon $|d_{z}|$ [cm]", + #"dze": r"Muon $d_{z}$ error [cm]", + "dzsig": r"Muon $|d_{z}|/\sigma_{z}$", + "nhitsbeforesv": r"Hits before SV", + "isGlobal": r"Muon isGlobal", + "isTracker": r"Muon isTracker", + #"isStandAlone": r"Muon isStandAlone", + "pixHits": r"Pixel hits", + "stripHits": r"Strip hits", + "pixLayers": r"Pixel layers", + "trkLayers": r"Tracker layers", + #"saHits": r"SA muon hits", + #"saMatchedStats": r"SA matched stations", + "muHits": r"Muon hits", + "muChambs": r"Muon chambers", + "muCSCDT": r"CSC/DT chambers", + #"muMatch": r"Muon matches", + #"muMatchedStats": r"Matched stations", + #"muExpMatchedStats": r"Expected matched stations", + #"muMatchedRPC": r"Matched RPC layers", + "ecalIso": r"ECAL isolation [GeV]", + "hcalIso": r"HCAL isolation [GeV]", + "trackIso": r"Track isolation [GeV]", + "ecalRelIso": r"ECAL isolation / $p_{T}$", + "hcalRelIso": r"HCAL isolation / $p_{T}$", + "trackRelIso": r"Track isolation / $p_{T}$", + #"PFIsoChg0p3": r"PF-chg iso. ($\Delta R<0.3$) [GeV]", + "PFIsoAll0p3": r"PF-all iso. ($\Delta R<0.3$) [GeV]", + #"PFRelIsoChg0p3": r"PF-chg rel. iso. ($\Delta R<0.3$)", + "PFRelIsoAll0p3": r"PF-all rel. iso. ($\Delta R<0.3$)", + #"mindrPF0p3": r"min $\Delta R(\mu,\mathrm{PF\ cand.})\ [\Delta R<0.3]$", + #"PFIsoChg0p4": r"PF-chg iso. ($\Delta R<0.4$) [GeV]", + #"PFIsoAll0p4": r"PF-all iso. ($\Delta R<0.4$) [GeV]", + #"PFRelIsoChg0p4": r"PF-chg rel. iso. ($\Delta R<0.4$)", + #"PFRelIsoAll0p4": r"PF-all rel. iso. ($\Delta R<0.4$)", + #"mindrPF0p4": r"min $\Delta R(\mu,\mathrm{PF\ cand.})\ [\Delta R<0.4]$", + "mindr": r"min $\Delta R(\mu_{i},\mu_{j})$", + "maxdr": r"max $\Delta R(\mu_{i},\mu_{j})$", + #"mindrJet": r"min $\Delta R(\mu,\mathrm{PF\ jet})$", + #"mindphiJet": r"$\Delta\phi(\mu,\mathrm{nearest\ PF\ jet})$ [rad]", + #"mindetaJet": r"$\Delta\eta(\mu,\mathrm{nearest\ PF\ jet})$", + #"ncompatible": r"$n_\mathrm{compatible}$", + #"ncompatibletotal": r"$n_\mathrm{compatible\ total}$", + #"nexpectedhits": r"$n_\mathrm{expected\ hits}$", + #"nexpectedhitsmultiple": r"$n_\mathrm{expected\ hits\ (multiple)}$", + #"nexpectedhitsmultipletotal":r"$n_\mathrm{expected\ hits\ (multiple,\ total)}$", + #"nexpectedhitstotal": r"$n_\mathrm{expected\ hits\ total}$", +} + +def _build_axis_labels(): + labels = {} + for sv in ("SV1", "SV2"): + for stem, label in _SV_STEM_LABELS.items(): + labels[f"{sv}_{stem}"] = f"{sv} {label}" + for mu in ("mu1", "mu2"): + for stem, label in _MU_STEM_LABELS.items(): + labels[f"{sv}_{mu}_{stem}"] = f"{sv} {mu} {label}" + return labels + +AXIS_LABELS = _build_axis_labels() + +# --------------------------------------------------------------------------- + +def load_all(files): + arrays = {} + for path in files: + with uproot.open(path) as f: + t = f[TREE_NAME] + for branch in t.keys(): + arr = t[branch].array(library="np") + arrays.setdefault(branch, []).append(arr) + return {k: np.concatenate(v) for k, v in arrays.items()} + + +print("Loading background...") +bkg = load_all(BKG_FILES) + +for mpi in mpi_values: + mA_vals = sorted(set(k[1] for k in mpi_mA_files if k[0] == mpi)) + for mA in mA_vals: + ctau_dict = mpi_mA_files[(mpi, mA)] + ctau_vals = sorted(ctau_dict) + outdir = str(_HERE / f"signal_vs_bkg_mpi{mpi:g}" / f"mA_{mA:g}") + os.makedirs(outdir, exist_ok=True) + + sig_per_ctau = {} + for ctau in ctau_vals: + print(f" Loading mpi={mpi}, mA={mA}, ctau={ctau} ({len(ctau_dict[ctau])} files)...") + sig_per_ctau[ctau] = load_all(ctau_dict[ctau]) + + common = set(bkg.keys()) + for s in sig_per_ctau.values(): + common &= set(s.keys()) + variables = [v for v in AXIS_LABELS if v in common] + print(f" ({len(variables)} vars) -> {outdir}/") + + for var in variables: + all_vals = [bkg[var].astype(float)] + for s in sig_per_ctau.values(): + all_vals.append(s[var].astype(float)) + combined = np.concatenate(all_vals) + combined = combined[np.isfinite(combined)] + if len(combined) == 0: + continue + + lo = np.percentile(combined, 1) + hi = np.percentile(combined, 99) + if lo == hi: + lo, hi = combined.min(), combined.max() + if lo == hi: + continue + + b = bkg[var].astype(float) + hb, edges = np.histogram(b, bins=N_BINS, range=(lo, hi)) + bsum = hb.sum() + hb_norm = hb / bsum if bsum > 0 else hb.astype(float) + + widths = np.diff(edges) + xlabel = AXIS_LABELS.get(var, var) + + fig, ax = plt.subplots(figsize=FIGSIZE) + + ax.bar(edges[:-1], hb_norm, width=widths, align='edge', + color=BKG_FACE, edgecolor=BKG_EDGE, linewidth=0.6, + label="Background", zorder=1) + + for i, ctau in enumerate(ctau_vals): + s = sig_per_ctau[ctau][var].astype(float) + hs, _ = np.histogram(s, bins=N_BINS, range=(lo, hi)) + if hs.sum() > 0: + hs = hs / hs.sum() + ax.stairs(hs, edges, color=MASS_COLORS[i % len(MASS_COLORS)], linewidth=1.6, + label=rf"$c\tau = {ctau:g}$ mm", zorder=3 + i) + + pos_vals = hb_norm[hb_norm > 0] + if pos_vals.size > 0: + ax.set_yscale('log') + ymin = max(pos_vals.min() * 0.3, 1e-6) + ax.set_ylim(bottom=ymin) + + ax.set_title(var, fontsize=11) + ax.set_xlabel(xlabel) + ax.set_ylabel("a.u.") + + ax.text(0.02, 0.97, "Preliminary", + transform=ax.transAxes, + fontsize=11, fontstyle="italic", fontweight="bold", + va="top", ha="left") + ax.legend(loc="best", framealpha=0.9, + title=rf"$m_\pi = {mpi:g}$ GeV, $m_A = {mA:g}$ GeV", + title_fontsize=10) + ax.tick_params(direction="in", top=True, right=True, which="both") + fig.tight_layout() + fig.savefig(os.path.join(outdir, f"{var}.png"), dpi=130) + plt.close(fig) + + print(" Done.") + +print("\nAll done.") \ No newline at end of file diff --git a/BDT/workingpoint.py b/BDT/workingpoint.py new file mode 100644 index 0000000..2be96a4 --- /dev/null +++ b/BDT/workingpoint.py @@ -0,0 +1,1076 @@ +#!/usr/bin/env python3 +"""Working point analysis for the scouting BDT. + +Trains the global BDT, finds the BDT score threshold at a configurable +background rejection target, then plots input variable distributions for +events passing the working point cut (signal vs background). + +Usage (e.g): + python3 BDT/workingpoint.py --bkg minbias --tuples-dir tuples_L1_info --require-l1 +""" + +import argparse +import gc +import glob +import json +import re +import uproot +import xgboost +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D +import mplhep as hep +import numpy as np +import pandas as pd +import os +from pathlib import Path + +from sklearn.model_selection import train_test_split +from sklearn.metrics import roc_curve, roc_auc_score +from scipy.stats import norm +from xgboost import XGBClassifier + +hep.style.use("CMS") +plt.rcParams.update({ + "font.size": 13, + "axes.labelsize": 13, + "axes.titlesize": 13, + "xtick.labelsize": 11, + "ytick.labelsize": 11, + "legend.fontsize": 9, + "legend.title_fontsize": 10, + "axes.linewidth": 1.0, + "xtick.major.size": 5, + "ytick.major.size": 5, + "xtick.minor.size": 3, + "ytick.minor.size": 3, + "xtick.major.width": 0.9, + "ytick.major.width": 0.9, +}) + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- +_parser = argparse.ArgumentParser(description="Working point analysis for the scouting BDT.") +_parser.add_argument("--model-tag", default="A", help="Signal scenario tag for output naming.") +_parser.add_argument("--conditional", action=argparse.BooleanOptionalAction, default=False, + help="Train a parametric (conditional) BDT with the signal mass point " + "(param_ctau, param_mA, param_mpi) as extra inputs. Background is " + "replicated once per signal point, as in BDT_training.py.") +_parser.add_argument("--table", action=argparse.BooleanOptionalAction, default=True, + help="Build the Asimov significance table: for a set of target FPRs, " + "find the BDT threshold at which the xsec-weighted SV-selected QCD " + "has that FPR, then tabulate p0 = 1 - Phi(Z) (with Z) per signal w" + "point. Writes .tex per (mpi, mA). " + "Default: on; pass --no-table to skip.") +_parser.add_argument("--fpr-targets", type=float, nargs="+", default=[1e-4], + help="Target false-positive rate(s), i.e. background-efficiency working " + "point(s). One BDT-score threshold is computed per target (rows of " + "the --table output; one significance line per target). Accepts a " + "space-separated list, e.g. --fpr-targets 1e-4 1e-3 1e-2 1e-1. " + "Default: 1e-4 (the highest-significance operating point).") +_parser.add_argument("--mass-window-gev", type=float, default=1.0, + help="ABSOLUTE SV1 dimuon mass window half-width in GeV. Used only when " + "--mass-window-rel is 0. Signal and background yields are counted " + "within [mA - W, mA + W] around the nominal signal mass mA. " + "Set both this and --mass-window-rel to 0 to disable (full mass range).") +_parser.add_argument("--mass-window-rel", type=float, default=0.1, + help="RELATIVE (mass-dependent) SV1 dimuon mass window half-width, as a " + "fraction of mA: the window is [mA*(1-r), mA*(1+r)]. Takes precedence " + "over --mass-window-gev when > 0 (default r = 0.1 = +/-10%% of mA). " + "Set to 0 to fall back to the absolute --mass-window-gev window.") +_parser.add_argument("--bkg", choices=["minbias", "qcd", "both"], default="minbias", + help="Which background(s) to run. One shared BDT is trained on the " + "selected background(s). Output files are suffixed _minBias, _QCD, " + "or _both. In 'both' mode the significance-vs-ctau and shape plots " + "overlay MinBias and QCD as separate lines (ROC/discriminant remain " + "single-line, from the shared BDT).") +_parser.add_argument("--tuples-dir", default="tuples_parking_nochi2", + help="Name of the tuple directory under the repo root. Use " + "'tuples_L1_info' for the re-filled (unskimmed, collection-OR) tuples " + "that carry the 'passL1' branch; pair it with --require-l1.") +_parser.add_argument("--require-l1", action=argparse.BooleanOptionalAction, default=False, + help="Keep only events with passL1 != 0 (the L1-seed decision), applied to " + "BOTH signal and background right after loading -- i.e. fold the L1 " + "trigger efficiency into every yield/efficiency. Requires tuples with a " + "'passL1' branch (see --tuples-dir tuples_L1_info). Outputs go to a " + "separate 'significance_plots_L1req/' tree so they don't clobber the " + "no-L1 results.") +_args = _parser.parse_args() + +use_conditional = _args.conditional +model_tag = _args.model_tag +MASS_WINDOW_GEV = _args.mass_window_gev # absolute SV1 mass window half-width [GeV] around mA +MASS_WINDOW_REL = _args.mass_window_rel # relative half-width (fraction of mA); takes precedence +MASS_WINDOW_ACTIVE = bool((MASS_WINDOW_REL and MASS_WINDOW_REL > 0.0) or + (MASS_WINDOW_GEV and MASS_WINDOW_GEV > 0.0)) +REQUIRE_L1 = _args.require_l1 # keep only passL1 != 0 events (sig + bkg) +TUPLES_SUBDIR = _args.tuples_dir # tuple directory name under the repo root + +def _mass_window_halfwidth(mA): + """SV1 dimuon mass window half-width around mA [GeV]. + Relative (mass-dependent) window takes precedence; else absolute; else None.""" + if MASS_WINDOW_REL and MASS_WINDOW_REL > 0.0: + return MASS_WINDOW_REL * mA + if MASS_WINDOW_GEV and MASS_WINDOW_GEV > 0.0: + return MASS_WINDOW_GEV + return None + +def _win_label_str(): + """LaTeX label describing the active mass window, or '' if disabled.""" + if MASS_WINDOW_REL and MASS_WINDOW_REL > 0.0: + return rf'$|m_{{\mu\mu}}-m_A|<{MASS_WINDOW_REL:g}\,m_A$' + if MASS_WINDOW_GEV and MASS_WINDOW_GEV > 0.0: + return rf'$|m_{{\mu\mu}}-m_A|<{MASS_WINDOW_GEV:g}$ GeV' + return '' + +LUMI_FB = 109.95 # 2024 Luminosity [fb^-1] +SIG_XSEC_PB = 0.439 + + +SIG_NGEN = { + (10.0, 1.0, 0.1): 997140, + (10.0, 1.0, 1.0): 999293, + (10.0, 1.0, 10.0): 980749, + (10.0, 1.0, 100.0): 954880, + (4.0, 1.33, 0.1): 987256, + (4.0, 1.33, 1.0): 932645, + (4.0, 1.33, 10.0): 978670, + (4.0, 1.33, 100.0): 951079, + (4.0, 0.40, 0.1): 903440, + (4.0, 0.40, 1.0): 998576, + (4.0, 0.40, 10.0): 952631, + (4.0, 0.40, 100.0): 997861, + (1.0, 0.33, 0.1): 916220, + (1.0, 0.33, 1.0): 955918, + (1.0, 0.33, 10.0): 934864, + (1.0, 0.33, 100.0): 998564, +} + + +BR_A_MUMU = { + (1.0, 0.33): 0.464, # 0.458/0.988 + (4.0, 0.40): 0.440, # 0.436/0.992 + (2.0, 0.67): 0.193, # not in current tuples + (10.0, 1.00): 0.307, # 0.293/0.95343 + (4.0, 1.33): 0.317, # 0.305/0.9623 +} + +# Trigger label shown in the significance-plot header. +TRIGGER_LABEL = "Scouting Asymptotic Significance" + +MINBIAS_FILES = [ + "tuples_MinBias_Fil-DoubleMuOS43_2024_2024.root", +] + +QCD_FILES = [ + "tuples_QCD_Bin-PT-15to20_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-20to30_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-30to50_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-50to80_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-80to120_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-120to170_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-170to300_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-300to470_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-470to600_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-600to800_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-800to1000_Fil-MuEnriched_2024_2024.root", + "tuples_QCD_Bin-PT-1000_Fil-MuEnriched_2024_2024.root", +] + +# Background cross sections [pb] (https://cmsweb.cern.ch/das/request?view=list&limit=50&instance=prod%2Fglobal&input=%2FQCD_*MuEnriched*%2F*Summer24MiniAODv6*%2FMINIAODSIM) +BKG_XSEC = { + "tuples_QCD_Bin-PT-15to20_Fil-MuEnriched_2024_2024.root": 3018000.0, + "tuples_QCD_Bin-PT-20to30_Fil-MuEnriched_2024_2024.root": 2701000.0, + "tuples_QCD_Bin-PT-30to50_Fil-MuEnriched_2024_2024.root": 1461000.0, + "tuples_QCD_Bin-PT-50to80_Fil-MuEnriched_2024_2024.root": 407600.0, + "tuples_QCD_Bin-PT-80to120_Fil-MuEnriched_2024_2024.root": 96070.0, + "tuples_QCD_Bin-PT-120to170_Fil-MuEnriched_2024_2024.root": 23140.0, + "tuples_QCD_Bin-PT-170to300_Fil-MuEnriched_2024_2024.root": 7754.0, + "tuples_QCD_Bin-PT-300to470_Fil-MuEnriched_2024_2024.root": 699.6, + "tuples_QCD_Bin-PT-470to600_Fil-MuEnriched_2024_2024.root": 67.67, + "tuples_QCD_Bin-PT-600to800_Fil-MuEnriched_2024_2024.root": 21.27, + "tuples_QCD_Bin-PT-800to1000_Fil-MuEnriched_2024_2024.root": 3.89, + "tuples_QCD_Bin-PT-1000_Fil-MuEnriched_2024_2024.root": 1.323, + + "tuples_MinBias_Fil-DoubleMuOS43_2024_2024.root": 1.051e7 * (409318867 / 8.31e9), # ~= 5.18e5 pb +} + +# Number of generated events per background sample (from DAS nevents) +BKG_NGEN = { + "tuples_QCD_Bin-PT-15to20_Fil-MuEnriched_2024_2024.root": 125036760, + #"tuples_QCD_Bin-PT-20to30_Fil-MuEnriched_2024_2024.root": 93304586, #ZOMBIE (for now) + "tuples_QCD_Bin-PT-30to50_Fil-MuEnriched_2024_2024.root": 95305920, + "tuples_QCD_Bin-PT-50to80_Fil-MuEnriched_2024_2024.root": 107449521, + "tuples_QCD_Bin-PT-80to120_Fil-MuEnriched_2024_2024.root": 94128199, + "tuples_QCD_Bin-PT-120to170_Fil-MuEnriched_2024_2024.root": 99824346, + "tuples_QCD_Bin-PT-170to300_Fil-MuEnriched_2024_2024.root": 94338762, + "tuples_QCD_Bin-PT-300to470_Fil-MuEnriched_2024_2024.root": 79815908, + "tuples_QCD_Bin-PT-470to600_Fil-MuEnriched_2024_2024.root": 71786916, + "tuples_QCD_Bin-PT-600to800_Fil-MuEnriched_2024_2024.root": 85842835, + "tuples_QCD_Bin-PT-800to1000_Fil-MuEnriched_2024_2024.root": 81930100, + "tuples_QCD_Bin-PT-1000_Fil-MuEnriched_2024_2024.root": 87293168, + "tuples_MinBias_Fil-DoubleMuOS43_2024_2024.root": 409318867, +} + + +_MINBIAS_LOOPER_COV = 114448466 / 409318867 # ~= 0.2796 +BKG_FRACTION = {"tuples_MinBias_Fil-DoubleMuOS43_2024_2024.root":_MINBIAS_LOOPER_COV} + +def _bkg_fraction(fname): + return BKG_FRACTION.get(fname, 1.0) + +def _eff_ngen(fname): + """Effective generated-event denominator for the loaded subsample: frac * N_gen. + Using this everywhere N_gen enters the yield keeps subsampling unbiased.""" + return _bkg_fraction(fname) * BKG_NGEN[fname] + +PB_TO_FB = 1.0e3 # 1 pb = 1000 fb (for the QCD background cross sections) + +_HERE = Path(__file__).resolve().parent +tuples_dir = _HERE.parent / TUPLES_SUBDIR + +# Filename suffix tagging the background set used (appended to every plot/table). +OUT_TAG = "_minBias" + +# --------------------------------------------------------------------------- +# Per-lxy-bin cut-and-count +# --------------------------------------------------------------------------- +CNC_LXY_DIRS = { + (10.0, 1.00): "mpi10_mA1p00", + (4.0, 1.33): "mpi4_mA1p33", +} + +def _parse_cnc_lxy_table(path): + """Parse one cutncount table into {lxy_label -> (sig_eff, bkg_rej)} for the + per-lxy-bin rows (cutncount label style, e.g. '0p0-0p2'). Rows look like: + ' 0p0-0p2 | 6993/18961 0.369 | 973828/5061994 0.808'.""" + out = {} + for line in Path(path).read_text().splitlines(): + if line.count("|") != 2: + continue + left, mid, right = line.split("|") + # Per-lxy (and 'full range') rows carry 'pass/total' fractions; skip the + # cutflow/header rows, which never do. + if "/" not in mid or "/" not in right: + continue + try: + eff = float(mid.split()[-1]) + rej = float(right.split()[-1]) + except (ValueError, IndexError): + continue + out[left.strip()] = (eff, rej) + return out + +_CNC_DIR = _HERE / "cnc_tables" +_CNC_SUBDIR_RE = re.compile(r"^mpi(\w+)_mA(\w+)$") + +def _cnc_bkg_for_tag(tag): + """Map a plot tag (OUT_TAG / SUB_BKGS) to the cutncount --bkg filename token.""" + return {"_minBias": "minbias", "_QCD": "qcd"}.get(tag) + +def _iter_cnc_tables(bkg): + """Yield (mpi, mA, ctau, path) for every cnc_tables/mpi*_mA*/cnc_ctau-*mm_*.txt. + Matches window-tagged names too (cnc_ctau-1p0mm_qcd_mwinRel0p1.txt); if both a + windowed and a full-mass table exist for a point, the last in sorted order wins.""" + if not bkg: + return + for sub in sorted(_CNC_DIR.glob("mpi*_mA*")): + m = _CNC_SUBDIR_RE.match(sub.name) + if not m: + continue + mpi_v = float(m.group(1).replace("p", ".")) + mA_v = float(m.group(2).replace("p", ".")) + for fp in sorted(sub.glob(f"cnc_ctau-*mm_{bkg}*.txt")): + name = fp.name + ctau_v = float(name[len("cnc_ctau-"):name.index("mm_")].replace("p", ".")) + yield mpi_v, mA_v, ctau_v, fp + +def _load_cnc_fullrange_points(bkg): + """{(mpi, mA, ctau): (sig_eff, bkg_rej)} from the 'full range' row of each table.""" + pts = {} + for mpi_v, mA_v, ctau_v, fp in _iter_cnc_tables(bkg): + rows = _parse_cnc_lxy_table(fp) + if "full range" in rows: + pts[(mpi_v, mA_v, ctau_v)] = rows["full range"] + return pts + +def _load_cnc_lxy_points(bkg): + """Build {(mpi, mA): {lxy_label(cutncount) -> {ctau -> (eff, rej)}}} for the + CNC_LXY_DIRS mass points, from the per-lxy-bin rows of the tables.""" + pts = {} + for mpi_v, mA_v, ctau_v, fp in _iter_cnc_tables(bkg): + if (mpi_v, mA_v) not in CNC_LXY_DIRS: + continue + for lbl, (eff, rej) in _parse_cnc_lxy_table(fp).items(): + if lbl == "full range": + continue + pts.setdefault((mpi_v, mA_v), {}).setdefault(lbl, {})[ctau_v] = (eff, rej) + return pts + +# Full-range C&C operating points, auto-loaded per background tag +CNC_POINTS = { + "_minBias": _load_cnc_fullrange_points("minbias"), + "_QCD": _load_cnc_fullrange_points("qcd"), +} + +FIGSIZE = (8.5, 6.5) +MASS_COLORS = ["#d62728", "#ff7f0e", "#2ca02c", "#1f77b4", "#e377c2"] +BKG_FACE = "#7fc7c4" +BKG_EDGE = "#2f5f5d" + +_SIG_RE = re.compile(r"tuples_Signal_ScenarioA_Par_2024_mpi-(\w+)_mA-(\w+)_ctau-(\w+)mm_2024(?:_\w+)?\.root") + +def _p2f(s): + return float(s.replace("p", ".")) + +def _flabel(f): + return f"{f:g}".replace(".", "p") + +sig_file_params = [] +for fpath in sorted(glob.glob(str(tuples_dir / "tuples_Signal_ScenarioA_Par_2024_*.root"))): + m = _SIG_RE.search(os.path.basename(fpath)) + if m: + mpi_s, mA_s, ctau_s = m.groups() + sig_file_params.append((fpath, _p2f(mpi_s), _p2f(mA_s), _p2f(ctau_s))) + +param_grid = [(p[3], p[2], p[1]) for p in sig_file_params] # (ctau, mA, mpi) + + +# --------------------------------------------------------------------------- +# BDT variables — mirrors BDT_training.py +# --------------------------------------------------------------------------- +def _make_bdt_vars(): + sv_stems = [ + "chi2Ndof", "d3d_mumu_SV", "dphi_mumu_SV", "l3d", "lxy", + "prob", "ptmm", "x", "xErr", "y", "yErr", "z", "zErr", + "dr_mumu", "dphi_mumu", "deta_mumu", "deta_mumu_SV", + "sindphi_lxy", "a3d_mumu", + ] + + mu_stems = [ + "dxy", "dxysig", "dxy_lxy", "dz", "dzsig", + "eta", "isGlobal", "isTracker", "isvtx", "maxdr", + "mindr", "muCSCDT", "muChambs", "muHits", "nhitsbeforesv", + "normChi2", "phi", "phiCorr", "pixHits", "pixLayers", + "pt", "stripHits", "trkLayers", "PFIsoAll0p3", "PFRelIsoAll0p3", + ] + vars_ = [] + for sv in ("SV1", "SV2"): + for s in sv_stems: + vars_.append(f"{sv}_{s}") + for mu in ("mu1", "mu2"): + for s in mu_stems: + vars_.append(f"{sv}_{mu}_{s}") + return vars_ + +BDT_VARIABLES = _make_bdt_vars() +_LOAD_BRANCHES = list(dict.fromkeys(BDT_VARIABLES + ["SV1_lxy", "SV1_mass", "SV2_mass", "passL1"])) + +# --------------------------------------------------------------------------- +# Data loading helpers +# --------------------------------------------------------------------------- +def read_flat(path): + with uproot.open(path) as f: + t = f['tuples'] + available = set(t.keys()) + branches = [b for b in _LOAD_BRANCHES if b in available] + df = t.arrays(branches, library='pd') + return df.astype('float32') + +def apply_l1(df, src): + """If --require-l1, keep only events with passL1 != 0 (L1-seed decision), folding the + L1 trigger efficiency into the sample. Errors out if the branch is missing (wrong dir).""" + if not REQUIRE_L1: + return df + if 'passL1' not in df.columns: + raise SystemExit(f"--require-l1 set but '{src}' has no passL1 branch " + f"(use --tuples-dir tuples_L1_info).") + return df[df['passL1'] > 0.5].reset_index(drop=True) + +def compute_sample_weights(df): + y = df['label'].values + w = np.ones(len(df), dtype=float) + bkg_mask = (y == 0) + w[bkg_mask] = df.loc[bkg_mask, 'xsec_weight'].values + + # Class balancing + n_sig = int((y == 1).sum()) + sum_bkg = float(w[bkg_mask].sum()) + if n_sig > 0 and sum_bkg > 0: + w[y == 1] = sum_bkg / n_sig + return w + + +def add_dxy_lxy(df): + for sv in ("SV1", "SV2"): + denom = df[f"{sv}_lxy"] * df[f"{sv}_mass"] / df[f"{sv}_ptmm"] + denom = np.where(denom > 1e-9, denom, np.float32(1e-9)) + for mu in ("mu1", "mu2"): + df[f"{sv}_{mu}_dxy_lxy"] = np.abs(df[f"{sv}_{mu}_dxy"]) / denom + +sig_frames = [] +for fpath, mpi_val, mA_val, ctau_val in sig_file_params: + if not Path(fpath).exists(): + continue + df = apply_l1(read_flat(fpath), Path(fpath).name) + df['param_ctau'] = float(ctau_val) + df['param_mA'] = float(mA_val) + df['param_mpi'] = float(mpi_val) + df['label'] = 1 + sig_frames.append(df) +df_sig = pd.concat(sig_frames, ignore_index=True) + + +if _args.bkg == "minbias": + OUT_TAG = "_minBias" + ACTIVE_FILES = list(MINBIAS_FILES) + SUB_BKGS = [("_minBias", list(MINBIAS_FILES), "-", "MinBias")] +elif _args.bkg == "qcd": + OUT_TAG = "_QCD" + ACTIVE_FILES = list(QCD_FILES) + SUB_BKGS = [("_QCD", list(QCD_FILES), "-", "QCD")] +else: + OUT_TAG = "_both" + ACTIVE_FILES = list(MINBIAS_FILES) + list(QCD_FILES) + SUB_BKGS = [("_minBias", list(MINBIAS_FILES), "-", "MinBias"), + ("_QCD", list(QCD_FILES), "--", "QCD")] + +for BKG_FILES, OUT_TAG in [(ACTIVE_FILES, OUT_TAG)]: + bkg_frames = [] + for fname in BKG_FILES: + fpath = tuples_dir / fname + if not fpath.exists(): + continue + if BKG_XSEC.get(fname) is None or BKG_NGEN.get(fname) is None: + continue + df = apply_l1(read_flat(fpath), fname) + df['label'] = 0 + df['bkg_file'] = fname + df['xsec_weight'] = BKG_XSEC[fname] / _eff_ngen(fname) + bkg_frames.append(df) + + df_bkg = pd.concat(bkg_frames, ignore_index=True) + + add_dxy_lxy(df_sig) + add_dxy_lxy(df_bkg) + + # --------------------------------------------------------------------------- + # Lxy binning (cm) + # --------------------------------------------------------------------------- + lxy_bins = [0.0, 0.2, 1.0, 2.4, 3.1, 7.0, 11.0, 16.0, 70.0] # Match Scouting analysis + lxy_labels = ["0p0to0p2", "0p2to1p0", "1p0to2p4", "2p4to3p1", "3p1to7p0", "7p0to11p0", "11p0to16p0", "16p0to70p0"] + # Human-readable range per label for plot titles/legends, e.g. "[0, 0.2]" (cm). + lxy_pretty = {lbl: rf'[{lxy_bins[i]:g}, {lxy_bins[i+1]:g}]' + for i, lbl in enumerate(lxy_labels)} + #lxy_bins = [0.0, 1.0, 10.0, 100.0] # Match Parking analysis + #lxy_labels = ["0to1", "1to10", "10to100"] + + df_sig['lxy_bin'] = pd.cut(df_sig['SV1_lxy'], bins=lxy_bins, labels=lxy_labels, include_lowest=True) + df_bkg['lxy_bin'] = pd.cut(df_bkg['SV1_lxy'], bins=lxy_bins, labels=lxy_labels, include_lowest=True) + + available = set(df_sig.columns) & set(df_bkg.columns) + input_vars = [v for v in BDT_VARIABLES if v in available] + missing = [v for v in BDT_VARIABLES if v not in available] + + + + cond_vars = ['param_ctau', 'param_mA', 'param_mpi'] if use_conditional else [] + out_dir = _HERE / ('significance_plots_L1req' if REQUIRE_L1 else 'significance_plots') + os.makedirs(out_dir, exist_ok=True) + + # --------------------------------------------------------------------------- + # Train global BDT + # --------------------------------------------------------------------------- + if use_conditional: + # Parametric BDT: replicate the full background once per signal point + _combined = [] + for ctau_val, mA_val, mpi_val in param_grid: + sig_pt = df_sig[ + (df_sig['param_ctau'] == ctau_val) & + (df_sig['param_mA'] == mA_val) & + (df_sig['param_mpi'] == mpi_val) + ].copy() + bkg_cp = df_bkg.copy() + bkg_cp['param_ctau'] = float(ctau_val) + bkg_cp['param_mA'] = float(mA_val) + bkg_cp['param_mpi'] = float(mpi_val) + _combined.append(pd.concat([sig_pt, bkg_cp], ignore_index=True)) + df_global = pd.concat(_combined, ignore_index=True) + else: + df_global = pd.concat([df_sig, df_bkg], ignore_index=True) + + # df_sig/df_bkg are fully folded into df_global now; free them so their rows + # aren't held alongside the copies below (df_global, X_g, the split). + del df_sig, df_bkg + gc.collect() + + X_g = df_global[input_vars + cond_vars] + y_g = df_global['label'] + w_g = compute_sample_weights(df_global) + df_global['weight'] = w_g + + X_train, X_test, y_train, y_test, w_train, w_test = train_test_split(X_g, y_g, w_g, test_size=0.3, random_state=42, stratify=y_g) + + del X_g + gc.collect() + + bdt = XGBClassifier( + n_estimators=100, max_depth=3, learning_rate=0.1, + use_label_encoder=False, eval_metric='logloss', + tree_method='hist', n_jobs=4, + ) + bdt.fit(X_train, y_train, sample_weight=w_train) + + # --------------------------------------------------------------------------- + # Find working point + # --------------------------------------------------------------------------- + y_score = bdt.predict_proba(X_test)[:, 1] + fpr, tpr, thresholds = roc_curve(y_test, y_score, sample_weight=w_test) + auc = roc_auc_score(y_test, y_score, sample_weight=w_test) + + # --------------------------------------------------------------------------- + # Persist the trained model so it can be APPLIED TO DATA later (see BDT/apply_bdt_to_data.py) + # --------------------------------------------------------------------------- + _model_dir = out_dir / "models" + os.makedirs(_model_dir, exist_ok=True) + _l1_tag = "_L1req" if REQUIRE_L1 else "" + _model_stub = f"bdt_global{OUT_TAG}{_l1_tag}" + _model_path = _model_dir / f"{_model_stub}.json" + bdt.save_model(str(_model_path)) + + # Working-point thresholds: BDT score cut giving each target FPR on the + # xsec-weighted background ROC (closest grid point; same rule as the table). + _wp_thresholds = {} + for f_t in sorted(_args.fpr_targets, reverse=True): + _j = int(np.argmin(np.abs(fpr - f_t))) + _wp_thresholds[f"{f_t:g}"] = { + "threshold": float(thresholds[_j]), + "fpr_achieved": float(fpr[_j]), + "tpr_achieved": float(tpr[_j]), + } + + _manifest = { + "model_file": _model_path.name, + "features": list(input_vars + cond_vars), # EXACT training column order + "conditional": bool(use_conditional), + "cond_vars": list(cond_vars), + "bkg": OUT_TAG.lstrip("_"), + "require_l1": bool(REQUIRE_L1), + "tuples_subdir": TUPLES_SUBDIR, + "auc": float(auc), + "mass_window_rel": float(MASS_WINDOW_REL), + "mass_window_gev": float(MASS_WINDOW_GEV), + "wp_thresholds": _wp_thresholds, + "xgboost_version": xgboost.__version__, + } + _manifest_path = _model_dir / f"{_model_stub}_manifest.json" + with open(_manifest_path, "w") as _mf: + json.dump(_manifest, _mf, indent=2) + print(f"[workingpoint] saved BDT model -> {_model_path}") + print(f"[workingpoint] saved BDT manifest -> {_manifest_path}") + + # Single output directory holding the FPR-scanned significance tables/plots. + if True: + out_dir = _HERE / ('significance_plots_L1req' if REQUIRE_L1 else 'significance_plots') + os.makedirs(out_dir, exist_ok=True) + + def _mp_dir(mpi_val, mA_val, lxy_label=None): + """Per-mass-point output directory out_dir/mpi/mA_[/lxy_