-
Notifications
You must be signed in to change notification settings - Fork 120
/
test_nanoevents.py
195 lines (161 loc) · 6.16 KB
/
test_nanoevents.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
import os
from pathlib import Path
import awkward as ak
import pytest
from distributed import Client
from coffea.nanoevents import NanoAODSchema, NanoEventsFactory
def genroundtrips(genpart):
# check genpart roundtrip
assert ak.all(genpart.children.parent.pdgId == genpart.pdgId)
assert ak.all(
ak.any(
genpart.parent.children.pdgId == genpart.pdgId, axis=-1, mask_identity=True
)
)
# distinctParent should be distinct and it should have a relevant child
assert ak.all(genpart.distinctParent.pdgId != genpart.pdgId)
assert ak.all(
ak.any(
genpart.distinctParent.children.pdgId == genpart.pdgId,
axis=-1,
mask_identity=True,
)
)
# distinctChildren should be distinct
assert ak.all(genpart.distinctChildren.pdgId != genpart.pdgId)
# their distinctParent's should be the particle itself
assert ak.all(genpart.distinctChildren.distinctParent.pdgId == genpart.pdgId)
# parents in decay chains (same pdg id) should never have distinctChildrenDeep
parents_in_decays = genpart[genpart.parent.pdgId == genpart.pdgId]
assert ak.all(ak.num(parents_in_decays.distinctChildrenDeep, axis=2) == 0)
# parents at the top of decay chains that have children should always have distinctChildrenDeep
real_parents_at_top = genpart[
(genpart.parent.pdgId != genpart.pdgId) & (ak.num(genpart.children, axis=2) > 0)
]
assert ak.all(ak.num(real_parents_at_top.distinctChildrenDeep, axis=2) > 0)
# distinctChildrenDeep whose parent pdg id is the same must not have children
children_in_decays = genpart.distinctChildrenDeep[
genpart.distinctChildrenDeep.pdgId == genpart.distinctChildrenDeep.parent.pdgId
]
assert ak.all(ak.num(children_in_decays.children, axis=3) == 0)
# exercise hasFlags
genpart.hasFlags(["isHardProcess"])
genpart.hasFlags(["isHardProcess", "isDecayedLeptonHadron"])
def crossref(events):
# check some cross-ref roundtrips (some may not be true always but they are for the test file)
assert ak.all(events.Jet.matched_muons.matched_jet.pt == events.Jet.pt)
assert ak.all(
events.Electron.matched_photon.matched_electron.r9 == events.Electron.r9
)
# exercise LorentzVector.nearest
assert ak.all(
events.Muon.matched_jet.delta_r(events.Muon.nearest(events.Jet)) == 0.0
)
suffixes = [
"root",
# "parquet",
]
@pytest.mark.parametrize("suffix", suffixes)
def test_read_nanomc(suffix):
path = os.path.abspath(f"tests/samples/nano_dy.{suffix}")
# parquet files were converted from even older nanoaod
nanoversion = NanoAODSchema
factory = getattr(NanoEventsFactory, f"from_{suffix}")(
{path: "Events"},
schemaclass=nanoversion,
delayed=False,
)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all(
(abs(events.Electron.matched_gen.pdgId) == 11)
| (events.Electron.matched_gen.pdgId == 22)
)
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
# test issue 409
assert ak.to_list(events[[]].Photon.mass) == []
if suffix == "root":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
True,
True,
False,
False,
False,
False,
False,
]
if suffix == "parquet":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
True,
True,
False,
False,
False,
False,
False,
]
@pytest.mark.parametrize("suffix", suffixes)
def test_read_from_uri(suffix):
"Make sure we can properly open the file when a uri is used"
path = Path(os.path.abspath(f"tests/samples/nano_dy.{suffix}")).as_uri()
nanoversion = NanoAODSchema
factory = getattr(NanoEventsFactory, f"from_{suffix}")(
{path: "Events"},
schemaclass=nanoversion,
delayed=False,
)
events = factory.events()
assert len(events) == 40 if suffix == "root" else 10
@pytest.mark.parametrize("suffix", suffixes)
def test_read_nanodata(suffix):
path = os.path.abspath(f"tests/samples/nano_dimuon.{suffix}")
# parquet files were converted from even older nanoaod
nanoversion = NanoAODSchema
factory = getattr(NanoEventsFactory, f"from_{suffix}")(
{path: "Events"},
schemaclass=nanoversion,
delayed=False,
)
events = factory.events()
crossref(events)
crossref(events[ak.num(events.Jet) > 2])
def test_missing_eventIds_error():
path = os.path.abspath("tests/samples/missing_luminosityBlock.root") + ":Events"
with pytest.raises(RuntimeError):
factory = NanoEventsFactory.from_root(
path, schemaclass=NanoAODSchema, delayed=False
)
factory.events()
def test_missing_eventIds_warning():
path = os.path.abspath("tests/samples/missing_luminosityBlock.root") + ":Events"
with pytest.warns(
RuntimeWarning, match=r"Missing event_ids \: \[\'luminosityBlock\'\]"
):
NanoAODSchema.error_missing_event_ids = False
factory = NanoEventsFactory.from_root(
path, schemaclass=NanoAODSchema, delayed=False
)
factory.events()
def test_missing_eventIds_warning_dask():
path = os.path.abspath("tests/samples/missing_luminosityBlock.root") + ":Events"
NanoAODSchema.error_missing_event_ids = False
with Client() as _:
events = NanoEventsFactory.from_root(
path,
schemaclass=NanoAODSchema,
delayed=True,
).events()
events.Muon.pt.compute()