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/workflows/base.py

@@ -277,9 +277,10 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
             warnings.warn(f"Could not instantiate diphoton MVA: {e}")
             self.diphoton_mva = None
 
-    def process_extra(self, events: awkward.Array, photons: awkward.Array, diphotons:awkward.Array, signal: bool, variation: str) -> awkward.Array:
+    def process_extra(self, Nevents: dict, events: awkward.Array, photons: awkward.Array, diphotons:awkward.Array, signal: bool, variation: str) -> awkward.Array:
         # Run selections for pseudoscalars
-        photons, diphotons, cuts = self.produce_and_select_ps(
+        photons, diphotons, cuts, Nevents = self.produce_and_select_ps(
+            Nevents,
             photons=photons,
             diphotons=diphotons,
             signal=signal,
@@ -287,7 +288,7 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
             signal_mass = float(events.metadata["dataset"].replace("2018_","")[7:9]) if signal else None
         )
 
-        return photons, diphotons, cuts, {}
+        return photons, diphotons, cuts, {}, Nevents
 
     def apply_filters_and_triggers(self, events: awkward.Array) -> awkward.Array:
         # met filters
@@ -318,7 +319,8 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
 
         # Need Nevents for selection efficiency checks with signal samples
         # Number of initial events
-        Nevents = {"initial_events": awkward.sum(events.genWeight) if "Signal" in dataset_name else len(events)}
+        #Nevents = {"initial_events": awkward.sum(events.genWeight) if "Signal" in dataset_name else len(events)}
+        Nevents = {"initial_events": len(events)}
 
         # data or monte carlo?
         self.data_kind = "mc" if hasattr(events, "GenPart") else "data"
@@ -380,7 +382,8 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
         events = self.apply_filters_and_triggers(events)
         
         # Filters + triggers efficiency
-        Nevents.update({"filters_triggers": awkward.sum(events.genWeight) if "Signal" in dataset_name else len(events)})
+        #Nevents.update({"filters_triggers": awkward.sum(events.genWeight) if "Signal" in dataset_name else len(events)})
+        Nevents.update({"filters_triggers": len(events)})
 
         # remove events affected by EcalBadCalibCrystal
         if self.data_kind == "data":
@@ -733,15 +736,16 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
             """
 
             # Pre-selections efficiency
-            Nevents.update({"pre_selections": awkward.sum(photons.genWeight) if "Signal" in dataset_name else len(events)})
+            #Nevents.update({"pre_selections": awkward.sum(awkward.firsts(photons.genWeight)) if self.data_kind == "mc" else len(photons)})
+            Nevents.update({"pre_selections": len(photons)})
 
             # H4g-specific processing
-            events = events[(awkward.num(events.Photon) >= 4)]
-            Nevents.update({"pre_selections+4photon": awkward.sum(photons.genWeight) if "Signal" in dataset_name else len(events)})
-            pseudos, diphotons, extra_cuts, process_extra = self.process_extra(events, photons, diphotons, True if "Signal" in dataset_name else False, variation)
+            #Nevents.update({"pre_selections+4photon": awkward.sum(awkward.firsts(photons.genWeight)) if self.data_kind == "mc" else len(photons)})
+            pseudos, diphotons, extra_cuts, process_extra, Nevents = self.process_extra(Nevents, events, photons, diphotons, True if "Signal" in dataset_name else False, variation)
 
             histos_etc.update(process_extra)
-            Nevents.update({"selections": awkward.sum(photons.genWeight) if "Signal" in dataset_name else len(events)})
+            #Nevents.update({"selections": awkward.sum(pseudos.genWeight) if self.data_kind == "mc" else len(pseudos)})
+            Nevents.update({"selections": len(pseudos)})
 
             # return if there is no surviving events (in this case, want at least one pseudoscalar)
             if len(pseudos) == 0:
@@ -894,10 +898,13 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
         return corr_mva
 
 
-    def produce_and_select_ps(self, photons: awkward.Array, diphotons: awkward.Array, signal: bool, variation: str, signal_mass: float = None) -> awkward.Array:
+    def produce_and_select_ps(self, Nevents, photons: awkward.Array, diphotons: awkward.Array, signal: bool, variation: str, signal_mass: float = None) -> awkward.Array:
             # Ensure 4 photons in photons/diphotons arrays
             photons = photons[awkward.num(photons, axis=1) >= 4]
             diphotons = diphotons[awkward.num(photons, axis=1) >= 4]
+            
+            # Add efficiency values to Nevents
+            Nevents.update({"pre_selections+4photon": len(photons)})
 
             # Sort photons by pt
             photons = photons[awkward.argsort(photons.pt, ascending=False, axis=1)]
@@ -919,6 +926,7 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
 
             # Add diphotons to ensure cuts are applied the same way
             pseudos["diphotons"] = diphotons
+            pseudos["genWeight"] = awkward.firsts(photons.genWeight)
 
             # Add sumPt
             pseudos["sumPt_gggg"] = pseudos.pho1.pt + pseudos.pho2.pt + pseudos.pho3.pt + pseudos.pho4.pt
@@ -945,7 +953,11 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
 
             # Combine all cuts
             all_cuts = pt_cuts & eta_cuts & mass_cuts
+            print(all_cuts)
+            print(all_cuts)
             all_cuts = awkward.flatten(all_cuts)
+            print(all_cuts)
+            print(all_cuts)
 
             # Start dM mixing procedure
             # Permutation A - (1,2,3,4)
@@ -1091,4 +1103,4 @@ class HggBaseProcessor(processor.ProcessorABC):  # type: ignore
             pseudos = pseudos[all_cuts]
             pseudos = pseudos[~awkward.is_none(pseudos)]
 
-            return pseudos, pseudos.diphotons, all_cuts
+            return pseudos, pseudos.diphotons, all_cuts, Nevents