add photon supercluster eta calculation

src/egamma_tnp/nanoaod_efficiency.py

@@ -5,7 +5,7 @@
 from coffea.nanoevents import NanoAODSchema
 
 from egamma_tnp._base_tagnprobe import BaseTagNProbe
-from egamma_tnp.utils import custom_delta_r
+from egamma_tnp.utils import custom_delta_r, dask_calculate_photon_SC_eta
 
 
 class ElectronTagNProbeFromNanoAOD(BaseTagNProbe):
@@ -406,6 +406,8 @@ def __repr__(self):
 
     def find_probes(self, events, cut_and_count, vars):
         if self.use_sc_eta:
+            if "superclusterEta" not in events.Photon.fields:
+                events["Photon", "superclusterEta"] = dak.map_partitions(dask_calculate_photon_SC_eta, events.Photon, events.PV)
             events["Photon", "eta_to_use"] = events.Photon.superclusterEta
             if self.egm_nano:
                 events["Electron", "eta_to_use"] = events.Electron.superclusterEta

src/egamma_tnp/utils/__init__.py

@@ -11,7 +11,7 @@
     fill_pt_eta_phi_mll_histograms,
     get_ratio_histogram,
 )
-from egamma_tnp.utils.misc import custom_delta_r, delta_r_SC
+from egamma_tnp.utils.misc import calculate_photon_SC_eta, custom_delta_r, dask_calculate_photon_SC_eta, delta_r_SC
 
 __all__ = (
     "redirect_files",
@@ -25,6 +25,8 @@
     "get_ratio_histogram",
     "delta_r_SC",
     "custom_delta_r",
+    "calculate_photon_SC_eta",
+    "dask_calculate_photon_SC_eta",
 )
 
 

src/egamma_tnp/utils/misc.py

@@ -3,6 +3,7 @@
 import json
 import re
 
+import awkward as ak
 import numba
 import numpy as np
 
@@ -47,6 +48,74 @@ def delta_r_SC(electron, other):
     return delta_r(electron.eta + electron.deltaEtaSC, electron.phi, other.eta, other.phi)
 
 
+def calculate_photon_SC_eta(photons, PV):
+    """Calculate photon supercluster eta, following the implementation from https://github.com/bartokm/GbbMET/blob/026dac6fde5a1d449b2cfcaef037f704e34d2678/analyzer/Analyzer.h#L2487
+    In the current NanoAODv11, there is only the photon eta which is the SC eta corrected by the PV position.
+    The SC eta is needed to correctly apply a number of corrections and systematics.
+    """
+
+    PV_x = PV.x.to_numpy()
+    PV_y = PV.y.to_numpy()
+    PV_z = PV.z.to_numpy()
+
+    mask_barrel = photons.isScEtaEB
+    mask_endcap = photons.isScEtaEE
+
+    tg_theta_over_2 = np.exp(-photons.eta)
+    # avoid dividion by zero
+    tg_theta_over_2 = np.where(tg_theta_over_2 == 1.0, 1 - 1e-10, tg_theta_over_2)
+    tg_theta = 2 * tg_theta_over_2 / (1 - tg_theta_over_2 * tg_theta_over_2)  # tg(a+b) = tg(a)+tg(b) / (1-tg(a)*tg(b))
+
+    # calculations for EB
+    R = 130.0
+    angle_x0_y0 = np.zeros_like(PV_x)
+
+    angle_x0_y0[PV_x > 0] = np.arctan(PV_y[PV_x > 0] / PV_x[PV_x > 0])
+    angle_x0_y0[PV_x < 0] = np.pi + np.arctan(PV_y[PV_x < 0] / PV_x[PV_x < 0])
+    angle_x0_y0[((PV_x == 0) & (PV_y >= 0))] = np.pi / 2
+    angle_x0_y0[((PV_x == 0) & (PV_y < 0))] = -np.pi / 2
+
+    alpha = angle_x0_y0 + (np.pi - photons.phi)
+    sin_beta = np.sqrt(PV_x**2 + PV_y**2) / R * np.sin(alpha)
+    beta = np.abs(np.arcsin(sin_beta))
+    gamma = np.pi / 2 - alpha - beta
+    length = np.sqrt(R**2 + PV_x**2 + PV_y**2 - 2 * R * np.sqrt(PV_x**2 + PV_y**2) * np.cos(gamma))
+    z0_zSC = length / tg_theta
+
+    tg_sctheta = np.copy(tg_theta)
+    # correct values for EB
+    tg_sctheta = ak.where(mask_barrel, R / (PV_z + z0_zSC), tg_sctheta)
+
+    # calculations for EE
+    intersection_z = np.where(photons.eta > 0, 310.0, -310.0)
+    base = intersection_z - PV_z
+    r = base * tg_theta
+    crystalX = PV_x + r * np.cos(photons.phi)
+    crystalY = PV_y + r * np.sin(photons.phi)
+    # correct values for EE
+    tg_sctheta = ak.where(mask_endcap, np.sqrt(crystalX**2 + crystalY**2) / intersection_z, tg_sctheta)
+
+    sctheta = np.arctan(tg_sctheta)
+    sctheta = ak.where(sctheta < 0, np.pi + sctheta, sctheta)
+    ScEta = -np.log(np.tan(sctheta / 2))
+
+    return ScEta
+
+
+def dask_calculate_photon_SC_eta(photons, PV):
+    """Wrapper for calculate_photon_SC_eta to be used with map_partitions"""
+    ak.typetracer.length_zero_if_typetracer(photons.eta)
+    ak.typetracer.length_zero_if_typetracer(photons.phi)
+    ak.typetracer.length_zero_if_typetracer(photons.isScEtaEB)
+    ak.typetracer.length_zero_if_typetracer(photons.isScEtaEE)
+    ak.typetracer.length_zero_if_typetracer(PV.x)
+    ak.typetracer.length_zero_if_typetracer(PV.y)
+    ak.typetracer.length_zero_if_typetracer(PV.z)
+    if ak.backend(photons, PV) == "typetracer":
+        return ak.Array(ak.Array([[0.0]]).layout.to_typetracer(forget_length=True))
+    return calculate_photon_SC_eta(photons, PV)
+
+
 def merge_goldenjsons(files, outfile):
     """Merge multiple golden jsons into one.