Merging full workflow into Tanay's HiggsDNA

Don't merge this yet! Let's do some testing first before we merge to master

This should work out of the box (fingers crossed). This will have the full h->aa->4gamma workflow included in the base class. This should also not affect any corrections that you are using currently. That still needs to be tested. There may be bugs that I haven't caught yet, so I need to do some testing... I also included some nice helper scripts for you in the scripts folder.

higgs_dna/selections/photon_selections.py

@@ -161,7 +161,7 @@ def photon_preselection(
         )
         return photons
     else:
-        #return photons[  --orig
+        # return photons[  --orig
         #    (photons.electronVeto > e_veto)
         #    & (photons.pt > self.min_pt_photon)
         #    & (photons.isScEtaEB | photons.isScEtaEE)
@@ -175,11 +175,85 @@ def photon_preselection(
         #        | (rel_iso < self.max_chad_rel_iso)
         #    )
         #    & (isEB_high_r9 | isEB_low_r9 | isEE_high_r9 | isEE_low_r9)
-        #]
+        # ]
         return photons[
-            (photons.pt > self.min_pt_photon) #to be removed later
+            (photons.pt > self.min_pt_photon)  #to be removed later
             & (photons.mvaID > self.min_mvaid)
-            & (abs(photons.eta)<=2.5)
-            & ((abs(photons.eta)>=1.566) | (abs(photons.eta)<=1.4442))
+            & (abs(photons.eta) <= 2.5)
+            & ((abs(photons.eta) >= 1.566) | (abs(photons.eta) <= 1.4442))
         ]
 
+
+def photon_preselection_h4g(
+    self, photons: awkward.Array, events: awkward.Array, year="2022"
+) -> awkward.Array:
+
+    #run3 = True if "2022" in year or "2023" in year else False
+    run3 = True
+    if not run3:
+        rho = events.fixedGridRhoAll * awkward.ones_like(photons.pt)
+    else:
+        rho = events.Rho.fixedGridRhoAll * awkward.ones_like(photons.pt)
+
+    photons = photons[awkward.argsort(photons.pt, ascending=False, axis=1)]
+
+    pt_cuts = (photons[:,0].pt > self.min_lead_pho_pt) & (photons[:,1].pt > self.min_sublead_pho_pt)
+
+    ###   HLT-mimicking cuts   ###
+    # Tracker fiducial region
+    photon_abs_eta = numpy.abs(photons.eta)
+    eta = (photon_abs_eta < self.max_eta) & ((photon_abs_eta < self.gap_barrel_eta) | (photon_abs_eta > self.gap_endcap_eta))
+
+    # Electron veto
+    #e_veto = photons.electronVeto > self.e_veto  # should be > 0.5
+    e_veto = photons.pixelSeed < self.e_veto  # should be < 0.5 for pixelSeed
+
+    # Classifying HLT mimicking selections
+    # All photons eta & R9
+    eb_low_region = (photons.eta < self.gap_barrel_eta) & (photons.r9 <= self.hlt_r9_eb_low)
+    eb_high_region = (photons.eta < self.gap_barrel_eta) & (photons.r9 > self.hlt_r9_eb_high)
+    ee_region = (photons.eta > self.gap_endcap_eta) & (photons.r9 > self.hlt_r9_ee)
+
+    # H/E
+    hoe_eb_low = photons.hoe < self.hoe_eb_low
+    hoe_eb_high = photons.hoe < self.hoe_eb_high
+    hoe_ee = photons.hoe < self.hoe_ee
+
+    # Sigma_ieie
+    sigma_ieie_eb_low = photons.sieie < self.sigma_ieie_eb_low
+    sigma_ieie_eb_high = photons.sieie < self.sigma_ieie_eb_high
+    sigma_ieie_ee = photons.sieie < self.sigma_ieie_ee
+
+    # R9
+    r9_eb_low = photons.r9 > self.r9_eb_low
+    r9_eb_high = photons.r9 > self.r9_eb_high
+    r9_ee = photons.r9 > self.r9_ee
+
+    # pfPhoIso
+    pfPhoIso_eb_low = (photons.pfPhoIso03 - rho * self.rho_low_eta_corr) < self.pfPhoIso_eb_low
+    pfPhoIso_eb_high = (photons.pfPhoIso03 - rho * self.rho_low_eta_corr) < self.pfPhoIso_eb_high
+    pfPhoIso_ee = (photons.pfPhoIso03 - rho * self.rho_high_eta_corr) < self.pfPhoIso_ee
+
+    # TrackerIso
+    # Check this... maybe can use pfRelIso03_all_quadratic instead?
+    trackerIso_eb_low = photons.trkSumPtHollowConeDR03 < self.trackerIso_eb_low
+    trackerIso_eb_high = photons.trkSumPtHollowConeDR03 < self.trackerIso_eb_high
+    trackerIso_ee = photons.trkSumPtHollowConeDR03 < self.trackerIso_ee
+
+    # Combine HLT cuts
+    hlt_cut = photons.pt < 0  # initialize to all False
+    hlt_eb_low = hlt_cut | (eb_low_region & hoe_eb_low & sigma_ieie_eb_low & r9_eb_low & pfPhoIso_eb_low & trackerIso_eb_low)
+    hlt_eb_high = hlt_cut | (eb_high_region & hoe_eb_high & sigma_ieie_eb_high & r9_eb_high & pfPhoIso_eb_high & trackerIso_eb_high)
+    hlt_ee = hlt_cut | (ee_region & hoe_ee & sigma_ieie_ee & r9_ee & pfPhoIso_ee & trackerIso_ee)
+    hlt_all = hlt_eb_low | hlt_eb_high | hlt_ee
+
+    # Multiple option isolation cuts
+    # pfChargedIso seems to be chargedHadronIso for newer nanoAOD
+    iso_r9 = photons.r9 > self.iso_min_full5x5_r9
+    iso_charged = (photons.chargedHadronIso if not run3 else photons.pfChargedIso) < self.iso_max_chad
+    iso_charged_rel = (((photons.chargedHadronIso if not run3 else photons.pfChargedIso) / photons.pt) < self.iso_max_chad_rel) & (photons.pt > self.iso_chad_rel_pt) & (photons.hoe < self.iso_chad_rel_hoe)
+    iso = iso_r9 | iso_charged | iso_charged_rel
+
+    all_cuts = pt_cuts & eta & e_veto & hlt_all & iso
+
+    return events, photons[all_cuts], all_cuts