diff --git a/src/coffea/nanoevents/__init__.py b/src/coffea/nanoevents/__init__.py index 58acec1e5..348ca0d8e 100644 --- a/src/coffea/nanoevents/__init__.py +++ b/src/coffea/nanoevents/__init__.py @@ -10,6 +10,7 @@ PDUNESchema, PFNanoAODSchema, PHYSLITESchema, + ScoutingNanoAODSchema, TreeMakerSchema, ) @@ -22,4 +23,5 @@ "PHYSLITESchema", "DelphesSchema", "PDUNESchema", + "ScoutingNanoAODSchema", ] diff --git a/src/coffea/nanoevents/schemas/__init__.py b/src/coffea/nanoevents/schemas/__init__.py index a1d93f4cf..71f92825e 100644 --- a/src/coffea/nanoevents/schemas/__init__.py +++ b/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", ] diff --git a/src/coffea/nanoevents/schemas/nanoaod.py b/src/coffea/nanoevents/schemas/nanoaod.py index ecf9bbaee..601b7f34e 100644 --- a/src/coffea/nanoevents/schemas/nanoaod.py +++ b/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} diff --git a/tests/samples/scouting_nano.root b/tests/samples/scouting_nano.root new file mode 100644 index 000000000..53f99bceb Binary files /dev/null and b/tests/samples/scouting_nano.root differ diff --git a/tests/test_nanoevents_scoutingnano.py b/tests/test_nanoevents_scoutingnano.py new file mode 100644 index 000000000..b039c13eb --- /dev/null +++ b/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)