diff --git a/higgs_dna/metaconditions/Era2017_legacy_v1.json b/higgs_dna/metaconditions/Era2017_legacy_v1.json
index 2b4a3d72907a8a61a58e0e1b6795648ff65fa62d..e24ef337e42acdc3091e9f8de73a59990f617bfe 100644
--- a/higgs_dna/metaconditions/Era2017_legacy_v1.json
+++ b/higgs_dna/metaconditions/Era2017_legacy_v1.json
@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
-oid sha256:253243b86ce5aa3a6229e7b5a0fb97281f4d92adbedee2bf4174a0ad0a978c95
-size 17606
+oid sha256:d94b99c1a5b7ae1e404c935931d100fbc72a56c944a850de8b478374cab76355
+size 17757
diff --git a/higgs_dna/metaconditions/Era2022_v1.json b/higgs_dna/metaconditions/Era2022_v1.json
index d8dbee1f86ce259e934cd6f517409a254a371c40..b6d688d50633fb74d98b8ab37169fbeba7d54d34 100755
--- a/higgs_dna/metaconditions/Era2022_v1.json
+++ b/higgs_dna/metaconditions/Era2022_v1.json
@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
-oid sha256:b40b09b78f538b159bdc85ac29c359108bf96dc5990345d7d361819d401dd619
-size 17127
+oid sha256:5a67cb173e47b7f234c1a47e2819ebc0358f8bf8fc4b635e325ecb1825fe0b27
+size 17237
diff --git a/higgs_dna/selections/photon_selections.py b/higgs_dna/selections/photon_selections.py
index c1a3a2f833a13e1cdcc27651fec3fad8a6d754e7..be148e6ecb9048aa1778fc4ad4a773d9369c2cc6 100644
--- a/higgs_dna/selections/photon_selections.py
+++ b/higgs_dna/selections/photon_selections.py
@@ -11,6 +11,7 @@ def photon_preselection(
events: awkward.Array,
apply_electron_veto=True,
year="2023",
+ IsFlag=False
) -> awkward.Array:
"""
Apply preselection cuts to photons.
@@ -141,18 +142,37 @@ def photon_preselection(
)
# not apply electron veto for for TnP workflow
e_veto = self.e_veto if apply_electron_veto else -1
- return photons[
- (photons.electronVeto > e_veto)
- & (photons.pt > self.min_pt_photon)
- & (photons.isScEtaEB | photons.isScEtaEE)
- & (photons.mvaID > self.min_mvaid)
- & (photons.hoe < self.max_hovere)
- & (
- (photons.r9 > self.min_full5x5_r9)
- | (
- rel_iso * photons.pt < self.max_chad_iso
+
+ if IsFlag:
+ photons["PassPresel"] = (
+ (photons.electronVeto > e_veto)
+ & (photons.pt > self.min_pt_photon)
+ & (photons.isScEtaEB | photons.isScEtaEE)
+ & (photons.mvaID > self.min_mvaid)
+ & (photons.hoe < self.max_hovere)
+ & (
+ (photons.r9 > self.min_full5x5_r9)
+ | (
+ rel_iso * photons.pt < self.max_chad_iso
+ )
+ | (rel_iso < self.max_chad_rel_iso)
)
- | (rel_iso < self.max_chad_rel_iso)
+ & (isEB_high_r9 | isEB_low_r9 | isEE_high_r9 | isEE_low_r9)
)
- & (isEB_high_r9 | isEB_low_r9 | isEE_high_r9 | isEE_low_r9)
- ]
+ return photons
+ else:
+ return photons[
+ (photons.electronVeto > e_veto)
+ & (photons.pt > self.min_pt_photon)
+ & (photons.isScEtaEB | photons.isScEtaEE)
+ & (photons.mvaID > self.min_mvaid)
+ & (photons.hoe < self.max_hovere)
+ & (
+ (photons.r9 > self.min_full5x5_r9)
+ | (
+ rel_iso * photons.pt < self.max_chad_iso
+ )
+ | (rel_iso < self.max_chad_rel_iso)
+ )
+ & (isEB_high_r9 | isEB_low_r9 | isEE_high_r9 | isEE_low_r9)
+ ]
diff --git a/higgs_dna/workflows/__init__.py b/higgs_dna/workflows/__init__.py
index 86db257c5b56f6c7615cda4378f6f435d0e18207..11aa4aa0ae72f78560c4bcd6de77e944f78e526f 100644
--- a/higgs_dna/workflows/__init__.py
+++ b/higgs_dna/workflows/__init__.py
@@ -8,6 +8,7 @@ from higgs_dna.workflows.particleLevel import ParticleLevelProcessor
from higgs_dna.workflows.top import TopProcessor
from higgs_dna.workflows.Zmmy import ZmmyProcessor, ZmmyHist, ZmmyZptHist
from higgs_dna.workflows.hpc_processor import HplusCharmProcessor
+from higgs_dna.workflows.zee_processor import ZeeProcessor
from higgs_dna.workflows.lowmass import lowmassProcessor
workflows = {}
@@ -21,6 +22,7 @@ workflows["zmmy"] = ZmmyProcessor
workflows["zmmyHist"] = ZmmyHist
workflows["zmmyZptHist"] = ZmmyZptHist
workflows["hpc"] = HplusCharmProcessor
+workflows["zee"] = ZeeProcessor
workflows["lowmass"] = lowmassProcessor
__all__ = ["workflows", "taggers", "DYStudiesProcessor"]
diff --git a/higgs_dna/workflows/zee_processor.py b/higgs_dna/workflows/zee_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e489bd9cdfadafd1bb2156ef1b2e3a1981ccba2
--- /dev/null
+++ b/higgs_dna/workflows/zee_processor.py
@@ -0,0 +1,571 @@
+from higgs_dna.workflows.base import HggBaseProcessor
+from higgs_dna.systematics import object_systematics as available_object_systematics
+from higgs_dna.systematics import object_corrections as available_object_corrections
+from higgs_dna.systematics import weight_systematics as available_weight_systematics
+from higgs_dna.systematics import weight_corrections as available_weight_corrections
+from higgs_dna.selections.photon_selections import photon_preselection
+from higgs_dna.selections.lumi_selections import select_lumis
+from higgs_dna.tools.SC_eta import add_photon_SC_eta
+from higgs_dna.tools.flow_corrections import calculate_flow_corrections
+from higgs_dna.tools.EcalBadCalibCrystal_events import remove_EcalBadCalibCrystal_events
+from higgs_dna.tools.sigma_m_tools import compute_sigma_m
+from typing import Any, Dict, List, Optional
+import awkward
+import logging
+import functools
+import warnings
+import numpy
+import sys
+from coffea.analysis_tools import Weights
+from higgs_dna.selections.lepton_selections import select_electrons, select_muons
+from higgs_dna.selections.jet_selections import select_jets, jetvetomap
+from higgs_dna.tools.EELeak_region import veto_EEleak_flag
+from copy import deepcopy
+from higgs_dna.utils.misc_utils import choose_jet
+
+from higgs_dna.utils.dumping_utils import (
+ diphoton_ak_array,
+ dump_ak_array,
+ diphoton_list_to_pandas,
+ dump_pandas,
+ get_obj_syst_dict,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ZeeProcessor(HggBaseProcessor):
+ def __init__(
+ self,
+ metaconditions: Dict[str, Any],
+ systematics: Dict[str, List[Any]] = None,
+ corrections: Optional[Dict[str, List[str]]] = None,
+ apply_trigger: bool = False,
+ output_location: Optional[str] = None,
+ taggers: Optional[List[Any]] = None,
+ trigger_group: str = ".*DoubleEG.*",
+ analysis: str = "mainAnalysis",
+ skipCQR: bool = False,
+ skipJetVetoMap: bool = False,
+ year: Optional[Dict[str, List[str]]] = None,
+ fiducialCuts: str = "classical",
+ doDeco: bool = False,
+ Smear_sigma_m: bool = False,
+ doFlow_corrections: bool = False,
+ output_format: str = "parquet"
+ ) -> None:
+ super().__init__(
+ metaconditions,
+ systematics=systematics,
+ corrections=corrections,
+ apply_trigger=apply_trigger,
+ output_location=output_location,
+ taggers=taggers,
+ trigger_group=".*DoubleEG.*",
+ analysis="Dielectron",
+ skipCQR=skipCQR,
+ skipJetVetoMap=skipJetVetoMap,
+ year=year if year is not None else {},
+ fiducialCuts=fiducialCuts,
+ doDeco=doDeco,
+ Smear_sigma_m=Smear_sigma_m,
+ doFlow_corrections=doFlow_corrections,
+ output_format=output_format
+ )
+
+ # diphoton preselection cuts - Based on Dielectron trigger
+ self.min_pt_photon = 23.0
+ self.min_pt_lead_photon = 12.0
+
+ def postprocess(self, accumulant: Dict[Any, Any]) -> Any:
+ pass
+
+ def process(self, events: awkward.Array) -> Dict[Any, Any]:
+ dataset_name = events.metadata["dataset"]
+
+ # data or monte carlo?
+ self.data_kind = "mc" if hasattr(events, "GenPart") else "data"
+
+ # here we start recording possible coffea accumulators
+ # most likely histograms, could be counters, arrays, ...
+ histos_etc = {}
+ histos_etc[dataset_name] = {}
+ if self.data_kind == "mc":
+ histos_etc[dataset_name]["nTot"] = int(
+ awkward.num(events.genWeight, axis=0)
+ )
+ histos_etc[dataset_name]["nPos"] = int(awkward.sum(events.genWeight > 0))
+ histos_etc[dataset_name]["nNeg"] = int(awkward.sum(events.genWeight < 0))
+ histos_etc[dataset_name]["nEff"] = int(
+ histos_etc[dataset_name]["nPos"] - histos_etc[dataset_name]["nNeg"]
+ )
+ histos_etc[dataset_name]["genWeightSum"] = float(
+ awkward.sum(events.genWeight)
+ )
+ else:
+ histos_etc[dataset_name]["nTot"] = int(len(events))
+ histos_etc[dataset_name]["nPos"] = int(histos_etc[dataset_name]["nTot"])
+ histos_etc[dataset_name]["nNeg"] = int(0)
+ histos_etc[dataset_name]["nEff"] = int(histos_etc[dataset_name]["nTot"])
+ histos_etc[dataset_name]["genWeightSum"] = float(len(events))
+
+ # lumi mask
+ if self.data_kind == "data":
+ try:
+ lumimask = select_lumis(self.year[dataset_name][0], events, logger)
+ events = events[lumimask]
+ except:
+ logger.info(
+ f"[ lumimask ] Skip now! Unable to find year info of {dataset_name}"
+ )
+ # apply jetvetomap: only retain events that without any jets in the EE leakage region
+ if not self.skipJetVetoMap:
+ events = jetvetomap(
+ events, logger, dataset_name, year=self.year[dataset_name][0]
+ )
+ # metadata array to append to higgsdna output
+ metadata = {}
+
+ if self.data_kind == "mc":
+ # Add sum of gen weights before selection for normalisation in postprocessing
+ metadata["sum_genw_presel"] = str(awkward.sum(events.genWeight))
+ else:
+ metadata["sum_genw_presel"] = "Data"
+
+ # apply filters and triggers
+ events = self.apply_filters_and_triggers(events)
+
+ # remove events affected by EcalBadCalibCrystal
+ if self.data_kind == "data":
+ events = remove_EcalBadCalibCrystal_events(events)
+
+ # Matching photons eta and phi to electrons
+ events["Photon"] = events.Photon[events.Photon.electronIdx > -1]
+ events = events[awkward.num(events.Photon) >= 2]
+
+ photons = events.Photon
+ matched_electrons = events.Electron[photons.electronIdx]
+
+ # Keeping the photon eta and phi
+ photons["phoeta"] = photons.eta
+ photons["phophi"] = photons.phi
+
+ # Substituting the photon eta and phi with the matched electron eta and phi
+ photons["eta"] = matched_electrons.eta
+ photons["phi"] = matched_electrons.phi
+
+ photons["ele_pt"] = matched_electrons.pt
+ photons["ele_energy"] = matched_electrons.energy
+ photons["ele_ecalEnergy"] = matched_electrons.ecalEnergy
+ photons["ele_ecalEnergyError"] = matched_electrons.ecalEnergyError
+ photons["ele_ScEta"] = matched_electrons.eta + matched_electrons.deltaEtaSC
+
+ events.Photon = photons
+
+ # we need ScEta for corrections and systematics, it is present in NanoAODv13+ and can be calculated using PV for older versions
+ events.Photon = add_photon_SC_eta(events.Photon, events.PV)
+
+ # add veto EE leak branch for photons, could also be used for electrons
+ if (
+ self.year[dataset_name][0] == "2022EE"
+ or self.year[dataset_name][0] == "2022postEE"
+ ):
+ events.Photon = veto_EEleak_flag(self, events.Photon)
+
+ # read which systematics and corrections to process
+ try:
+ correction_names = self.corrections[dataset_name]
+ except KeyError:
+ correction_names = []
+ try:
+ systematic_names = self.systematics[dataset_name]
+ except KeyError:
+ systematic_names = []
+
+ # If --Smear_sigma_m == True and no Smearing correction in .json for MC throws an error, since the pt scpectrum need to be smeared in order to properly calculate the smeared sigma_m_m
+ if (
+ self.data_kind == "mc"
+ and self.Smear_sigma_m
+ and not ("Smearing" in correction_names)
+ ):
+ warnings.warn(
+ "Smearing should be specified in the corrections field in .json in order to smear the mass!"
+ )
+ sys.exit(0)
+
+ # Since now we are applying Smearing term to the sigma_m_over_m i added this portion of code
+ # specially for the estimation of smearing terms for the data events [data pt/energy] are not smeared!
+ if self.data_kind == "data" and self.Smear_sigma_m:
+ # Since we smear sigma_m also in data, we also need to take the smearing term
+ correction_name = "Smearing"
+
+ logger.info(
+ f"\nApplying correction {correction_name} to dataset {dataset_name}\n"
+ )
+ varying_function = available_object_corrections[correction_name]
+ events = varying_function(events=events, year=self.year[dataset_name][0])
+
+ for correction_name in correction_names:
+ if correction_name in available_object_corrections.keys():
+ logger.info(
+ f"Applying correction {correction_name} to dataset {dataset_name}"
+ )
+ varying_function = available_object_corrections[correction_name]
+ events = varying_function(
+ events=events, year=self.year[dataset_name][0]
+ )
+ elif correction_name in available_weight_corrections:
+ # event weight corrections will be applied after photon preselection / application of further taggers
+ continue
+ else:
+ # may want to throw an error instead, needs to be discussed
+ warnings.warn(f"Could not process correction {correction_name}.")
+ continue
+
+ original_photons = events.Photon
+ # NOTE: jet jerc systematics are added in the correction functions and handled later
+ original_jets = events.Jet
+
+ # Computing the normalizing flow correction
+ if self.data_kind == "mc" and self.doFlow_corrections:
+
+ # Applyting the Flow corrections to all photons before pre-selection
+ counts = awkward.num(original_photons)
+ corrected_inputs,var_list = calculate_flow_corrections(original_photons, events, self.meta["flashggPhotons"]["flow_inputs"], self.meta["flashggPhotons"]["Isolation_transform_order"], year=self.year[dataset_name][0])
+
+ # Store the raw nanoAOD value and update photon ID MVA value for preselection
+ original_photons["mvaID_nano"] = original_photons["mvaID"]
+
+ # Store the raw values of the inputs and update the input values with the corrections since some variables used in the preselection
+ for i in range(len(var_list)):
+ original_photons["raw_" + str(var_list[i])] = original_photons[str(var_list[i])]
+ original_photons[str(var_list[i])] = awkward.unflatten(corrected_inputs[:,i] , counts)
+
+ original_photons["mvaID"] = awkward.unflatten(self.add_photonid_mva_run3(original_photons, events), counts)
+
+ # systematic object variations
+ for systematic_name in systematic_names:
+ if systematic_name in available_object_systematics.keys():
+ systematic_dct = available_object_systematics[systematic_name]
+ if systematic_dct["object"] == "Photon":
+ logger.info(
+ f"Adding systematic {systematic_name} to photons collection of dataset {dataset_name}"
+ )
+ original_photons.add_systematic(
+ # passing the arguments here explicitly since I want to pass the events to the varying function. If there is a more elegant / flexible way, just change it!
+ name=systematic_name,
+ kind=systematic_dct["args"]["kind"],
+ what=systematic_dct["args"]["what"],
+ varying_function=functools.partial(
+ systematic_dct["args"]["varying_function"],
+ events=events,
+ year=self.year[dataset_name][0],
+ )
+ # name=systematic_name, **systematic_dct["args"]
+ )
+ # to be implemented for other objects here
+ elif systematic_name in available_weight_systematics:
+ # event weight systematics will be applied after photon preselection / application of further taggers
+ continue
+ else:
+ # may want to throw an error instead, needs to be discussed
+ warnings.warn(
+ f"Could not process systematic variation {systematic_name}."
+ )
+ continue
+
+ # Applying systematic variations
+ photons_dct = {}
+ photons_dct["nominal"] = original_photons
+ logger.debug(original_photons.systematics.fields)
+ for systematic in original_photons.systematics.fields:
+ for variation in original_photons.systematics[systematic].fields:
+ # deepcopy to allow for independent calculations on photon variables with CQR
+ photons_dct[f"{systematic}_{variation}"] = deepcopy(
+ original_photons.systematics[systematic][variation]
+ )
+
+ # NOTE: jet jerc systematics are added in the corrections, now extract those variations and create the dictionary
+ jerc_syst_list, jets_dct = get_obj_syst_dict(original_jets, ["pt", "mass"])
+ # object systematics dictionary
+ logger.debug(f"[ jerc systematics ] {jerc_syst_list}")
+
+ # Build the flattened array of all possible variations
+ variations_combined = []
+ variations_combined.append(original_photons.systematics.fields)
+ # NOTE: jet jerc systematics are not added with add_systematics
+ variations_combined.append(jerc_syst_list)
+ # Flatten
+ variations_flattened = sum(variations_combined, []) # Begin with empty list and keep concatenating
+ # Attach _down and _up
+ variations = [item + suffix for item in variations_flattened for suffix in ['_down', '_up']]
+ # Add nominal to the list
+ variations.append('nominal')
+ logger.debug(f"[systematics variations] {variations}")
+
+ for variation in variations:
+ photons, jets = photons_dct["nominal"], events.Jet
+ if variation == "nominal":
+ pass # Do nothing since we already get the unvaried, but nominally corrected objets above
+ elif variation in [*photons_dct]: # [*dict] gets the keys of the dict since Python >= 3.5
+ photons = photons_dct[variation]
+ elif variation in [*jets_dct]:
+ jets = jets_dct[variation]
+ do_variation = variation # We can also simplify this a bit but for now it works
+
+ if self.chained_quantile is not None:
+ photons = self.chained_quantile.apply(photons, events)
+ # recompute photonid_mva on the fly
+ if self.photonid_mva_EB and self.photonid_mva_EE:
+ photons = self.add_photonid_mva(photons, events)
+
+ # photon preselection
+ photons = photon_preselection(self, photons, events, year=self.year[dataset_name][0], apply_electron_veto=False, IsFlag=True)
+
+ # sort photons in each event descending in pt
+ # make descending-pt combinations of photons
+ photons = photons[awkward.argsort(photons.pt, ascending=False)]
+ photons["charge"] = awkward.zeros_like(
+ photons.pt
+ ) # added this because charge is not a property of photons in nanoAOD v11. We just assume every photon has charge zero...
+ diphotons = awkward.combinations(
+ photons, 2, fields=["pho_lead", "pho_sublead"]
+ )
+
+ # the remaining cut is to select the leading photons
+ # the previous sort assures the order
+ diphotons = diphotons[
+ diphotons["pho_lead"].pt > self.min_pt_lead_photon
+ ]
+
+ # now turn the diphotons into candidates with four momenta and such
+ diphoton_4mom = diphotons["pho_lead"] + diphotons["pho_sublead"]
+ diphotons["pt"] = diphoton_4mom.pt
+ diphotons["eta"] = diphoton_4mom.eta
+ diphotons["phi"] = diphoton_4mom.phi
+ diphotons["mass"] = diphoton_4mom.mass
+ diphotons["charge"] = diphoton_4mom.charge
+
+ diphoton_pz = diphoton_4mom.z
+ diphoton_e = diphoton_4mom.energy
+
+ diphotons["rapidity"] = 0.5 * numpy.log((diphoton_e + diphoton_pz) / (diphoton_e - diphoton_pz))
+
+ diphotons = awkward.with_name(diphotons, "PtEtaPhiMCandidate")
+
+ # sort diphotons by pT
+ diphotons = diphotons[
+ awkward.argsort(diphotons.pt, ascending=False)
+ ]
+
+ # jet_variables
+ jets = awkward.zip(
+ {
+ "pt": jets.pt,
+ "eta": jets.eta,
+ "phi": jets.phi,
+ "mass": jets.mass,
+ "charge": awkward.zeros_like(
+ jets.pt
+ ), # added this because jet charge is not a property of photons in nanoAOD v11. We just need the charge to build jet collection.
+ "hFlav": jets.hadronFlavour
+ if self.data_kind == "mc"
+ else awkward.zeros_like(jets.pt),
+ "btagDeepFlav_B": jets.btagDeepFlavB,
+ "btagDeepFlav_CvB": jets.btagDeepFlavCvB,
+ "btagDeepFlav_CvL": jets.btagDeepFlavCvL,
+ "btagDeepFlav_QG": jets.btagDeepFlavQG,
+ "jetId": jets.jetId,
+ }
+ )
+ jets = awkward.with_name(jets, "PtEtaPhiMCandidate")
+
+ electrons = awkward.zip(
+ {
+ "pt": events.Electron.pt,
+ "eta": events.Electron.eta,
+ "phi": events.Electron.phi,
+ "mass": events.Electron.mass,
+ "charge": events.Electron.charge,
+ "cutBased": events.Electron.cutBased,
+ "mvaIso_WP90": events.Electron.mvaIso_WP90,
+ "mvaIso_WP80": events.Electron.mvaIso_WP80,
+ }
+ )
+ electrons = awkward.with_name(electrons, "PtEtaPhiMCandidate")
+
+ # Special cut for base workflow to replicate iso cut for electrons also for muons
+ events['Muon'] = events.Muon[events.Muon.pfRelIso03_all < 0.2]
+
+ muons = awkward.zip(
+ {
+ "pt": events.Muon.pt,
+ "eta": events.Muon.eta,
+ "phi": events.Muon.phi,
+ "mass": events.Muon.mass,
+ "charge": events.Muon.charge,
+ "tightId": events.Muon.tightId,
+ "mediumId": events.Muon.mediumId,
+ "looseId": events.Muon.looseId,
+ "isGlobal": events.Muon.isGlobal,
+ }
+ )
+ muons = awkward.with_name(muons, "PtEtaPhiMCandidate")
+
+ # lepton cleaning
+ sel_electrons = electrons[
+ select_electrons(self, electrons, diphotons)
+ ]
+ sel_muons = muons[select_muons(self, muons, diphotons)]
+
+ # jet selection and pt ordering
+ jets = jets[
+ select_jets(self, jets, diphotons, sel_muons, sel_electrons)
+ ]
+ jets = jets[awkward.argsort(jets.pt, ascending=False)]
+
+ # adding selected jets to events to be used in ctagging SF calculation
+ events["sel_jets"] = jets
+ n_jets = awkward.num(jets)
+ Njets2p5 = awkward.num(jets[(jets.pt > 30) & (numpy.abs(jets.eta) < 2.5)])
+
+ first_jet_pt = choose_jet(jets.pt, 0, -999.0)
+ first_jet_eta = choose_jet(jets.eta, 0, -999.0)
+ first_jet_phi = choose_jet(jets.phi, 0, -999.0)
+ first_jet_mass = choose_jet(jets.mass, 0, -999.0)
+ first_jet_charge = choose_jet(jets.charge, 0, -999.0)
+
+ second_jet_pt = choose_jet(jets.pt, 1, -999.0)
+ second_jet_eta = choose_jet(jets.eta, 1, -999.0)
+ second_jet_phi = choose_jet(jets.phi, 1, -999.0)
+ second_jet_mass = choose_jet(jets.mass, 1, -999.0)
+ second_jet_charge = choose_jet(jets.charge, 1, -999.0)
+
+ diphotons["first_jet_pt"] = first_jet_pt
+ diphotons["first_jet_eta"] = first_jet_eta
+ diphotons["first_jet_phi"] = first_jet_phi
+ diphotons["first_jet_mass"] = first_jet_mass
+ diphotons["first_jet_charge"] = first_jet_charge
+
+ diphotons["second_jet_pt"] = second_jet_pt
+ diphotons["second_jet_eta"] = second_jet_eta
+ diphotons["second_jet_phi"] = second_jet_phi
+ diphotons["second_jet_mass"] = second_jet_mass
+ diphotons["second_jet_charge"] = second_jet_charge
+
+ diphotons["n_jets"] = n_jets
+ diphotons["Njets2p5"] = Njets2p5
+
+ diphotons = awkward.firsts(diphotons)
+ # set diphotons as part of the event record
+ events[f"diphotons_{do_variation}"] = diphotons
+ # annotate diphotons with event information
+ diphotons["event"] = events.event
+ diphotons["lumi"] = events.luminosityBlock
+ diphotons["run"] = events.run
+ # nPV just for validation of pileup reweighting
+ diphotons["nPV"] = events.PV.npvs
+ diphotons["fixedGridRhoAll"] = events.Rho.fixedGridRhoAll
+ # annotate diphotons with dZ information (difference between z position of GenVtx and PV) as required by flashggfinalfits
+ if self.data_kind == "mc":
+ diphotons["genWeight"] = events.genWeight
+ diphotons["dZ"] = events.GenVtx.z - events.PV.z
+ # Fill zeros for data because there is no GenVtx for data, obviously
+ else:
+ diphotons["dZ"] = awkward.zeros_like(events.PV.z)
+
+ # drop events without a preselected diphoton candidate
+ # drop events without a tag, if there are tags
+
+ selection_mask = ~awkward.is_none(diphotons)
+ diphotons = diphotons[selection_mask]
+
+ # return if there is no surviving events
+ if len(diphotons) == 0:
+ logger.debug("No surviving events in this run, return now!")
+ return histos_etc
+ if self.data_kind == "mc":
+ # initiate Weight container here, after selection, since event selection cannot easily be applied to weight container afterwards
+ event_weights = Weights(size=len(events[selection_mask]))
+
+ # corrections to event weights:
+ for correction_name in correction_names:
+ if correction_name in available_weight_corrections:
+ logger.info(
+ f"Adding correction {correction_name} to weight collection of dataset {dataset_name}"
+ )
+ varying_function = available_weight_corrections[
+ correction_name
+ ]
+ event_weights = varying_function(
+ events=events[selection_mask],
+ photons=events[f"diphotons_{do_variation}"][
+ selection_mask
+ ],
+ weights=event_weights,
+ dataset_name=dataset_name,
+ year=self.year[dataset_name][0],
+ )
+
+ diphotons["weight_central"] = event_weights.weight()
+ # Store variations with respect to central weight
+ if do_variation == "nominal":
+ if len(event_weights.variations):
+ logger.info(
+ "Adding systematic weight variations to nominal output file."
+ )
+ for modifier in event_weights.variations:
+ diphotons["weight_" + modifier] = event_weights.weight(
+ modifier=modifier
+ )
+
+ # Multiply weight by genWeight for normalisation in post-processing chain
+ event_weights._weight = (
+ events["genWeight"][selection_mask]
+ * diphotons["weight_central"]
+ )
+ diphotons["weight"] = event_weights.weight()
+
+ # Add weight variables (=1) for data for consistent datasets
+ else:
+ diphotons["weight_central"] = awkward.ones_like(
+ diphotons["event"]
+ )
+ diphotons["weight"] = awkward.ones_like(diphotons["event"])
+
+ # Compute and store the different variations of sigma_m_over_m
+ diphotons = compute_sigma_m(diphotons, processor='base', flow_corrections=self.doFlow_corrections, smear=self.Smear_sigma_m)
+
+ if self.output_location is not None:
+ if self.output_format == "root":
+ df = diphoton_list_to_pandas(self, diphotons)
+ else:
+ akarr = diphoton_ak_array(self, diphotons)
+
+ # Remove fixedGridRhoAll from photons to avoid having event-level info per photon
+ akarr = akarr[
+ [
+ field
+ for field in akarr.fields
+ if "lead_fixedGridRhoAll" not in field
+ ]
+ ]
+
+ fname = (
+ events.behavior[
+ "__events_factory__"
+ ]._partition_key.replace("/", "_")
+ + ".%s" % self.output_format
+ )
+ subdirs = []
+ if "dataset" in events.metadata:
+ subdirs.append(events.metadata["dataset"])
+ subdirs.append(do_variation)
+ if self.output_format == "root":
+ dump_pandas(self, df, fname, self.output_location, subdirs)
+ else:
+ dump_ak_array(
+ self, akarr, fname, self.output_location, metadata, subdirs,
+ )
+
+ return histos_etc