atlas-schema 0.2.4__tar.gz → 0.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/PKG-INFO +68 -18
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/README.md +65 -15
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/pyproject.toml +3 -2
- atlas_schema-0.4.0/src/atlas_schema/_version.py +34 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/enums.py +1 -1
- atlas_schema-0.4.0/src/atlas_schema/methods.py +363 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/schema.py +234 -82
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/utils.py +6 -7
- atlas_schema-0.2.4/src/atlas_schema/_version.py +0 -16
- atlas_schema-0.2.4/src/atlas_schema/methods.py +0 -253
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/.gitignore +0 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/LICENSE +0 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/__init__.py +0 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/_version.pyi +0 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/py.typed +0 -0
- {atlas_schema-0.2.4 → atlas_schema-0.4.0}/src/atlas_schema/typing_compat.py +0 -0
@@ -1,11 +1,11 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: atlas-schema
|
3
|
-
Version: 0.
|
3
|
+
Version: 0.4.0
|
4
4
|
Summary: Helper python package for ATLAS Common NTuple Analysis work.
|
5
5
|
Project-URL: Homepage, https://github.com/scipp-atlas/atlas-schema
|
6
6
|
Project-URL: Bug Tracker, https://github.com/scipp-atlas/atlas-schema/issues
|
7
7
|
Project-URL: Discussions, https://github.com/scipp-atlas/atlas-schema/discussions
|
8
|
-
Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.
|
8
|
+
Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.4.0/
|
9
9
|
Project-URL: Releases, https://github.com/scipp-atlas/atlas-schema/releases
|
10
10
|
Project-URL: Release Notes, https://atlas-schema.readthedocs.io/en/latest/history.html
|
11
11
|
Author-email: Giordon Stark <kratsg@gmail.com>
|
@@ -227,7 +227,7 @@ Classifier: Programming Language :: Python :: 3.12
|
|
227
227
|
Classifier: Topic :: Scientific/Engineering
|
228
228
|
Classifier: Typing :: Typed
|
229
229
|
Requires-Python: >=3.9
|
230
|
-
Requires-Dist: coffea[dask]>=
|
230
|
+
Requires-Dist: coffea[dask]>=2025.7.0
|
231
231
|
Requires-Dist: particle>=0.25.0
|
232
232
|
Provides-Extra: dev
|
233
233
|
Requires-Dist: pytest-cov>=3; extra == 'dev'
|
@@ -251,7 +251,7 @@ Requires-Dist: tbump>=6.7.0; extra == 'test'
|
|
251
251
|
Requires-Dist: twine; extra == 'test'
|
252
252
|
Description-Content-Type: text/markdown
|
253
253
|
|
254
|
-
# atlas-schema v0.
|
254
|
+
# atlas-schema v0.4.0
|
255
255
|
|
256
256
|
[![Actions Status][actions-badge]][actions-link]
|
257
257
|
[![Documentation Status][rtd-badge]][rtd-link]
|
@@ -335,11 +335,9 @@ like below:
|
|
335
335
|
|
336
336
|
```python
|
337
337
|
import awkward as ak
|
338
|
-
import
|
339
|
-
import hist.dask as had
|
338
|
+
from hist import Hist
|
340
339
|
import matplotlib.pyplot as plt
|
341
340
|
from coffea import processor
|
342
|
-
from coffea.nanoevents import NanoEventsFactory
|
343
341
|
from distributed import Client
|
344
342
|
|
345
343
|
from atlas_schema.schema import NtupleSchema
|
@@ -352,7 +350,7 @@ class MyFirstProcessor(processor.ProcessorABC):
|
|
352
350
|
def process(self, events):
|
353
351
|
dataset = events.metadata["dataset"]
|
354
352
|
h_ph_pt = (
|
355
|
-
|
353
|
+
Hist.new.StrCat(["all", "pass", "fail"], name="isEM")
|
356
354
|
.Regular(200, 0.0, 2000.0, name="pt", label="$pt_{\gamma}$ [GeV]")
|
357
355
|
.Int64()
|
358
356
|
)
|
@@ -376,17 +374,18 @@ class MyFirstProcessor(processor.ProcessorABC):
|
|
376
374
|
if __name__ == "__main__":
|
377
375
|
client = Client()
|
378
376
|
|
379
|
-
|
380
|
-
events = NanoEventsFactory.from_root(
|
381
|
-
{fname: "analysis"},
|
382
|
-
schemaclass=NtupleSchema,
|
383
|
-
metadata={"dataset": "700352.Zqqgamma.mc20d.v1"},
|
384
|
-
).events()
|
377
|
+
fileset = {"700352.Zqqgamma.mc20d.v1": {"files": {"ntuple.root": "analysis"}}}
|
385
378
|
|
386
|
-
|
387
|
-
|
388
|
-
|
389
|
-
|
379
|
+
run = processor.Runner(
|
380
|
+
executor=processor.IterativeExecutor(compression=None),
|
381
|
+
schema=NtupleSchema,
|
382
|
+
savemetrics=True,
|
383
|
+
)
|
384
|
+
|
385
|
+
out, metrics = run(fileset, processor_instance=MyFirstProcessor())
|
386
|
+
|
387
|
+
print(out)
|
388
|
+
print(metrics)
|
390
389
|
|
391
390
|
fig, ax = plt.subplots()
|
392
391
|
computed["700352.Zqqgamma.mc20d.v1"]["ph_pt"].plot1d(ax=ax)
|
@@ -400,6 +399,57 @@ which produces
|
|
400
399
|
|
401
400
|
<img src="https://raw.githubusercontent.com/scipp-atlas/atlas-schema/main/docs/_static/img/ph_pt.png" alt="three stacked histograms of photon pT, with each stack corresponding to: no selection, requiring the isEM flag, and inverting the isEM requirement" width="500" style="display: block; margin-left: auto; margin-right: auto;">
|
402
401
|
|
402
|
+
## Processing with Systematic Variations
|
403
|
+
|
404
|
+
For analyses requiring systematic uncertainty evaluation, you can easily iterate
|
405
|
+
over all systematic variations using the new `events["NOSYS"]` alias and
|
406
|
+
`systematic_names` property:
|
407
|
+
|
408
|
+
```python
|
409
|
+
import awkward as ak
|
410
|
+
from hist import Hist
|
411
|
+
from coffea import processor
|
412
|
+
from atlas_schema.schema import NtupleSchema
|
413
|
+
|
414
|
+
|
415
|
+
class SystematicsProcessor(processor.ProcessorABC):
|
416
|
+
def __init__(self):
|
417
|
+
self.h = (
|
418
|
+
Hist.new.StrCat([], name="variation", growth=True)
|
419
|
+
.Regular(50, 0.0, 500.0, name="jet_pt", label="Leading Jet $p_T$ [GeV]")
|
420
|
+
.Int64()
|
421
|
+
)
|
422
|
+
|
423
|
+
def process(self, events):
|
424
|
+
dsid = events.metadata["dataset"]
|
425
|
+
|
426
|
+
# Process all systematic variations including nominal ("NOSYS")
|
427
|
+
for variation in events.systematic_names:
|
428
|
+
event_view = events[variation]
|
429
|
+
|
430
|
+
# Fill histogram with leading jet pT for this systematic variation
|
431
|
+
leading_jet_pt = event_view.jet.pt[:, 0] / 1_000 # Convert MeV to GeV
|
432
|
+
weights = (
|
433
|
+
event_view.weight.mc
|
434
|
+
if hasattr(event_view, "weight")
|
435
|
+
else ak.ones_like(leading_jet_pt)
|
436
|
+
)
|
437
|
+
|
438
|
+
self.h.fill(variation=variation, jet_pt=leading_jet_pt, weight=weights)
|
439
|
+
|
440
|
+
return {
|
441
|
+
"hist": self.h,
|
442
|
+
"meta": {"sumw": {dsid: {(events.metadata["fileuuid"], ak.sum(weights))}}},
|
443
|
+
}
|
444
|
+
|
445
|
+
def postprocess(self, accumulator):
|
446
|
+
return accumulator
|
447
|
+
```
|
448
|
+
|
449
|
+
This approach allows you to seamlessly process both nominal and systematic
|
450
|
+
variations in a single loop, eliminating the need for special-case handling of
|
451
|
+
the nominal variation.
|
452
|
+
|
403
453
|
<!-- SPHINX-END -->
|
404
454
|
|
405
455
|
## Developer Notes
|
@@ -1,4 +1,4 @@
|
|
1
|
-
# atlas-schema v0.
|
1
|
+
# atlas-schema v0.4.0
|
2
2
|
|
3
3
|
[![Actions Status][actions-badge]][actions-link]
|
4
4
|
[![Documentation Status][rtd-badge]][rtd-link]
|
@@ -82,11 +82,9 @@ like below:
|
|
82
82
|
|
83
83
|
```python
|
84
84
|
import awkward as ak
|
85
|
-
import
|
86
|
-
import hist.dask as had
|
85
|
+
from hist import Hist
|
87
86
|
import matplotlib.pyplot as plt
|
88
87
|
from coffea import processor
|
89
|
-
from coffea.nanoevents import NanoEventsFactory
|
90
88
|
from distributed import Client
|
91
89
|
|
92
90
|
from atlas_schema.schema import NtupleSchema
|
@@ -99,7 +97,7 @@ class MyFirstProcessor(processor.ProcessorABC):
|
|
99
97
|
def process(self, events):
|
100
98
|
dataset = events.metadata["dataset"]
|
101
99
|
h_ph_pt = (
|
102
|
-
|
100
|
+
Hist.new.StrCat(["all", "pass", "fail"], name="isEM")
|
103
101
|
.Regular(200, 0.0, 2000.0, name="pt", label="$pt_{\gamma}$ [GeV]")
|
104
102
|
.Int64()
|
105
103
|
)
|
@@ -123,17 +121,18 @@ class MyFirstProcessor(processor.ProcessorABC):
|
|
123
121
|
if __name__ == "__main__":
|
124
122
|
client = Client()
|
125
123
|
|
126
|
-
|
127
|
-
events = NanoEventsFactory.from_root(
|
128
|
-
{fname: "analysis"},
|
129
|
-
schemaclass=NtupleSchema,
|
130
|
-
metadata={"dataset": "700352.Zqqgamma.mc20d.v1"},
|
131
|
-
).events()
|
124
|
+
fileset = {"700352.Zqqgamma.mc20d.v1": {"files": {"ntuple.root": "analysis"}}}
|
132
125
|
|
133
|
-
|
134
|
-
|
135
|
-
|
136
|
-
|
126
|
+
run = processor.Runner(
|
127
|
+
executor=processor.IterativeExecutor(compression=None),
|
128
|
+
schema=NtupleSchema,
|
129
|
+
savemetrics=True,
|
130
|
+
)
|
131
|
+
|
132
|
+
out, metrics = run(fileset, processor_instance=MyFirstProcessor())
|
133
|
+
|
134
|
+
print(out)
|
135
|
+
print(metrics)
|
137
136
|
|
138
137
|
fig, ax = plt.subplots()
|
139
138
|
computed["700352.Zqqgamma.mc20d.v1"]["ph_pt"].plot1d(ax=ax)
|
@@ -147,6 +146,57 @@ which produces
|
|
147
146
|
|
148
147
|
<img src="https://raw.githubusercontent.com/scipp-atlas/atlas-schema/main/docs/_static/img/ph_pt.png" alt="three stacked histograms of photon pT, with each stack corresponding to: no selection, requiring the isEM flag, and inverting the isEM requirement" width="500" style="display: block; margin-left: auto; margin-right: auto;">
|
149
148
|
|
149
|
+
## Processing with Systematic Variations
|
150
|
+
|
151
|
+
For analyses requiring systematic uncertainty evaluation, you can easily iterate
|
152
|
+
over all systematic variations using the new `events["NOSYS"]` alias and
|
153
|
+
`systematic_names` property:
|
154
|
+
|
155
|
+
```python
|
156
|
+
import awkward as ak
|
157
|
+
from hist import Hist
|
158
|
+
from coffea import processor
|
159
|
+
from atlas_schema.schema import NtupleSchema
|
160
|
+
|
161
|
+
|
162
|
+
class SystematicsProcessor(processor.ProcessorABC):
|
163
|
+
def __init__(self):
|
164
|
+
self.h = (
|
165
|
+
Hist.new.StrCat([], name="variation", growth=True)
|
166
|
+
.Regular(50, 0.0, 500.0, name="jet_pt", label="Leading Jet $p_T$ [GeV]")
|
167
|
+
.Int64()
|
168
|
+
)
|
169
|
+
|
170
|
+
def process(self, events):
|
171
|
+
dsid = events.metadata["dataset"]
|
172
|
+
|
173
|
+
# Process all systematic variations including nominal ("NOSYS")
|
174
|
+
for variation in events.systematic_names:
|
175
|
+
event_view = events[variation]
|
176
|
+
|
177
|
+
# Fill histogram with leading jet pT for this systematic variation
|
178
|
+
leading_jet_pt = event_view.jet.pt[:, 0] / 1_000 # Convert MeV to GeV
|
179
|
+
weights = (
|
180
|
+
event_view.weight.mc
|
181
|
+
if hasattr(event_view, "weight")
|
182
|
+
else ak.ones_like(leading_jet_pt)
|
183
|
+
)
|
184
|
+
|
185
|
+
self.h.fill(variation=variation, jet_pt=leading_jet_pt, weight=weights)
|
186
|
+
|
187
|
+
return {
|
188
|
+
"hist": self.h,
|
189
|
+
"meta": {"sumw": {dsid: {(events.metadata["fileuuid"], ak.sum(weights))}}},
|
190
|
+
}
|
191
|
+
|
192
|
+
def postprocess(self, accumulator):
|
193
|
+
return accumulator
|
194
|
+
```
|
195
|
+
|
196
|
+
This approach allows you to seamlessly process both nominal and systematic
|
197
|
+
variations in a single loop, eliminating the need for special-case handling of
|
198
|
+
the nominal variation.
|
199
|
+
|
150
200
|
<!-- SPHINX-END -->
|
151
201
|
|
152
202
|
## Developer Notes
|
@@ -28,7 +28,7 @@ classifiers = [
|
|
28
28
|
"Typing :: Typed",
|
29
29
|
]
|
30
30
|
dynamic = ["version"]
|
31
|
-
dependencies = ["coffea[dask] >=
|
31
|
+
dependencies = ["coffea[dask] >= 2025.7.0", "particle >= 0.25.0"]
|
32
32
|
|
33
33
|
[project.optional-dependencies]
|
34
34
|
test = [
|
@@ -60,7 +60,7 @@ docs = [
|
|
60
60
|
Homepage = "https://github.com/scipp-atlas/atlas-schema"
|
61
61
|
"Bug Tracker" = "https://github.com/scipp-atlas/atlas-schema/issues"
|
62
62
|
Discussions = "https://github.com/scipp-atlas/atlas-schema/discussions"
|
63
|
-
Documentation = "https://atlas-schema.readthedocs.io/en/v0.
|
63
|
+
Documentation = "https://atlas-schema.readthedocs.io/en/v0.4.0/"
|
64
64
|
Releases = "https://github.com/scipp-atlas/atlas-schema/releases"
|
65
65
|
"Release Notes" = "https://atlas-schema.readthedocs.io/en/latest/history.html"
|
66
66
|
|
@@ -112,6 +112,7 @@ addopts = [
|
|
112
112
|
xfail_strict = true
|
113
113
|
filterwarnings = [
|
114
114
|
"error",
|
115
|
+
"ignore:In version 2025.1.0 .*, this will be an error:FutureWarning",
|
115
116
|
]
|
116
117
|
|
117
118
|
log_cli_level = "INFO"
|
@@ -0,0 +1,34 @@
|
|
1
|
+
# file generated by setuptools-scm
|
2
|
+
# don't change, don't track in version control
|
3
|
+
|
4
|
+
__all__ = [
|
5
|
+
"__version__",
|
6
|
+
"__version_tuple__",
|
7
|
+
"version",
|
8
|
+
"version_tuple",
|
9
|
+
"__commit_id__",
|
10
|
+
"commit_id",
|
11
|
+
]
|
12
|
+
|
13
|
+
TYPE_CHECKING = False
|
14
|
+
if TYPE_CHECKING:
|
15
|
+
from typing import Tuple
|
16
|
+
from typing import Union
|
17
|
+
|
18
|
+
VERSION_TUPLE = Tuple[Union[int, str], ...]
|
19
|
+
COMMIT_ID = Union[str, None]
|
20
|
+
else:
|
21
|
+
VERSION_TUPLE = object
|
22
|
+
COMMIT_ID = object
|
23
|
+
|
24
|
+
version: str
|
25
|
+
__version__: str
|
26
|
+
__version_tuple__: VERSION_TUPLE
|
27
|
+
version_tuple: VERSION_TUPLE
|
28
|
+
commit_id: COMMIT_ID
|
29
|
+
__commit_id__: COMMIT_ID
|
30
|
+
|
31
|
+
__version__ = version = '0.4.0'
|
32
|
+
__version_tuple__ = version_tuple = (0, 4, 0)
|
33
|
+
|
34
|
+
__commit_id__ = commit_id = None
|
@@ -18,7 +18,7 @@ class MultipleEnumAccessMeta(EnumType):
|
|
18
18
|
Enum Metaclass to provide a way to access multiple values all at once.
|
19
19
|
"""
|
20
20
|
|
21
|
-
def __getitem__(
|
21
|
+
def __getitem__(cls: type[_E], key: str | tuple[str]) -> _E | list[_E]: # type:ignore[misc,override]
|
22
22
|
getitem = cast(Callable[[str], _E], super().__getitem__) # type:ignore[misc]
|
23
23
|
if isinstance(key, tuple):
|
24
24
|
return [getitem(name) for name in key]
|
@@ -0,0 +1,363 @@
|
|
1
|
+
"""Mixins for the Ntuple schema"""
|
2
|
+
|
3
|
+
from __future__ import annotations
|
4
|
+
|
5
|
+
from functools import reduce
|
6
|
+
from operator import ior
|
7
|
+
|
8
|
+
import awkward
|
9
|
+
import particle
|
10
|
+
from coffea.nanoevents.methods import base, candidate, vector
|
11
|
+
|
12
|
+
from atlas_schema.enums import PhotonID
|
13
|
+
from atlas_schema.typing_compat import Behavior
|
14
|
+
|
15
|
+
behavior: Behavior = {}
|
16
|
+
behavior.update(base.behavior)
|
17
|
+
# vector behavior is included in candidate behavior
|
18
|
+
behavior.update(candidate.behavior)
|
19
|
+
|
20
|
+
|
21
|
+
def _set_repr_name(classname):
|
22
|
+
def namefcn(_self):
|
23
|
+
return classname
|
24
|
+
|
25
|
+
behavior[("__typestr__", classname)] = classname[0].lower() + classname[1:]
|
26
|
+
behavior[classname].__repr__ = namefcn
|
27
|
+
|
28
|
+
|
29
|
+
class NtupleEvents(behavior["NanoEvents"]): # type: ignore[misc, valid-type, name-defined]
|
30
|
+
"""Individual systematic variation of events."""
|
31
|
+
|
32
|
+
def __repr__(self):
|
33
|
+
return f"<event {getattr(self, 'runNumber', '??')}:{getattr(self, 'eventNumber', '??')}:{getattr(self, 'mcChannelNumber', '??')}>"
|
34
|
+
|
35
|
+
def __getitem__(self, key):
|
36
|
+
"""Support accessing systematic variations via bracket notation.
|
37
|
+
|
38
|
+
Args:
|
39
|
+
key: The systematic variation name. "NOSYS" returns the nominal events.
|
40
|
+
|
41
|
+
Returns:
|
42
|
+
The requested systematic variation or nominal events for "NOSYS".
|
43
|
+
"""
|
44
|
+
if key == "NOSYS":
|
45
|
+
return self
|
46
|
+
return super().__getitem__(key)
|
47
|
+
|
48
|
+
@property
|
49
|
+
def systematic(self):
|
50
|
+
"""Get the systematic variation name for this event collection."""
|
51
|
+
return "nominal"
|
52
|
+
|
53
|
+
@property
|
54
|
+
def systematic_names(self):
|
55
|
+
"""Get all systematic variations available in this event collection.
|
56
|
+
|
57
|
+
Returns a list of systematic variation names, including 'NOSYS' for nominal.
|
58
|
+
"""
|
59
|
+
# Get systematics from metadata stored during schema building
|
60
|
+
systematics = self.metadata.get("systematics", [])
|
61
|
+
return ["NOSYS", *systematics]
|
62
|
+
|
63
|
+
@property
|
64
|
+
def systematics(self):
|
65
|
+
"""Get all systematic variations available in this event collection.
|
66
|
+
|
67
|
+
Returns a list of systematic variation names, excluding 'nominal'.
|
68
|
+
"""
|
69
|
+
# Get systematics from metadata stored during schema building
|
70
|
+
return [
|
71
|
+
getattr(self, systematic)
|
72
|
+
for systematic in self.systematic_names
|
73
|
+
if systematic != "NOSYS"
|
74
|
+
]
|
75
|
+
|
76
|
+
|
77
|
+
behavior["NtupleEvents"] = NtupleEvents
|
78
|
+
|
79
|
+
|
80
|
+
class NtupleEventsArray(behavior[("*", "NanoEvents")]): # type: ignore[misc, valid-type, name-defined]
|
81
|
+
"""Collection of NtupleEvents objects, one for each systematic variation."""
|
82
|
+
|
83
|
+
def __getitem__(self, key):
|
84
|
+
"""Support accessing systematic variations via bracket notation.
|
85
|
+
|
86
|
+
Args:
|
87
|
+
key: The systematic variation name. "NOSYS" returns the nominal events.
|
88
|
+
|
89
|
+
Returns:
|
90
|
+
The requested systematic variation or nominal events for "NOSYS".
|
91
|
+
"""
|
92
|
+
if key == "NOSYS":
|
93
|
+
return self
|
94
|
+
return super().__getitem__(key)
|
95
|
+
|
96
|
+
@property
|
97
|
+
def systematic_names(self):
|
98
|
+
"""Get all systematic variations available in this event collection.
|
99
|
+
|
100
|
+
Returns a list of systematic variation names, including 'NOSYS' for nominal.
|
101
|
+
"""
|
102
|
+
# Get systematics from metadata stored during schema building
|
103
|
+
systematics = self.metadata.get("systematics", [])
|
104
|
+
return ["NOSYS", *systematics]
|
105
|
+
|
106
|
+
@property
|
107
|
+
def systematics(self):
|
108
|
+
"""Get all systematic variations available in this event collection.
|
109
|
+
|
110
|
+
Returns a list of systematic variation names, excluding 'nominal'.
|
111
|
+
"""
|
112
|
+
# Get systematics from metadata stored during schema building
|
113
|
+
return [
|
114
|
+
getattr(self, systematic)
|
115
|
+
for systematic in self.systematic_names
|
116
|
+
if systematic != "NOSYS"
|
117
|
+
]
|
118
|
+
|
119
|
+
|
120
|
+
behavior[("*", "NtupleEvents")] = NtupleEventsArray
|
121
|
+
|
122
|
+
|
123
|
+
@awkward.mixin_class(behavior)
|
124
|
+
class Systematic(base.NanoCollection, base.Systematic):
|
125
|
+
"""Base class for systematic variations."""
|
126
|
+
|
127
|
+
@property
|
128
|
+
def metadata(self):
|
129
|
+
"""Arbitrary metadata"""
|
130
|
+
return self.layout.purelist_parameter("metadata") # pylint: disable=no-member
|
131
|
+
|
132
|
+
@property
|
133
|
+
def systematic(self):
|
134
|
+
"""Get the systematic variation name for this event collection."""
|
135
|
+
return self.metadata["systematic"]
|
136
|
+
|
137
|
+
def __repr__(self):
|
138
|
+
return f"<event {self.systematic}>"
|
139
|
+
|
140
|
+
|
141
|
+
_set_repr_name("Systematic")
|
142
|
+
|
143
|
+
|
144
|
+
@awkward.mixin_class(behavior)
|
145
|
+
class Weight(base.NanoCollection, base.Systematic): ...
|
146
|
+
|
147
|
+
|
148
|
+
_set_repr_name("Weight")
|
149
|
+
|
150
|
+
|
151
|
+
@awkward.mixin_class(behavior)
|
152
|
+
class Pass(base.NanoCollection, base.Systematic): ...
|
153
|
+
|
154
|
+
|
155
|
+
_set_repr_name("Pass")
|
156
|
+
|
157
|
+
behavior.update(
|
158
|
+
awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) # pylint: disable=protected-access
|
159
|
+
)
|
160
|
+
|
161
|
+
|
162
|
+
@awkward.mixin_class(behavior)
|
163
|
+
class Particle(vector.PtEtaPhiMLorentzVector):
|
164
|
+
"""Generic particle collection that has Lorentz vector properties
|
165
|
+
|
166
|
+
Also handles the following additional branches:
|
167
|
+
- '{obj}_select'
|
168
|
+
"""
|
169
|
+
|
170
|
+
def passes(self, name):
|
171
|
+
return self[f"select_{name}"] == 1
|
172
|
+
|
173
|
+
# NB: fields with the name 'pt' take precedence over this
|
174
|
+
# @dask_property
|
175
|
+
# def pt(self):
|
176
|
+
# print('inside non-dask prop')
|
177
|
+
# return self["pt_NOSYS"]
|
178
|
+
|
179
|
+
# @pt.dask
|
180
|
+
# def pt(self, dask_array):
|
181
|
+
# branch = 'pt'
|
182
|
+
# print('inside dask prop')
|
183
|
+
# variation = dask_array._events().metadata.get("systematic", "NOSYS")
|
184
|
+
# with contextlib.suppress(Exception):
|
185
|
+
# return dask_array[f"{branch}_{variation}"]
|
186
|
+
|
187
|
+
# if variation != "NOSYS":
|
188
|
+
# with contextlib.suppress(Exception):
|
189
|
+
# return dask_array[f"{branch}_NOSYS"]
|
190
|
+
|
191
|
+
# return dask_array[branch]
|
192
|
+
|
193
|
+
|
194
|
+
_set_repr_name("Particle")
|
195
|
+
|
196
|
+
ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
197
|
+
ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
198
|
+
ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 # pylint: disable=undefined-variable
|
199
|
+
ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
200
|
+
|
201
|
+
|
202
|
+
behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior)) # pylint: disable=protected-access
|
203
|
+
|
204
|
+
|
205
|
+
@awkward.mixin_class(behavior)
|
206
|
+
class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic):
|
207
|
+
"""Missing transverse energy collection."""
|
208
|
+
|
209
|
+
@property
|
210
|
+
def r(self):
|
211
|
+
"""Distance from origin in XY plane"""
|
212
|
+
return self["met"]
|
213
|
+
|
214
|
+
|
215
|
+
_set_repr_name("MissingET")
|
216
|
+
|
217
|
+
MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821 # pylint: disable=undefined-variable
|
218
|
+
MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
219
|
+
MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
220
|
+
MissingETArray.MomentumClass = MissingETArray # noqa: F821 # pylint: disable=undefined-variable
|
221
|
+
|
222
|
+
behavior.update(awkward._util.copy_behaviors("Particle", "Photon", behavior)) # pylint: disable=protected-access
|
223
|
+
|
224
|
+
|
225
|
+
@awkward.mixin_class(behavior)
|
226
|
+
class Photon(Particle, base.NanoCollection, base.Systematic):
|
227
|
+
"""Photon particle collection."""
|
228
|
+
|
229
|
+
@property
|
230
|
+
def mass(self):
|
231
|
+
"""Return zero mass for photon."""
|
232
|
+
return awkward.zeros_like(self.pt)
|
233
|
+
|
234
|
+
@property
|
235
|
+
def charge(self):
|
236
|
+
"""Return zero charge for photon."""
|
237
|
+
return awkward.zeros_like(self.pt)
|
238
|
+
|
239
|
+
@property
|
240
|
+
def isEM(self):
|
241
|
+
return self.isEM_syst.NOSYS == 0 # pylint: disable=no-member
|
242
|
+
|
243
|
+
def pass_isEM(self, words: list[PhotonID]):
|
244
|
+
# 0 is pass, 1 is fail
|
245
|
+
return (
|
246
|
+
self.isEM_syst.NOSYS & reduce(ior, (1 << word.value for word in words)) # pylint: disable=no-member
|
247
|
+
) == 0
|
248
|
+
|
249
|
+
|
250
|
+
_set_repr_name("Photon")
|
251
|
+
|
252
|
+
PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
253
|
+
PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
254
|
+
PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 # pylint: disable=undefined-variable
|
255
|
+
PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
256
|
+
|
257
|
+
behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) # pylint: disable=protected-access
|
258
|
+
|
259
|
+
|
260
|
+
@awkward.mixin_class(behavior)
|
261
|
+
class Electron(Particle, base.NanoCollection, base.Systematic):
|
262
|
+
"""Electron particle collection."""
|
263
|
+
|
264
|
+
@property
|
265
|
+
def mass(self):
|
266
|
+
"""Electron mass in MeV"""
|
267
|
+
return awkward.ones_like(self.pt) * particle.literals.e_minus.mass # pylint: disable=no-member
|
268
|
+
|
269
|
+
|
270
|
+
_set_repr_name("Electron")
|
271
|
+
|
272
|
+
ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
273
|
+
ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
274
|
+
ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 # pylint: disable=undefined-variable
|
275
|
+
ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
276
|
+
|
277
|
+
behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) # pylint: disable=protected-access
|
278
|
+
|
279
|
+
|
280
|
+
@awkward.mixin_class(behavior)
|
281
|
+
class Muon(Particle, base.NanoCollection, base.Systematic):
|
282
|
+
"""Muon particle collection."""
|
283
|
+
|
284
|
+
@property
|
285
|
+
def mass(self):
|
286
|
+
"""Muon mass in MeV"""
|
287
|
+
return awkward.ones_like(self.pt) * particle.literals.mu_minus.mass # pylint: disable=no-member
|
288
|
+
|
289
|
+
|
290
|
+
_set_repr_name("Muon")
|
291
|
+
|
292
|
+
MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
293
|
+
MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
294
|
+
MuonArray.ProjectionClass4D = MuonArray # noqa: F821 # pylint: disable=undefined-variable
|
295
|
+
MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
296
|
+
|
297
|
+
behavior.update(awkward._util.copy_behaviors("Particle", "Tau", behavior)) # pylint: disable=protected-access
|
298
|
+
|
299
|
+
|
300
|
+
@awkward.mixin_class(behavior)
|
301
|
+
class Tau(Particle, base.NanoCollection, base.Systematic):
|
302
|
+
"""Tau particle collection."""
|
303
|
+
|
304
|
+
@property
|
305
|
+
def mass(self):
|
306
|
+
"""Tau mass in MeV"""
|
307
|
+
return awkward.ones_like(self.pt) * particle.literals.tau_minus.mass # pylint: disable=no-member
|
308
|
+
|
309
|
+
|
310
|
+
_set_repr_name("Tau")
|
311
|
+
|
312
|
+
TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
313
|
+
TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
314
|
+
TauArray.ProjectionClass4D = TauArray # noqa: F821 # pylint: disable=undefined-variable
|
315
|
+
TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
316
|
+
|
317
|
+
|
318
|
+
behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior)) # pylint: disable=protected-access
|
319
|
+
|
320
|
+
|
321
|
+
@awkward.mixin_class(behavior)
|
322
|
+
class Jet(Particle, base.NanoCollection, base.Systematic):
|
323
|
+
"""Jet particle collection."""
|
324
|
+
|
325
|
+
@property
|
326
|
+
def mass(self):
|
327
|
+
r"""Invariant mass (+, -, -, -)
|
328
|
+
|
329
|
+
:math:`\sqrt{t^2-x^2-y^2-z^2}`
|
330
|
+
"""
|
331
|
+
return self["m"]
|
332
|
+
|
333
|
+
|
334
|
+
_set_repr_name("Jet")
|
335
|
+
|
336
|
+
JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
337
|
+
JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
338
|
+
JetArray.ProjectionClass4D = JetArray # noqa: F821 # pylint: disable=undefined-variable
|
339
|
+
JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 # pylint: disable=undefined-variable,no-member
|
340
|
+
|
341
|
+
__all__ = [
|
342
|
+
"Electron",
|
343
|
+
"ElectronArray", # noqa: F822 # pylint: disable=undefined-all-variable
|
344
|
+
"ElectronRecord", # noqa: F822 # pylint: disable=undefined-all-variable
|
345
|
+
"Jet",
|
346
|
+
"JetArray", # noqa: F822 # pylint: disable=undefined-all-variable
|
347
|
+
"JetRecord", # noqa: F822 # pylint: disable=undefined-all-variable
|
348
|
+
"MissingET",
|
349
|
+
"MissingETArray", # noqa: F822 # pylint: disable=undefined-all-variable
|
350
|
+
"MissingETRecord", # noqa: F822 # pylint: disable=undefined-all-variable
|
351
|
+
"Muon",
|
352
|
+
"MuonArray", # noqa: F822 # pylint: disable=undefined-all-variable
|
353
|
+
"MuonRecord", # noqa: F822 # pylint: disable=undefined-all-variable
|
354
|
+
"NtupleEvents",
|
355
|
+
"Particle",
|
356
|
+
"ParticleArray", # noqa: F822 # pylint: disable=undefined-all-variable
|
357
|
+
"ParticleRecord", # noqa: F822 # pylint: disable=undefined-all-variable
|
358
|
+
"Pass",
|
359
|
+
"Photon",
|
360
|
+
"PhotonArray", # noqa: F822 # pylint: disable=undefined-all-variable
|
361
|
+
"PhotonRecord", # noqa: F822 # pylint: disable=undefined-all-variable
|
362
|
+
"Weight",
|
363
|
+
]
|