Adding ScoutingNanoAODSchema and basic mixins that match collections …

…in the ScoutingNANO format.

Add a small test suite for ScoutingNanoAODSchema

Trimmed example root file down to 157K

src/coffea/nanoevents/__init__.py

@@ -10,6 +10,7 @@
     PDUNESchema,
     PFNanoAODSchema,
     PHYSLITESchema,
+    ScoutingNanoAODSchema,
     TreeMakerSchema,
 )
 
@@ -22,4 +23,5 @@
     "PHYSLITESchema",
     "DelphesSchema",
     "PDUNESchema",
+    "ScoutingNanoAODSchema",
 ]

src/coffea/nanoevents/schemas/__init__.py

@@ -1,6 +1,6 @@
 from .base import BaseSchema
 from .delphes import DelphesSchema
-from .nanoaod import NanoAODSchema, PFNanoAODSchema
+from .nanoaod import NanoAODSchema, PFNanoAODSchema, ScoutingNanoAODSchema
 from .pdune import PDUNESchema
 from .physlite import PHYSLITESchema
 from .treemaker import TreeMakerSchema
@@ -13,4 +13,5 @@
     "PHYSLITESchema",
     "DelphesSchema",
     "PDUNESchema",
+    "ScoutingNanoAODSchema",
 ]

src/coffea/nanoevents/schemas/nanoaod.py

@@ -375,3 +375,26 @@ class PFNanoAODSchema(NanoAODSchema):
         "JetSVs_sVIdx": "SV",
         "SubJet_subGenJetAK8Idx": "SubGenJetAK8",
     }
+
+
+class ScoutingNanoAODSchema(NanoAODSchema):
+    """ScoutingNano schema builder
+
+    ScoutingNano is a NanoAOD format that includes Scouting objects
+    """
+
+    mixins = {
+        **NanoAODSchema.mixins,
+        "ScoutingJet": "Jet",
+        "ScoutingFatJet": "FatJet",
+        "ScoutingMET": "MissingET",
+        "ScoutingMuonNoVtxDisplacedVertex": "Vertex",
+        "ScoutingMuonVtxDisplacedVertex": "Vertex",
+        "ScoutingPrimaryVertex": "Vertex",
+        "ScoutingElectron": "Electron",
+        "ScoutingPhoton": "Photon",
+        "ScoutingMuonNoVtx": "Muon",
+        "ScoutingMuonVtx": "Muon",
+    }
+
+    all_cross_references = {**NanoAODSchema.all_cross_references}

tests/test_nanoevents_scoutingnano.py

@@ -0,0 +1,41 @@
+import os
+
+import pytest
+
+from coffea.nanoevents import NanoEventsFactory, ScoutingNanoAODSchema
+
+
+@pytest.fixture(scope="module")
+def events(tests_directory):
+    path = os.path.join(tests_directory, "samples/scouting_nano.root")
+    ScoutingNanoAODSchema.warn_missing_crossrefs = False
+    events = NanoEventsFactory.from_root(
+        {path: "Events"},
+        schemaclass=ScoutingNanoAODSchema,
+        delayed=True,
+    ).events()
+    return events
+
+
+@pytest.mark.parametrize(
+    "field",
+    [
+        # Jet Sanity check
+        "ScoutingJet",
+        # associated PF candidate
+        "ScoutingJet.pt",
+        # FatJet sanity check
+        "ScoutingFatJet",
+        # Subjet collection (secondary sanity check)
+        "ScoutingFatJet.pt",
+    ],
+)
+def test_nested_collections(events, field):
+    def check_fields_recursive(coll, field):
+        if "." not in field:
+            assert hasattr(coll, field)
+        else:
+            split = field.split(".")
+            return check_fields_recursive(getattr(coll, split[0]), ".".join(split[1:]))
+
+    check_fields_recursive(events, field)