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,191 @@ 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 addRhoCorrections(
+    self,
+    photons: awkward.Array,
+    events: awkward.Array,
+    run3: bool
+) -> awkward.Array:
+    """
+    Adds a new field for pfPhoIso03 with the EA rho corrections for Run 3.
+    """
+
+    if not run3:
+        rho = events.fixedGridRhoAll * awkward.ones_like(photons.pt)
+    else:
+        rho = events.Rho.fixedGridRhoAll * awkward.ones_like(photons.pt)
+
+    # Get absolute eta to keep cuts simple
+    photon_abs_eta = numpy.abs(photons.eta)
+
+    # EB regions (rhoCorr_EB_X is the region, mask_EB_X is the mask with only that region)
+    rhoCorr_EB_0 = (photon_abs_eta > 0.0) & (photon_abs_eta < 1.0)
+    mask_EB_0 = awkward.mask(photons.pfPhoIso03, rhoCorr_EB_0)
+    EB_0 = mask_EB_0 - (rho * self.EA1_EB1) - (rho * rho * self.EA2_EB1)
+
+    rhoCorr_EB_1 = (photon_abs_eta > 1.0) & (photon_abs_eta < 1.4442)
+    mask_EB_1 = awkward.mask(photons.pfPhoIso03, rhoCorr_EB_1)
+    EB_1 = mask_EB_1 - (rho * self.EA1_EB2) - (rho * rho * self.EA2_EB2)
+
+    # EE regions (rhoCorr_EE_X is the region, mask_EE_X is the mask with only that region)
+    rhoCorr_EE_0 = (photon_abs_eta > 1.566) & (photon_abs_eta < 2.0)
+    mask_EE_0 = awkward.mask(photons.pfPhoIso03, rhoCorr_EE_0)
+    EE_0 = mask_EE_0 - (rho * self.EA1_EE1) - (rho * rho * self.EA2_EE1)
+
+    rhoCorr_EE_1 = (photon_abs_eta > 2.0) & (photon_abs_eta < 2.2)
+    mask_EE_1 = awkward.mask(photons.pfPhoIso03, rhoCorr_EE_1)
+    EE_1 = mask_EE_1 - (rho * self.EA1_EE2) - (rho * rho * self.EA2_EE2)
+
+    rhoCorr_EE_2 = (photon_abs_eta > 2.2) & (photon_abs_eta < 2.3)
+    mask_EE_2 = awkward.mask(photons.pfPhoIso03, rhoCorr_EE_2)
+    EE_2 = mask_EE_2 - (rho * self.EA1_EE3) - (rho * rho * self.EA2_EE3)
+
+    rhoCorr_EE_3 = (photon_abs_eta > 2.3) & (photon_abs_eta < 2.4)
+    mask_EE_3 = awkward.mask(photons.pfPhoIso03, rhoCorr_EE_3)
+    EE_3 = mask_EE_3 - (rho * self.EA1_EE4) - (rho * rho * self.EA2_EE4)
+
+    rhoCorr_EE_4 = (photon_abs_eta > 2.4) & (photon_abs_eta < 2.5)
+    mask_EE_4 = awkward.mask(photons.pfPhoIso03, rhoCorr_EE_4)
+    EE_4 = mask_EE_4 - (rho * self.EA1_EE5) - (rho * rho * self.EA2_EE5)
+
+    EB_EE_else = awkward.mask(photons.pfPhoIso03, (~rhoCorr_EB_0) & (~rhoCorr_EB_1) & (~rhoCorr_EE_0) & (~rhoCorr_EE_1) & (~rhoCorr_EE_2) & (~rhoCorr_EE_3) & (~rhoCorr_EE_4))
+
+    # Apply corrections to appropriate regions and remove extra Nones to get back original shape
+    #print(EB_0, "\n", EB_1, "\n", EE_0, "\n", EE_1, "\n", EE_2, "\n", EE_3, "\n", EE_4, "\n", EB_EE_else)
+    pfPhoIso03_rhoCorrected = awkward.concatenate([EB_0, EB_1, EE_0, EE_1, EE_2, EE_3, EE_4, EB_EE_else], axis=1)
+    pfPhoIso03_rhoCorrected = pfPhoIso03_rhoCorrected[~awkward.is_none(pfPhoIso03_rhoCorrected, axis=1)]
+
+    assert len(photons.pfPhoIso03) == len(pfPhoIso03_rhoCorrected)
+    assert awkward.all(awkward.num(photons.pfPhoIso03, axis=1) == awkward.num(pfPhoIso03_rhoCorrected, axis=1))
+    photons["pfPhoIso03_rhoCorrected"] = pfPhoIso03_rhoCorrected
+
+    return photons
+
+
+def photon_preselection_h4g(
+    self, photons: awkward.Array, events: awkward.Array, year="2022"
+) -> awkward.Array:
+
+    # Adding in *rho corrected* pfPhoIso03 to make proper plots
+    run3 = True
+    photons = addRhoCorrections(self, photons, events, run3)
+    assert "pfPhoIso03_rhoCorrected" in photons.fields
+
+    ###   HLT-mimicking cuts   ###
+    # Sort photons by pT
+    photons = photons[awkward.argsort(photons.pt, ascending=False, axis=1)]
+
+    # pT cuts
+    photons = photons[awkward.num(photons, axis=1) >= 2]
+    pt_cuts = (photons[:,0].pt > self.min_lead_pho_pt) & (photons[:,1].pt > self.min_sublead_pho_pt)
+    photons = photons[pt_cuts]
+
+    # H/E
+    hoe_eb = photons.hoe < self.hoe_eb
+    hoe_ee = photons.hoe < self.hoe_ee
+
+    # All photons eta & R9
+    eb_region = photons.isScEtaEB & (photons.r9 > self.hlt_r9_eb)
+    ee_region = photons.isScEtaEE & (photons.r9 > self.hlt_r9_ee)
+
+    # Sigma_ieie
+    sigma_ieie_eb = photons.sieie < self.sigma_ieie_eb
+    sigma_ieie_ee = photons.sieie < self.sigma_ieie_ee
+
+    # pfPhoIso
+    # quadratic EA corrections in Run3 : https://indico.cern.ch/event/1204277/contributions/5064356/attachments/2538496/4369369/CutBasedPhotonID_20221031.pdf
+    photon_abs_eta = numpy.abs(photons.eta)
+    pass_phoIso_rho_corr_EB = (
+        ((photon_abs_eta > 0.0) & (photon_abs_eta < 1.0))
+        & (
+            photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_eb
+        )
+    ) | (
+        ((photon_abs_eta > 1.0) & (photon_abs_eta < 1.4442))
+        & (
+            photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_eb
+        )
+    )
+
+    pass_phoIso_rho_corr_EE = (
+        (
+            ((photon_abs_eta > 1.566) & (photon_abs_eta < 2.0))
+            & (
+                photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_ee
+            )
+        )
+        | (
+            ((photon_abs_eta > 2.0) & (photon_abs_eta < 2.2))
+            & (
+                photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_ee
+            )
+        )
+        | (
+            ((photon_abs_eta > 2.2) & (photon_abs_eta < 2.3))
+            & (
+                photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_ee
+            )
+        )
+        | (
+            ((photon_abs_eta > 2.3) & (photon_abs_eta < 2.4))
+            & (
+                photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_ee
+            )
+        )
+        | (
+            ((photon_abs_eta > 2.4) & (photon_abs_eta < 2.5))
+            & (
+                photons.pfPhoIso03_rhoCorrected < self.pfPhoIso_ee
+            )
+        )
+    )
+
+    # TrackerIso
+    trackerIso_eb = photons.trkSumPtHollowConeDR03 < self.trackerIso_eb
+    trackerIso_ee = photons.trkSumPtHollowConeDR03 < self.trackerIso_ee
+
+    # Combine HLT cuts
+    hlt_cut = photons.pt < 0  # initialize to all False
+    hlt_eb = hlt_cut | (eb_region & hoe_eb & sigma_ieie_eb & pass_phoIso_rho_corr_EB & trackerIso_eb)
+    hlt_ee = hlt_cut | (ee_region & hoe_ee & sigma_ieie_ee & pass_phoIso_rho_corr_EE & trackerIso_ee)
+    hlt_all = hlt_eb | hlt_ee
+    """
+    print(hlt_cut, eb_region, hoe_eb, sigma_ieie_eb, pass_phoIso_rho_corr_EB, trackerIso_eb, hlt_all)
+    print(hlt_cut, eb_region, hoe_eb, sigma_ieie_eb, pass_phoIso_rho_corr_EB, trackerIso_eb, hlt_all)
+    print(hlt_cut, eb_region, hoe_eb, sigma_ieie_eb, pass_phoIso_rho_corr_EB, trackerIso_eb, hlt_all)
+    """
+
+    # Tracker fiducial region
+    eta = photons.isScEtaEB | photons.isScEtaEE
+
+    # Electron veto (use pixelSeed in Run 3)
+    #e_veto = photons.electronVeto > self.e_veto  # should be > 0.5 for passing electron veto (i.e., is a photon)
+    e_veto = photons.pixelSeed == False  # should be False to reject any pixelSeed (i.e., is a photon)
+
+    # Multiple option isolation cuts
+    # No iso check cut included since it fully overlaps with previous pT and H/E cuts anyway. Also check all photons to maintain consistent phase space with cuts.
+    #iso_check = (photons.pt > self.iso_chad_rel_pt) & (photons.hoe < self.iso_chad_rel_hoe)
+    iso_r9 = photons.r9 > self.iso_min_full5x5_r9
+    iso_charged = (photons.chargedHadronIso if not run3 else photons.pfRelIso03_chg_quadratic) < self.iso_max_chad
+    iso_charged_rel = (((photons.chargedHadronIso if not run3 else photons.pfRelIso03_chg_quadratic) * photons.pt) < self.iso_max_chad_rel)
+    iso = iso_r9 | iso_charged | iso_charged_rel
+
+    # Apply cuts to photons and events
+    all_cuts = hlt_all & eta & e_veto & iso
+    photons = photons[all_cuts]
+
+    # Remove "None" after cuts are applied
+    photons = photons[~awkward.is_none(photons)]
+    photons = photons[awkward.num(photons, axis=1) >= 2]
+
+    return photons