Merge UCSD efforts into main

preselection/bdt/data.py

@@ -0,0 +1,50 @@
+import ROOT as r
+import awkward as ak
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+r.EnableImplicitMT(96)
+
+r.gInterpreter.Declare("""
+auto pt_m_jj = [](const ROOT::RVec<float>& jet1_pt, const ROOT::RVec<float>& jet1_eta, const ROOT::RVec<float>& jet1_phi, const ROOT::RVec<float>& jet1_mass, const ROOT::RVec<float>& jet2_pt, const ROOT::RVec<float>& jet2_eta, const ROOT::RVec<float>& jet2_phi, const ROOT::RVec<float>& jet2_mass) {
+        ROOT::VecOps::RVec<float> pt_jj;
+        ROOT::VecOps::RVec<float> m_jj;
+        for (size_t i = 0; i < jet1_pt.size(); ++i) {
+            auto v_jj = ROOT::Math::PtEtaPhiMVector(jet1_pt[i], jet1_eta[i], jet1_phi[i], jet1_mass[i]) + ROOT::Math::PtEtaPhiMVector(jet2_pt[i], jet2_eta[i], jet2_phi[i], jet2_mass[i]);
+            pt_jj.push_back(v_jj.Pt());
+            m_jj.push_back(v_jj.M());
+        }
+    return std::make_pair(pt_jj, m_jj);
+};
+""")
+
+def process_signal():
+    df = r.RDataFrame("Events", ["/data/userdata/aaarora/spanet_training/run2.root", "/data/userdata/aaarora/spanet_training/run3.root"])
+    r.RDF.Experimental.AddProgressBar(df)
+
+    df = df.Filter("jet_pt.size() >= 2", "At least 2 jets")
+    df = df.Filter("truth_vbs1_idx >= 0 && truth_vbs2_idx >= 0", "Valid truth VBS jet indices")
+
+    df = df.Define("jet_pair_idx", "ROOT::VecOps::Combinations(jet_pt, 2)") \
+        .Define("jet1_pt", "ROOT::VecOps::Take(jet_pt, jet_pair_idx[0])") \
+        .Define("jet2_pt", "ROOT::VecOps::Take(jet_pt, jet_pair_idx[1])") \
+        .Define("jet1_eta", "ROOT::VecOps::Take(jet_eta, jet_pair_idx[0])") \
+        .Define("jet2_eta", "ROOT::VecOps::Take(jet_eta, jet_pair_idx[1])") \
+        .Define("jet1_phi", "ROOT::VecOps::Take(jet_phi, jet_pair_idx[0])") \
+        .Define("jet2_phi", "ROOT::VecOps::Take(jet_phi, jet_pair_idx[1])") \
+        .Define("jet1_mass", "ROOT::VecOps::Take(jet_mass, jet_pair_idx[0])") \
+        .Define("jet2_mass", "ROOT::VecOps::Take(jet_mass, jet_pair_idx[1])") \
+        .Define("pt_m_jj", "pt_m_jj(jet1_pt, jet1_eta, jet1_phi, jet1_mass, jet2_pt, jet2_eta, jet2_phi, jet2_mass)") \
+        .Define("pt_jj", "pt_m_jj.first") \
+        .Define("m_jj", "pt_m_jj.second") \
+        .Define("deta_jj", "abs(jet1_eta - jet2_eta)") \
+        .Define("dphi_jj", "ROOT::VecOps::DeltaPhi(jet1_phi, jet2_phi)") \
+        .Define("labels", "(jet_pair_idx[0] == truth_vbs1_idx && jet_pair_idx[1] == truth_vbs2_idx || jet_pair_idx[1] == truth_vbs1_idx && jet_pair_idx[0] == truth_vbs2_idx)")
+
+    return df
+
+if __name__ == "__main__":
+    output_cols = ["jet1_pt", "jet2_pt", "jet1_eta", "jet2_eta", "jet1_phi", "jet2_phi", 
+                "jet1_mass", "jet2_mass", "pt_jj", "deta_jj", "dphi_jj", "m_jj", "labels"]
+
+    df_sig = ak.to_parquet(ak.from_rdataframe(process_signal(), output_cols), "sig.parquet")

preselection/bdt/train.py

@@ -0,0 +1,96 @@
+import numpy as np
+
+import xgboost as xgb
+
+from sklearn.metrics import roc_curve, auc
+from sklearn.model_selection import train_test_split
+
+import matplotlib.pyplot as plt
+
+import ROOT as r
+
+import awkward as ak
+
+data = ak.from_parquet("../data/sig.parquet")
+
+X = np.column_stack([ak.flatten(data[col]).to_numpy() for col in data.fields[:12]])
+y = ak.flatten(data["labels"]).to_numpy()
+
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+params = {
+    "objective": "binary:logistic",
+    "device": "cuda:1",
+    "tree_method": "hist",
+    "eval_metric": "auc",
+    "scale_pos_weight": sum(y_train == 0) / sum(y_train == 1),
+    "n_estimators": 500,
+    "early_stopping_rounds": 100,
+    "learning_rate": 0.03,
+    "max_depth": 4,
+    "min_child_weight": 6,
+    "gamma": 0.5,
+    "subsample": 0.85,
+    "colsample_bytree": 0.9,
+    "colsample_bylevel": 0.9,
+    "max_delta_step": 0,
+    "lambda": 5.0,
+    "alpha": 0.5,
+}
+
+bdt = xgb.XGBClassifier(**params)
+
+bdt.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=True)
+
+# move bdt to cpu
+bdt.set_params(device="cpu")
+
+y_pred_test = bdt.predict_proba(X_test)[:, 1]
+y_pred_train = bdt.predict_proba(X_train)[:, 1]
+
+fig, ax = plt.subplots(1,2)
+fig.set_size_inches(12, 6)
+
+ax[0].hist(y_pred_test[y_test == 0], bins=20, histtype="step", color="b", label="Non VBS Jets", density=True)
+ax[0].hist(y_pred_test[y_test == 1], bins=20, histtype="step", color="r", label="VBS Jets", density=True)
+ax[0].set_title("Test Set")
+ax[0].legend()
+
+ax[1].hist(y_pred_train[y_train == 0], bins=20, histtype="step", color="b", label="Non VBS Jets", density=True)
+ax[1].hist(y_pred_train[y_train == 1], bins=20, histtype="step", color="r", label="VBS Jets", density=True)
+ax[1].set_title("Train Set")
+ax[1].legend()
+
+plt.savefig("bdt_output.png", dpi=300)
+
+fpr_test, tpr_test, _ = roc_curve(y_test, y_pred_test)
+roc_auc_test = auc(fpr_test, tpr_test)
+
+fpr_train, tpr_train, _ = roc_curve(y_train, y_pred_train)
+roc_auc_train = auc(fpr_train, tpr_train)
+
+fig, ax = plt.subplots()
+ax.plot(fpr_test, tpr_test, label='Test (area = %0.2f)' % roc_auc_test)
+ax.plot(fpr_train, tpr_train, label='Train (area = %0.2f)' % roc_auc_train)
+
+ax.plot([0, 1], [0, 1], 'k--')
+ax.set_xlim([0.0, 1.0])
+ax.set_ylim([0.0, 1.05])
+ax.set_xlabel('Background Efficiency')
+ax.set_ylabel('Signal Efficiency')
+ax.legend()
+
+plt.savefig("roc_curve.png", dpi=300)
+
+features = data.fields
+features = [k for k in features if k not in ["weight", "labels"]]
+
+fig, ax = plt.subplots(figsize=(8, 6))
+xgb.plot_importance(bdt, ax=ax, height=0.5, importance_type="gain", show_values=False)
+ax.set_yticklabels(features) 
+
+r.TMVA.Experimental.SaveXGBoost(bdt, "VBS BDT", "BDT_Weights.root", num_inputs=X.shape[1])
+
+
+
+

preselection/condor/submit.py

@@ -193,6 +193,8 @@ def parse_args():
                         help="Store JES/JER variation branches (default: nominal only)")
     parser.add_argument("--no_jetveto", "--no-jetveto", dest="no_jetveto", action="store_true",
                         help="DEBUG ONLY: compute Jet_vetoMap but do not apply it")
+    parser.add_argument("--cutflow", action="store_true",
+                        help="Generate cutflow histograms")
     return parser.parse_args()
 
 
@@ -390,7 +392,7 @@ def create_tarball(preselection_dir: Path) -> Path:
     preselection_items = [
         "Makefile", "src", "include", "corrections", "applybtag.yaml",
         "etc/goldenJson", spanet_run2_dir, spanet_run3_dir,
-        bdt_dir
+        bdt_dir, "data"
     ]
 
     # Build tar command
@@ -428,6 +430,8 @@ def generate_submit_file(task_dir: Path, job_dir: Path, job_name: str,
         extra_flags += " --store_hlt"
     if args.skip_btag_sf:
         extra_flags += " --skip-btag-sf"
+    if args.cutflow:
+        extra_flags += " --cutflow"
 
     # Arguments passed to executable:
     # USER N_CPUS CONFIG_FILE OUTPUT_NAME ANALYSIS RUN_NUMBER SAMPLE_NAME JOB_IDX [EXTRA_FLAGS]