add ability to get all variables

ikrommyd · Jun 17, 2024 · 830bfcb · 830bfcb
1 parent 1fcde43
commit 830bfcb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/src/egamma_tnp/nanoaod_efficiency.py b/src/egamma_tnp/nanoaod_efficiency.py
@@ -219,6 +219,11 @@ def find_probes(self, events, cut_and_count, mass_range, vars):
             is_photon_filter=self.is_photon_filter,
         )
 
+        if vars == "all":
+            vars_tags = [f"tag_Ele_{var}" for var in all_probe_events.tag_Ele.fields]
+            vars_probes = [f"el_{var}" for var in all_probe_events.el.fields]
+            vars = vars_tags + vars_probes
+
         probe_dict = {}
         for var in vars:
             if var.startswith("el_"):
@@ -555,6 +560,11 @@ def find_probes(self, events, cut_and_count, mass_range, vars):
             start_from_diphotons=self.start_from_diphotons,
         )
 
+        if vars == "all":
+            vars_tags = [f"tag_Ele_{var}" for var in all_probe_events.tag_Ele.fields]
+            vars_probes = [f"ph_{var}" for var in all_probe_events.ph.fields]
+            vars = vars_tags + vars_probes
+
         probe_dict = {}
         for var in vars:
             if var.startswith("ph_"):

diff --git a/src/egamma_tnp/ntuple_efficiency.py b/src/egamma_tnp/ntuple_efficiency.py
@@ -163,6 +163,11 @@ def find_probes(self, events, cut_and_count, mass_range, vars):
 
         passing_locs, all_probe_events = self._find_passing_events(events, cut_and_count=cut_and_count, mass_range=mass_range)
 
+        if vars == "all":
+            vars_tags = [v for v in events.fields if v.startswith("tag_Ele_")]
+            vars_probes = [v for v in events.fields if v.startswith("el_")]
+            vars = vars_tags + vars_probes
+
         if cut_and_count:
             probes = dak.zip({var: all_probe_events[var] for var in vars} | passing_locs)
         else:
@@ -324,6 +329,11 @@ def find_probes(self, events, cut_and_count, mass_range, vars):
 
         passing_locs, all_probe_events = self._find_passing_events(events, cut_and_count=cut_and_count, mass_range=mass_range)
 
+        if vars == "all":
+            vars_tags = [v for v in events.fields if v.startswith("tag_Ele_")]
+            vars_probes = [v for v in events.fields if v.startswith("ph_")]
+            vars = vars_tags + vars_probes
+
         if cut_and_count:
             probes = dak.zip({var: all_probe_events[var] for var in vars} | passing_locs)
         else: