Merge pull request #1245 from prayagyadav/edm4hep-schema

feat: EDM4HEPSchema and Newstyle FCCSchema

Author: lgray

pyproject.toml

@@ -38,7 +38,7 @@ classifiers = [
 ]
 dependencies = [
   "awkward>=2.6.7",
-  "uproot>=5.3.11",
+  "uproot>=5.6.0",
   "dask[array]>=2024.3.0",
   "dask-awkward>=2025.2.0",
   "dask-histogram>=2025.2.0",

src/coffea/dataset_tools/preprocess.py

@@ -83,7 +83,7 @@ def get_steps(
         if save_form:
             form_str = uproot.dask(
                 tree,
-                ak_add_doc=True,
+                ak_add_doc={"__doc__": "title", "typename": "typename"},
                 filter_name=no_filter,
                 filter_typename=no_filter,
                 filter_branch=partial(_remove_not_interpretable, emit_warning=False),

src/coffea/nanoevents/__init__.py

@@ -5,7 +5,9 @@
     FCC,
     BaseSchema,
     DelphesSchema,
+    EDM4HEPSchema,
     FCCSchema,
+    FCCSchema_edm4hep1,
     NanoAODSchema,
     PDUNESchema,
     PFNanoAODSchema,
@@ -26,4 +28,6 @@
     "ScoutingNanoAODSchema",
     "FCC",
     "FCCSchema",
+    "FCCSchema_edm4hep1",
+    "EDM4HEPSchema",
 ]

src/coffea/nanoevents/assets/__init__.py

@@ -0,0 +1,32 @@
+import importlib
+import os
+from functools import partial
+
+import yaml
+
+root_dir = importlib.resources.files("coffea.nanoevents.assets")
+
+versions = [
+    "00-10-01",
+    "00-10-02",
+    "00-10-03",
+    "00-10-04",
+    "00-10-05",
+    "00-99-00",
+    "00-99-01",
+]
+
+
+def _load_edm4hep_version(yamlfile):
+    with open(yamlfile) as f:
+        loaded = yaml.safe_load(f)
+    return loaded
+
+
+edm4hep_ver = {
+    version: partial(
+        _load_edm4hep_version,
+        yamlfile=os.path.join(root_dir, f"edm4hep_v{version}.yaml"),
+    )
+    for version in versions
+}

src/coffea/nanoevents/assets/edm4hep_v00-10-01.yaml

@@ -117,7 +117,10 @@ def __call__(self, form):
         for ifield, field in enumerate(form.fields):
             iform = form.contents[ifield].to_dict()
             branch_forms[field] = _lazify_form(
-                iform, f"{field},!load", docstr=iform["parameters"]["__doc__"]
+                iform,
+                f"{field},!load",
+                docstr=iform["parameters"]["__doc__"],
+                typestr=iform["parameters"]["typename"],
             )
         lform = {
             "class": "RecordArray",
@@ -129,6 +132,7 @@ def __call__(self, form):
             },
             "form_key": None,
         }
+
         return (
             awkward.forms.form.from_dict(self.schemaclass(lform, self.version).form),
             self,
@@ -340,13 +344,12 @@ def from_root(
             to_open = file
             if isinstance(file, uproot.reading.ReadOnlyDirectory):
                 to_open = file[treepath]
-
             opener = partial(
                 uproot.dask,
                 to_open,
                 full_paths=True,
                 open_files=False,
-                ak_add_doc=True,
+                ak_add_doc={"__doc__": "title", "typename": "typename"},
                 filter_branch=_remove_not_interpretable,
                 steps_per_file=steps_per_file,
                 known_base_form=known_base_form,
@@ -372,6 +375,9 @@ def from_root(
         else:
             tree = uproot.open(file, **uproot_options)
 
+        # Get the typenames
+        typenames = tree.typenames()
+
         if entry_start is None or entry_start < 0:
             entry_start = 0
         if entry_stop is None or entry_stop > tree.num_entries:
@@ -396,6 +402,7 @@ def from_root(
         base_form = mapping._extract_base_form(
             tree, iteritems_options=iteritems_options
         )
+        base_form["typenames"] = typenames
 
         return cls._from_mapping(
             mapping,

src/coffea/nanoevents/assets/edm4hep_v00-10-02.yaml

@@ -28,24 +28,26 @@ class CannotBeNanoEvents(Exception):
     pass
 
 
-def _lazify_form(form, prefix, docstr=None):
+def _lazify_form(form, prefix, docstr=None, typestr=None):
     if not isinstance(form, dict) or "class" not in form:
         raise RuntimeError("form should have been normalized by now")
 
-    parameters = _lazify_parameters(form.get("parameters", {}), docstr=docstr)
+    parameters = _lazify_parameters(
+        form.get("parameters", {}), docstr=docstr, typestr=typestr
+    )
     if form["class"].startswith("ListOffset"):
         # awkward will add !offsets
         form["form_key"] = quote(prefix)
         form["content"] = _lazify_form(
-            form["content"], prefix + ",!content", docstr=docstr
+            form["content"], prefix + ",!content", docstr=docstr, typestr=typestr
         )
     elif form["class"] == "NumpyArray":
         form["form_key"] = quote(prefix)
         if parameters:
             form["parameters"] = parameters
     elif form["class"] == "RegularArray":
         form["content"] = _lazify_form(
-            form["content"], prefix + ",!content", docstr=docstr
+            form["content"], prefix + ",!content", docstr=docstr, typestr=typestr
         )
         if parameters:
             form["parameters"] = parameters
@@ -61,7 +63,10 @@ def _lazify_form(form, prefix, docstr=None):
         assert prefix.endswith("!load")
         form["form_key"] = quote(prefix + "allowmissing,!index")
         form["content"] = _lazify_form(
-            form["content"], prefix + "allowmissing,!content", docstr=docstr
+            form["content"],
+            prefix + "allowmissing,!content",
+            docstr=docstr,
+            typestr=typestr,
         )
         if parameters:
             form["parameters"] = parameters
@@ -88,12 +93,16 @@ def _lazify_form(form, prefix, docstr=None):
     return form
 
 
-def _lazify_parameters(form_parameters, docstr=None):
+def _lazify_parameters(form_parameters, docstr=None, typestr=None):
     parameters = {}
     if "__array__" in form_parameters:
         parameters["__array__"] = form_parameters["__array__"]
     if docstr is not None:
         parameters["__doc__"] = docstr
+    if typestr is not None:
+        parameters["typename"] = typestr
+    if "typename" in form_parameters:  # eager mode
+        parameters["typename"] = form_parameters["typename"]
     return parameters
 
 
@@ -159,13 +168,16 @@ def _extract_base_form(cls, tree, iteritems_options={}):
                 form.to_json()
             )  # normalizes form (expand NumpyArray classes)
             try:
-                form = _lazify_form(form, f"{key},!load", docstr=branch.title)
+                form = _lazify_form(
+                    form, f"{key},!load", docstr=branch.title, typestr=branch.typename
+                )
             except CannotBeNanoEvents as ex:
                 warnings.warn(
                     f"Skipping {key} as it is not interpretable by NanoEvents\nDetails: {ex}"
                 )
                 continue
             branch_forms[key] = form
+
         return {
             "class": "RecordArray",
             "contents": [item for item in branch_forms.values()],