feat: add CorrT1METJet to NanoAOD schema and also add zero TrigObj mass (#1480)

* add CorrT1METJet to nanoaod and add zero trigobj mass

* make tracing work with form changes?

* better implementation using awkward internals

* move report def

* remove commented out code

* add key assertion in tracing

* raise warning in case TrigObj mass is present

* alias rawmass -> mass only if present

Author: ikrommyd

src/coffea/nanoevents/schemas/nanoaod.py

@@ -74,6 +74,7 @@ class NanoAODSchema(BaseSchema):
         "LHEPart": "PtEtaPhiMCollection",
         "SubGenJetAK8": "PtEtaPhiMCollection",
         "SubJet": "PtEtaPhiMCollection",
+        "CorrT1METJet": "PtEtaPhiMCollection",
         # Candidate: lorentz + charge
         "Electron": "Electron",
         "LowPtElectron": "LowPtElectron",
@@ -314,6 +315,26 @@ def _build_collections(self, field_names, input_contents):
         if "Photon_energy" in branch_forms:
             branch_forms["Photon_regrEnergy"] = branch_forms.pop("Photon_energy")
 
+        # Alias CorrT1METJet_rawPt to CorrT1METJet_pt and CorrT1METJet_rawMass to CorrT1METJet_mass
+        if "oCorrT1METJet" in branch_forms:
+            if "CorrT1METJet_pt" not in branch_forms:
+                branch_forms["CorrT1METJet_pt"] = branch_forms["CorrT1METJet_rawPt"]
+            if (
+                "CorrT1METJet_mass" not in branch_forms
+                and "CorrT1METJet_rawMass" in branch_forms
+            ):
+                branch_forms["CorrT1METJet_mass"] = branch_forms["CorrT1METJet_rawMass"]
+
+        # Add zero mass to the trigger objects
+        if "oTrigObj" in branch_forms:
+            if "TrigObj_mass" in branch_forms:
+                warnings.warn(
+                    "TrigObj_mass branch is present but will be replaced with zeros"
+                )
+            branch_forms["TrigObj_mass"] = transforms.zeros_from_offsets_form(
+                branch_forms["oTrigObj"]
+            )
+
         output = {}
         for name in collections:
             mixin = self.mixins.get(name, "NanoCollection")

src/coffea/nanoevents/trace.py

@@ -36,30 +36,39 @@ def _make_length_zero_one_tracer(
 ) -> tuple[ak.Array, list]:
     form = ak.forms.from_dict(events.attrs["@form"])
     buffer_key = events.attrs["@buffer_key"]
-    buffer_keys = form.expected_from_buffers(buffer_key=buffer_key).keys()
+    expected_buffer_keys = form.expected_from_buffers(buffer_key=buffer_key).keys()
 
     if length == 0:
-        buffers = ak.to_buffers(
-            form.length_zero_array(),
-            byteorder=ak._util.native_byteorder,
-        )[2].values()
+        tmp_array = form.length_zero_array()
     elif length == 1:
-        buffers = ak.to_buffers(
-            form.length_one_array(),
-            byteorder=ak._util.native_byteorder,
-        )[2].values()
+        tmp_array = form.length_one_array()
     else:
         raise ValueError("length must be 0 or 1")
 
+    def getkey(layout, form, attribute):
+        return buffer_key(
+            form_key=form.form_key,
+            attribute=attribute,
+            form=form,
+        )
+
+    container = {}
+    tmp_array._to_buffers(
+        form, getkey, container, events._layout.backend, ak._util.native_byteorder
+    )
+
     report = []
 
     def generate(buffer, report, buffer_key):
         report.append(buffer_key)
         return buffer
 
-    container = {}
-    for key, buffer in zip(buffer_keys, buffers):
+    for key, buffer in container.items():
         container[key] = partial(generate, buffer=buffer, report=report, buffer_key=key)
+
+    assert list(container.keys()) == list(
+        expected_buffer_keys
+    ), "length zero/one array buffer keys do not match the expected ones"
     array = ak.from_buffers(
         form=form,
         length=length,