allow JetResolution to handle additional filename metadata

scikit-hep · Jan 8, 2025 · e8203cb · e8203cb
1 parent 1f69f3a
commit e8203cb
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 4 deletions.
diff --git a/coffea/jetmet_tools/JetResolution.py b/coffea/jetmet_tools/JetResolution.py
@@ -65,15 +65,16 @@ def __init__(self, **kwargs):
                     )
                 )
             info = name.split("_")
-            if len(info) != 5:
+            if len(info) > 6 or len(info) < 5:
                 raise Exception("Corrector name is not properly formatted!")
+            offset = len(info) - 5
 
             campaign = _checkConsistency(campaign, info[0])
             dataera = _checkConsistency(dataera, info[1])
-            datatype = _checkConsistency(datatype, info[2])
-            levels.append(info[3])
+            datatype = _checkConsistency(datatype, info[2 + offset])
+            levels.append(info[3 + offset])
             funcs.append(func)
-            jettype = _checkConsistency(jettype, info[4])
+            jettype = _checkConsistency(jettype, info[4 + offset])
 
         if campaign is None:
             raise Exception("Unable to determine production campaign of JECs!")

diff --git a/coffea/lookup_tools/extractor.py b/coffea/lookup_tools/extractor.py
@@ -26,6 +26,7 @@
         "jec": convert_jec_txt_file,
         "jersf": convert_jersf_txt_file,
         "jr": convert_jr_txt_file,
+        "jer": convert_jr_txt_file,
         "junc": convert_junc_txt_file,
         "ea": convert_effective_area_file,
         "pileup": convert_pileup_json_file,

diff --git a/tests/samples/Summer23Prompt23_RunCv1234_JRV1_MC_PtResolution_AK4PFPuppi.jr.txt.gz b/tests/samples/Summer23Prompt23_RunCv1234_JRV1_MC_PtResolution_AK4PFPuppi.jr.txt.gz
diff --git a/tests/test_jetmet_tools.py b/tests/test_jetmet_tools.py
@@ -27,6 +27,7 @@ def jetmet_evaluator():
             "* * tests/samples/RegroupedV2_Fall17_17Nov2017_V32_MC_UncertaintySources_AK4PFchs.junc.txt.gz",
             "* * tests/samples/Regrouped_Fall17_17Nov2017_V32_MC_UncertaintySources_AK4PFchs.junc.txt",
             "* * tests/samples/Spring16_25nsV10_MC_PtResolution_AK4PFPuppi.jr.txt.gz",
+            "* * tests/samples/Summer23Prompt23_RunCv1234_JRV1_MC_PtResolution_AK4PFPuppi.jr.txt.gz",
             "* * tests/samples/Spring16_25nsV10_MC_SF_AK4PFPuppi.jersf.txt.gz",
             "* * tests/samples/Autumn18_V7_MC_SF_AK4PFchs.jersf.txt.gz",
         ]
@@ -206,6 +207,9 @@ def test_jet_resolution():
     jer_names = ["Spring16_25nsV10_MC_PtResolution_AK4PFPuppi"]
     reso = JetResolution(**{name: evaluator[name] for name in jer_names})
 
+    check2023_names = ["Summer23Prompt23_RunCv1234_JRV1_MC_PtResolution_AK4PFPuppi"]
+    _ = JetResolution(**{name: evaluator[name] for name in check2023_names})
+
     print(reso)
 
     resos = reso.getResolution(JetEta=test_eta, Rho=test_Rho, JetPt=test_pt)