feat: implement transforms that do zeros like of branches at the schema level (#1466)

* fun with forms

* updates

* delete alias transform

* avoid bad egm nano name

* photon energy fix too

* use float32

* itemsize should be 4

* remove useless parameters

* fix zeros from content

Author: ikrommyd

src/coffea/nanoevents/schemas/nanoaod.py

@@ -296,6 +296,24 @@ def _build_collections(self, field_names, input_contents):
             if all(k in branch_forms for k in args):
                 branch_forms[name] = fcn(*(branch_forms[k] for k in args))
 
+        # Add mass and charge fields for Photon collection (always zero for photons)
+        if "oPhoton" in branch_forms:
+            if "Photon_mass" not in branch_forms:
+                branch_forms["Photon_mass"] = transforms.zeros_from_offsets_form(
+                    branch_forms["oPhoton"]
+                )
+            if "Photon_charge" not in branch_forms:
+                branch_forms["Photon_charge"] = transforms.zeros_from_offsets_form(
+                    branch_forms["oPhoton"]
+                )
+
+        # Rename Electron/Photon_energy to Electron/Photon_regrEnergy to avoid conflict with mixin
+        # Present in EGamma NanoAOD flavor
+        if "Electron_energy" in branch_forms:
+            branch_forms["Electron_regrEnergy"] = branch_forms.pop("Electron_energy")
+        if "Photon_energy" in branch_forms:
+            branch_forms["Photon_regrEnergy"] = branch_forms.pop("Photon_energy")
+
         output = {}
         for name in collections:
             mixin = self.mixins.get(name, "NanoCollection")

src/coffea/nanoevents/transforms.py

@@ -487,6 +487,61 @@ def eventindex(stack):
     stack.append(out)
 
 
+def zeros_from_content_form(source_form):
+    form = copy.deepcopy(source_form)
+    if not (form["class"] == "NumpyArray" or form["class"].startswith("ListOffset")):
+        raise RuntimeError
+    if form["class"] == "NumpyArray":
+        form["form_key"] = concat(source_form["form_key"], "!zeros_from_content")
+        form["parameters"].pop("__doc__", None)
+    elif form["class"].startswith("ListOffset"):
+        form["content"]["form_key"] = concat(
+            source_form["form_key"], "!zeros_from_content", "!content"
+        )
+        form["parameters"].pop("__doc__", None)
+        form["content"]["parameters"].pop("__doc__", None)
+    return form
+
+
+def zeros_from_content(stack):
+    source = stack.pop()
+    stack.append(awkward.zeros_like(source))
+
+
+def zeros_from_offsets_form(offsets_form):
+    if not offsets_form["class"].startswith("NumpyArray"):
+        raise RuntimeError
+
+    form = {
+        "class": "ListOffsetArray",
+        "offsets": "i64",
+        "content": {
+            "class": "NumpyArray",
+            "primitive": "float32",
+            "itemsize": 4,
+            "format": "i",
+            "form_key": concat(
+                offsets_form["form_key"], "!zeros_from_offsets", "!content"
+            ),
+        },
+        "form_key": concat(offsets_form["form_key"], "!zeros_from_offsets"),
+    }
+    return form
+
+
+def zeros_from_offsets(stack):
+    offsets = ensure_array(stack.pop())
+    n_elements = offsets[-1]
+    content = numpy.zeros(n_elements, dtype=numpy.float32)
+    out = awkward.Array(
+        awkward.contents.ListOffsetArray(
+            awkward.index.Index64(offsets),
+            awkward.contents.NumpyArray(content),
+        )
+    )
+    stack.append(out)
+
+
 # For EDM4HEPSchema and FCCSChema: