atlas-schema 0.2.3__tar.gz → 0.2.4__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.
@@ -1,11 +1,11 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: atlas-schema
3
- Version: 0.2.3
3
+ Version: 0.2.4
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.2.3/
8
+ Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.2.4/
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>
@@ -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.2.3
254
+ # atlas-schema v0.2.4
255
255
 
256
256
  [![Actions Status][actions-badge]][actions-link]
257
257
  [![Documentation Status][rtd-badge]][rtd-link]
@@ -1,4 +1,4 @@
1
- # atlas-schema v0.2.3
1
+ # atlas-schema v0.2.4
2
2
 
3
3
  [![Actions Status][actions-badge]][actions-link]
4
4
  [![Documentation Status][rtd-badge]][rtd-link]
@@ -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.2.3/"
63
+ Documentation = "https://atlas-schema.readthedocs.io/en/v0.2.4/"
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
 
@@ -111,14 +111,18 @@ addopts = [
111
111
  ]
112
112
  xfail_strict = true
113
113
  filterwarnings = [
114
- "error",
114
+ "error",
115
115
  ]
116
+
116
117
  log_cli_level = "INFO"
117
118
  testpaths = [
118
119
  "src",
119
120
  "tests",
120
121
  "docs",
121
122
  ]
123
+ norecursedirs = [
124
+ "tests/helpers"
125
+ ]
122
126
 
123
127
  [tool.coverage]
124
128
  run.source = ["atlas_schema"]
@@ -12,5 +12,5 @@ __version__: str
12
12
  __version_tuple__: VERSION_TUPLE
13
13
  version_tuple: VERSION_TUPLE
14
14
 
15
- __version__ = version = '0.2.3'
16
- __version_tuple__ = version_tuple = (0, 2, 3)
15
+ __version__ = version = '0.2.4'
16
+ __version_tuple__ = version_tuple = (0, 2, 4)
@@ -230,12 +230,24 @@ JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
230
230
 
231
231
  __all__ = [
232
232
  "Electron",
233
+ "ElectronArray", # noqa: F822
234
+ "ElectronRecord", # noqa: F822
233
235
  "Jet",
236
+ "JetArray", # noqa: F822
237
+ "JetRecord", # noqa: F822
234
238
  "MissingET",
239
+ "MissingETArray", # noqa: F822
240
+ "MissingETRecord", # noqa: F822
235
241
  "Muon",
242
+ "MuonArray", # noqa: F822
243
+ "MuonRecord", # noqa: F822
236
244
  "NtupleEvents",
237
245
  "Particle",
246
+ "ParticleArray", # noqa: F822
247
+ "ParticleRecord", # noqa: F822
238
248
  "Pass",
239
249
  "Photon",
250
+ "PhotonArray", # noqa: F822
251
+ "PhotonRecord", # noqa: F822
240
252
  "Weight",
241
253
  ]
@@ -0,0 +1,405 @@
1
+ from __future__ import annotations
2
+
3
+ import difflib
4
+ import warnings
5
+ from collections.abc import KeysView, ValuesView
6
+ from typing import Any, ClassVar
7
+
8
+ from coffea.nanoevents.schemas.base import BaseSchema, zip_forms
9
+
10
+ from atlas_schema.typing_compat import Behavior, Self
11
+
12
+
13
+ class NtupleSchema(BaseSchema): # type: ignore[misc]
14
+ """The schema for building ATLAS ntuples following the typical centralized formats.
15
+
16
+ This schema is built from all branches found in a tree in the supplied
17
+ file, based on the naming pattern of the branches. This naming pattern is
18
+ typically assumed to be
19
+
20
+ .. code-block:: bash
21
+
22
+ {collection:str}_{subcollection:str}_{systematic:str}
23
+
24
+ where:
25
+ * ``collection`` is assumed to be a prefix with typical characters, following the regex ``[a-zA-Z][a-zA-Z0-9]*``; that is starting with a case-insensitive letter, and proceeded by zero or more alphanumeric characters,
26
+ * ``subcollection`` is assumed to be anything with typical characters (allowing for underscores) following the regex ``[a-zA-Z_][a-zA-Z0-9_]*``; that is starting with a case-insensitive letter or underscore, and proceeded by zero or more alphanumeric characters including underscores, and
27
+ * ``systematic`` is assumed to be either ``NOSYS`` to indicate a branch with potential systematic variariations, or anything with typical characters (allowing for underscores) following the same regular expression as the ``subcollection``.
28
+
29
+ Here, a collection refers to the top-level entry to access an item - a collection called ``el`` will be accessible under the ``el`` attributes via ``events['el']`` or ``events.el``. A subcollection called ``pt`` will be accessible under that collection, such as ``events['el']['pt']`` or ``events.el.pt``. This is the power of the schema providing a more user-friendly (and programmatic) access to the underlying branches.
30
+
31
+ The above logic means that the following branches below will be categorized as follows:
32
+
33
+ +-------------------------------+-------------------+-----------------------+------------------+
34
+ | branch | collection | subcollection | systematic |
35
+ +===============================+===================+=======================+==================+
36
+ | ``'eventNumber'`` | ``'eventNumber'`` | ``None`` | ``None`` |
37
+ +-------------------------------+-------------------+-----------------------+------------------+
38
+ | ``'runNumber'`` | ``'runNumber'`` | ``None`` | ``None`` |
39
+ +-------------------------------+-------------------+-----------------------+------------------+
40
+ | ``'el_pt_NOSYS'`` | ``'el'`` | ``'pt'`` | ``'NOSYS'`` |
41
+ +-------------------------------+-------------------+-----------------------+------------------+
42
+ | ``'jet_cleanTightBad_NOSYS'`` | ``'jet'`` | ``'cleanTightBad'`` | ``'NOSYS'`` |
43
+ +-------------------------------+-------------------+-----------------------+------------------+
44
+ | ``'jet_select_btag_NOSYS'`` | ``'jet'`` | ``'select_btag'`` | ``'NOSYS'`` |
45
+ +-------------------------------+-------------------+-----------------------+------------------+
46
+ | ``'jet_e_NOSYS'`` | ``'jet'`` | ``'e'`` | ``'NOSYS'`` |
47
+ +-------------------------------+-------------------+-----------------------+------------------+
48
+ | ``'truthel_phi'`` | ``'truthel'`` | ``'phi'`` | ``None`` |
49
+ +-------------------------------+-------------------+-----------------------+------------------+
50
+ | ``'truthel_pt'`` | ``'truthel'`` | ``'pt'`` | ``None`` |
51
+ +-------------------------------+-------------------+-----------------------+------------------+
52
+ | ``'ph_eta'`` | ``'ph'`` | ``'eta'`` | ``None`` |
53
+ +-------------------------------+-------------------+-----------------------+------------------+
54
+ | ``'ph_phi_SCALE__1up'`` | ``'ph'`` | ``'phi'`` | ``'SCALE__1up'`` |
55
+ +-------------------------------+-------------------+-----------------------+------------------+
56
+ | ``'mu_TTVA_effSF_NOSYS'`` | ``'mu'`` | ``'TTVA_effSF'`` | ``'NOSYS'`` |
57
+ +-------------------------------+-------------------+-----------------------+------------------+
58
+ | ``'recojet_antikt4PFlow_pt'`` | ``'recojet'`` | ``'antikt4PFlow_pt'`` | ``'NOSYS'`` |
59
+ +-------------------------------+-------------------+-----------------------+------------------+
60
+ | ``'recojet_antikt10UFO_m'`` | ``'recojet'`` | ``'antikt10UFO_m'`` | ``None`` |
61
+ +-------------------------------+-------------------+-----------------------+------------------+
62
+
63
+ Sometimes this logic is not what you want, and there are ways to teach ``NtupleSchema`` how to group some of these better for atypical cases. We can address these case-by-case.
64
+
65
+ **Singletons**
66
+
67
+ Sometimes you have particular branches that you don't want to be treated as a collection (with subcollections). And sometimes you will see warnings about this (see :ref:`faq`). There are some pre-defined ``singletons`` stored under :attr:`event_ids`, and these will be lazily treated as a _singleton_. For other cases where you add your own branches, you can additionally extend this class to add your own :attr:`singletons`:
68
+
69
+ .. code-block:: python
70
+
71
+ from atlas_schema.schema import NtupleSchema
72
+
73
+
74
+ class MySchema(NtupleSchema):
75
+ singletons = {"RandomRunNumber"}
76
+
77
+ and use this schema in your analysis code. The rest of the logic will be handled for you, and you can access your singletons under ``events.RandomRunNumber`` as expected.
78
+
79
+ **Mixins (collections, subcollections)**
80
+
81
+ In more complicated scenarios, you might need to teach :class:`NtupleSchema` how to handle collections that end up having underscores in their name, or other characters that make the grouping non-trivial. In some other scenarios, you want to tell the schema to assign a certain set of behaviors to a collection - rather than the default :class:`atlas_schema.methods.Particle` behavior. This is where :attr:`mixins` comes in. Similar to how :attr:`singletons` are handled, you extend this schema to include your own ``mixins`` pointing them at one of the behaviors defined in :mod:`atlas_schema.methods`.
82
+
83
+ Let's demonstrate both cases. Imagine you want to have your ``truthel`` collections above treated as :class:`atlas_schema.methods.Electron`, then you would extend the existing :attr:`mixins`:
84
+
85
+ .. code-block:: python
86
+
87
+ from atlas_schema.schema import NtupleSchema
88
+
89
+
90
+ class MySchema(NtupleSchema):
91
+ mixins = {"truthel": "Electron", **NtupleSchema.mixins}
92
+
93
+ Now, ``events.truthel`` will give you arrays zipped up with :class:`atlas_schema.methods.Electron` behaviors.
94
+
95
+ If instead, you run into problems with mixing different branches in the same collection, because the default behavior of this schema described above is not smart enough to handle the atypical cases, you can explicitly fix this by defining your collections:
96
+
97
+ .. code-block:: python
98
+
99
+ from atlas_schema.schema import NtupleSchema
100
+
101
+
102
+ class MySchema(NtupleSchema):
103
+ mixins = {
104
+ "recojet_antikt4PFlow": "Jet",
105
+ "recojet_antikt10UFO": "Jet",
106
+ **NtupleSchema.mixins,
107
+ }
108
+
109
+ Now, ``events.recojet_antikt4PFlow`` and ``events.recojet_antikt10UFO`` will be separate collections, instead of a single ``events.recojet`` that incorrectly merged branches from each of these collections.
110
+ """
111
+
112
+ __dask_capable__: ClassVar[bool] = True
113
+
114
+ warn_missing_crossrefs: ClassVar[bool] = True
115
+
116
+ #: Treat missing event-level branches as error instead of warning (default is ``False``)
117
+ error_missing_event_ids: ClassVar[bool] = False
118
+ #: Determine closest behavior for a given branch or treat branch as :attr:`default_behavior` (default is ``True``)
119
+ identify_closest_behavior: ClassVar[bool] = True
120
+
121
+ #: event IDs to expect in data datasets
122
+ event_ids_data: ClassVar[set[str]] = {
123
+ "lumiBlock",
124
+ "averageInteractionsPerCrossing",
125
+ "actualInteractionsPerCrossing",
126
+ "dataTakingYear",
127
+ }
128
+ #: event IDs to expect in MC datasets
129
+ event_ids_mc: ClassVar[set[str]] = {
130
+ "mcChannelNumber",
131
+ "runNumber",
132
+ "eventNumber",
133
+ "mcEventWeights",
134
+ }
135
+ #: all event IDs to expect in the dataset
136
+ event_ids: ClassVar[set[str]] = {*event_ids_data, *event_ids_mc}
137
+
138
+ #: mixins defining the mapping from collection name to behavior to use for that collection
139
+ mixins: ClassVar[dict[str, str]] = {
140
+ "el": "Electron",
141
+ "jet": "Jet",
142
+ "met": "MissingET",
143
+ "mu": "Muon",
144
+ "pass": "Pass",
145
+ "ph": "Photon",
146
+ "trigPassed": "Trigger",
147
+ "weight": "Weight",
148
+ }
149
+
150
+ #: additional branches to pass-through with no zipping or additional interpretation (such as those stored as length-1 vectors)
151
+ singletons: ClassVar[set[str]] = set()
152
+
153
+ #: docstrings to assign for specific subcollections across the various collections identified by this schema
154
+ docstrings: ClassVar[dict[str, str]] = {
155
+ "charge": "charge",
156
+ "eta": "pseudorapidity",
157
+ "met": "missing transverse energy [MeV]",
158
+ "mass": "invariant mass [MeV]",
159
+ "pt": "transverse momentum [MeV]",
160
+ "phi": "azimuthal angle",
161
+ }
162
+
163
+ #: default behavior to use for any collection (default ``"NanoCollection"``, from :class:`coffea.nanoevents.methods.base.NanoCollection`)
164
+ default_behavior: ClassVar[str] = "NanoCollection"
165
+
166
+ def __init__(self, base_form: dict[str, Any], version: str = "latest"):
167
+ super().__init__(base_form)
168
+ self._version = version
169
+ if version == "latest":
170
+ pass
171
+ else:
172
+ pass
173
+ self._form["fields"], self._form["contents"] = self._build_collections(
174
+ self._form["fields"], self._form["contents"]
175
+ )
176
+ self._form["parameters"]["metadata"]["version"] = self._version
177
+
178
+ @classmethod
179
+ def v1(cls, base_form: dict[str, Any]) -> Self:
180
+ """Build the NtupleEvents
181
+
182
+ For example, one can use ``NanoEventsFactory.from_root("file.root", schemaclass=NtupleSchema.v1)``
183
+ to ensure NanoAODv7 compatibility.
184
+ """
185
+ return cls(base_form, version="1")
186
+
187
+ def _build_collections(
188
+ self, field_names: list[str], input_contents: list[Any]
189
+ ) -> tuple[KeysView[str], ValuesView[dict[str, Any]]]:
190
+ branch_forms = dict(zip(field_names, input_contents))
191
+
192
+ # parse into high-level records (collections, list collections, and singletons)
193
+ collections = {k.split("_")[0] for k in branch_forms}
194
+ collections -= self.event_ids
195
+ collections -= set(self.singletons)
196
+
197
+ # now handle any collections that we identified that are substrings of the items in the mixins
198
+ # convert all valid branch_forms into strings to make the lookups a bit faster
199
+ bf_str = ",".join(branch_forms.keys())
200
+ for mixin in self.mixins:
201
+ if mixin in collections:
202
+ continue
203
+ if f",{mixin}_" not in bf_str and not bf_str.startswith(f"{mixin}_"):
204
+ continue
205
+ if "_" in mixin:
206
+ warnings.warn(
207
+ f"I identified a mixin that I did not automatically identify as a collection because it contained an underscore: '{mixin}'. I will add this to the known collections. To suppress this warning next time, please create your ntuples with collections without underscores. [mixin-underscore]",
208
+ RuntimeWarning,
209
+ stacklevel=2,
210
+ )
211
+ collections.add(mixin)
212
+ for collection in list(collections):
213
+ if mixin.startswith(f"{collection}_"):
214
+ warnings.warn(
215
+ f"I found a misidentified collection: '{collection}'. I will remove this from the known collections. To suppress this warning next time, please create your ntuples with collections that are not similarly named with underscores. [collection-subset]",
216
+ RuntimeWarning,
217
+ stacklevel=2,
218
+ )
219
+ collections.remove(collection)
220
+ break
221
+
222
+ # rename needed because easyjet breaks the AMG assumptions
223
+ # https://gitlab.cern.ch/easyjet/easyjet/-/issues/246
224
+ for k in list(branch_forms):
225
+ if "NOSYS" not in k:
226
+ continue
227
+ branch_forms[k.replace("_NOSYS", "") + "_NOSYS"] = branch_forms.pop(k)
228
+
229
+ # these are collections with systematic variations
230
+ subcollections = {
231
+ k.split("__")[0].split("_", 1)[1].replace("_NOSYS", "")
232
+ for k in branch_forms
233
+ if "NOSYS" in k
234
+ }
235
+
236
+ # Check the presence of the event_ids
237
+ missing_event_ids = [
238
+ event_id for event_id in self.event_ids if event_id not in branch_forms
239
+ ]
240
+
241
+ if len(missing_event_ids) > 0:
242
+ if self.error_missing_event_ids:
243
+ msg = f"There are missing event ID fields: {missing_event_ids} \n\n\
244
+ The event ID fields {self.event_ids} are necessary to perform sub-run identification \
245
+ (e.g. for corrections and sub-dividing data during different detector conditions),\
246
+ to cross-validate MC and Data (i.e. matching events for comparison), and to generate event displays. \
247
+ It's advised to never drop these branches from the dataformat.\n\n\
248
+ This error can be demoted to a warning by setting the class level variable error_missing_event_ids to False."
249
+ raise RuntimeError(msg)
250
+ warnings.warn(
251
+ f"Missing event_ids : {missing_event_ids}",
252
+ RuntimeWarning,
253
+ stacklevel=2,
254
+ )
255
+
256
+ output = {}
257
+
258
+ # first, register singletons (event-level, others)
259
+ for name in {*self.event_ids, *self.singletons}:
260
+ if name in missing_event_ids:
261
+ continue
262
+ output[name] = branch_forms[name]
263
+
264
+ # next, go through and start grouping up collections
265
+ for name in collections:
266
+ content = {}
267
+ used = set()
268
+
269
+ for subname in subcollections:
270
+ prefix = f"{name}_{subname}_"
271
+ used.update({k for k in branch_forms if k.startswith(prefix)})
272
+ subcontent = {
273
+ k[len(prefix) :]: branch_forms[k]
274
+ for k in branch_forms
275
+ if k.startswith(prefix)
276
+ }
277
+ if subcontent:
278
+ # create the nominal version
279
+ content[subname] = branch_forms[f"{prefix}NOSYS"]
280
+ # create a collection of the systematic variations for the given variable
281
+ content[f"{subname}_syst"] = zip_forms(
282
+ subcontent, f"{name}_syst", record_name="NanoCollection"
283
+ )
284
+
285
+ content.update(
286
+ {
287
+ k[len(name) + 1 :]: branch_forms[k]
288
+ for k in branch_forms
289
+ if k.startswith(name + "_") and k not in used
290
+ }
291
+ )
292
+
293
+ if not used and not content:
294
+ warnings.warn(
295
+ f"I identified a branch that likely does not have any leaves: '{name}'. I will treat this as a 'singleton'. To suppress this warning next time, please define your singletons explicitly. [singleton-undefined]",
296
+ RuntimeWarning,
297
+ stacklevel=2,
298
+ )
299
+ self.singletons.add(name)
300
+ output[name] = branch_forms[name]
301
+
302
+ else:
303
+ behavior = self.mixins.get(name, "")
304
+ if not behavior:
305
+ behavior = self.suggested_behavior(name)
306
+ warnings.warn(
307
+ f"I found a collection with no defined mixin: '{name}'. I will assume behavior: '{behavior}'. To suppress this warning next time, please define mixins for your custom collections. [mixin-undefined]",
308
+ RuntimeWarning,
309
+ stacklevel=2,
310
+ )
311
+
312
+ output[name] = zip_forms(content, name, record_name=behavior)
313
+
314
+ output[name].setdefault("parameters", {})
315
+ output[name]["parameters"].update({"collection_name": name})
316
+
317
+ if output[name]["class"] == "ListOffsetArray":
318
+ if output[name]["class"] == "RecordArray":
319
+ parameters = output[name]["content"]["fields"]
320
+ contents = output[name]["content"]["contents"]
321
+ else:
322
+ # these are also singletons of another kind that we just pass through
323
+ continue
324
+ elif output[name]["class"] == "RecordArray":
325
+ parameters = output[name]["fields"]
326
+ contents = output[name]["contents"]
327
+ elif output[name]["class"] == "NumpyArray":
328
+ # these are singletons that we just pass through
329
+ continue
330
+ else:
331
+ msg = f"Unhandled class {output[name]['class']}"
332
+ raise RuntimeError(msg)
333
+
334
+ # update docstrings as needed
335
+ # NB: must be before flattening for easier logic
336
+ for index, parameter in enumerate(parameters):
337
+ if "parameters" not in contents[index]:
338
+ continue
339
+
340
+ parsed_name = parameter.replace("_NOSYS", "")
341
+ contents[index]["parameters"]["__doc__"] = self.docstrings.get(
342
+ parsed_name,
343
+ contents[index]["parameters"].get(
344
+ "__doc__", "no docstring available"
345
+ ),
346
+ )
347
+
348
+ return output.keys(), output.values()
349
+
350
+ @classmethod
351
+ def behavior(cls) -> Behavior:
352
+ """Behaviors necessary to implement this schema
353
+
354
+ Returns:
355
+ dict[str | tuple['*', str], type[awkward.Record]]: an :data:`awkward.behavior` dictionary
356
+ """
357
+ from atlas_schema.methods import behavior as roaster
358
+
359
+ return roaster
360
+
361
+ @classmethod
362
+ def suggested_behavior(cls, key: str, cutoff: float = 0.4) -> str:
363
+ """
364
+ Suggest e behavior to use for a provided collection or branch name.
365
+
366
+ Default behavior: :class:`~coffea.nanoevents.methods.base.NanoCollection`.
367
+
368
+ Note:
369
+ If :attr:`identify_closest_behavior` is ``False``, then this function will return the default behavior ``NanoCollection``.
370
+
371
+ Warning:
372
+ If no behavior is found above the *cutoff* score, then this function will return the default behavior.
373
+
374
+ Args:
375
+ key (str): collection name to suggest a matching behavior for
376
+ cutoff (float): o ptional argument cutoff (default ``0.4``) is a float in the range ``[0, 1]``. Possibilities that don't score at least that similar to *key* are ignored.
377
+
378
+ Returns:
379
+ str: suggested behavior to use by string
380
+
381
+ Example:
382
+ >>> from atlas_schema.schema import NtupleSchema
383
+ >>> NtupleSchema.suggested_behavior("truthjet")
384
+ 'Jet'
385
+ >>> NtupleSchema.suggested_behavior("SignalElectron")
386
+ 'Electron'
387
+ >>> NtupleSchema.suggested_behavior("generatorWeight")
388
+ 'Weight'
389
+ >>> NtupleSchema.suggested_behavior("aVeryStrangelyNamedBranchWithNoMatch")
390
+ 'NanoCollection'
391
+ """
392
+ if cls.identify_closest_behavior:
393
+ # lowercase everything to do case-insensitive matching
394
+ behaviors = [b for b in cls.behavior() if isinstance(b, str)]
395
+ behaviors_l = [b.lower() for b in behaviors]
396
+ results = difflib.get_close_matches(
397
+ key.lower(), behaviors_l, n=1, cutoff=cutoff
398
+ )
399
+ if not results:
400
+ return cls.default_behavior
401
+
402
+ behavior = results[0]
403
+ # need to identify the index and return the unlowered version
404
+ return behaviors[behaviors_l.index(behavior)]
405
+ return cls.default_behavior
@@ -5,7 +5,7 @@ Typing helpers.
5
5
  from __future__ import annotations
6
6
 
7
7
  import sys
8
- from typing import Annotated
8
+ from typing import Annotated, Literal, Union
9
9
 
10
10
  import awkward
11
11
 
@@ -19,6 +19,6 @@ if sys.version_info >= (3, 11):
19
19
  else:
20
20
  from typing_extensions import Self
21
21
 
22
- Behavior: TypeAlias = dict[str, type[awkward.Record]]
22
+ Behavior: TypeAlias = dict[Union[str, tuple[Literal["*"]], str], type[awkward.Record]]
23
23
 
24
24
  __all__ = ("Annotated", "Behavior", "Self")
@@ -21,12 +21,12 @@ def isin(element: Array, test_elements: dak.Array | ak.Array, axis: int = -1) ->
21
21
  comparison.
22
22
 
23
23
  Args:
24
- element (dak.Array or ak.Array): input array of values.
25
- test_elements (dak.Array or ak.Array): one-dimensional set of values against which to test each value of *element*.
24
+ element (dask_awkward.Array or ak.Array): input array of values.
25
+ test_elements (dask_awkward.Array or ak.Array): one-dimensional set of values against which to test each value of *element*.
26
26
  axis (int): the axis along which the comparison is performed
27
27
 
28
28
  Returns:
29
- dak.Array or ak.Array: result of comparison for test_elements in *element*
29
+ dask_awkward.Array or ak.Array: result of comparison for test_elements in *element*
30
30
 
31
31
  Example:
32
32
  >>> import awkward as ak
@@ -1,214 +0,0 @@
1
- from __future__ import annotations
2
-
3
- import warnings
4
- from collections.abc import KeysView, ValuesView
5
- from typing import Any, ClassVar
6
-
7
- from coffea.nanoevents.schemas.base import BaseSchema, zip_forms
8
-
9
- from atlas_schema.typing_compat import Behavior, Self
10
-
11
-
12
- class NtupleSchema(BaseSchema): # type: ignore[misc]
13
- """Ntuple schema builder
14
-
15
- The Ntuple schema is built from all branches found in the supplied file, based on
16
- the naming pattern of the branches. The following additional arrays are constructed:
17
-
18
- - n/a
19
- """
20
-
21
- __dask_capable__ = True
22
-
23
- warn_missing_crossrefs = True
24
- error_missing_event_ids = False
25
-
26
- event_ids_data: ClassVar[set[str]] = {
27
- "lumiBlock",
28
- "averageInteractionsPerCrossing",
29
- "actualInteractionsPerCrossing",
30
- "dataTakingYear",
31
- }
32
- event_ids_mc: ClassVar[set[str]] = {
33
- "mcChannelNumber",
34
- "runNumber",
35
- "eventNumber",
36
- "mcEventWeights",
37
- }
38
- event_ids: ClassVar[set[str]] = {*event_ids_data, *event_ids_mc}
39
-
40
- mixins: ClassVar[dict[str, str]] = {
41
- "el": "Electron",
42
- "jet": "Jet",
43
- "met": "MissingET",
44
- "mu": "Muon",
45
- "pass": "Pass",
46
- "ph": "Photon",
47
- "trigPassed": "Trigger",
48
- "weight": "Weight",
49
- }
50
-
51
- # These are stored as length-1 vectors unnecessarily
52
- singletons: ClassVar[set[str]] = set()
53
-
54
- docstrings: ClassVar[dict[str, str]] = {
55
- "charge": "charge",
56
- "eta": "pseudorapidity",
57
- "met": "missing transverse energy [MeV]",
58
- "mass": "invariant mass [MeV]",
59
- "pt": "transverse momentum [MeV]",
60
- "phi": "azimuthal angle",
61
- }
62
-
63
- def __init__(self, base_form: dict[str, Any], version: str = "latest"):
64
- super().__init__(base_form)
65
- self._version = version
66
- if version == "latest":
67
- pass
68
- else:
69
- pass
70
- self._form["fields"], self._form["contents"] = self._build_collections(
71
- self._form["fields"], self._form["contents"]
72
- )
73
- self._form["parameters"]["metadata"]["version"] = self._version
74
-
75
- @classmethod
76
- def v1(cls, base_form: dict[str, Any]) -> Self:
77
- """Build the NtupleEvents
78
-
79
- For example, one can use ``NanoEventsFactory.from_root("file.root", schemaclass=NtupleSchema.v1)``
80
- to ensure NanoAODv7 compatibility.
81
- """
82
- return cls(base_form, version="1")
83
-
84
- def _build_collections(
85
- self, field_names: list[str], input_contents: list[Any]
86
- ) -> tuple[KeysView[str], ValuesView[dict[str, Any]]]:
87
- branch_forms = dict(zip(field_names, input_contents))
88
-
89
- # parse into high-level records (collections, list collections, and singletons)
90
- collections = {k.split("_")[0] for k in branch_forms}
91
- collections -= self.event_ids
92
- collections -= set(self.singletons)
93
-
94
- # rename needed because easyjet breaks the AMG assumptions
95
- # https://gitlab.cern.ch/easyjet/easyjet/-/issues/246
96
- for k in list(branch_forms):
97
- if "NOSYS" not in k:
98
- continue
99
- branch_forms[k.replace("_NOSYS", "") + "_NOSYS"] = branch_forms.pop(k)
100
-
101
- # these are collections with systematic variations
102
- subcollections = {
103
- k.split("__")[0].split("_", 1)[1].replace("_NOSYS", "")
104
- for k in branch_forms
105
- if "NOSYS" in k
106
- }
107
-
108
- # Check the presence of the event_ids
109
- missing_event_ids = [
110
- event_id for event_id in self.event_ids if event_id not in branch_forms
111
- ]
112
-
113
- if len(missing_event_ids) > 0:
114
- if self.error_missing_event_ids:
115
- msg = f"There are missing event ID fields: {missing_event_ids} \n\n\
116
- The event ID fields {self.event_ids} are necessary to perform sub-run identification \
117
- (e.g. for corrections and sub-dividing data during different detector conditions),\
118
- to cross-validate MC and Data (i.e. matching events for comparison), and to generate event displays. \
119
- It's advised to never drop these branches from the dataformat.\n\n\
120
- This error can be demoted to a warning by setting the class level variable error_missing_event_ids to False."
121
- raise RuntimeError(msg)
122
- warnings.warn(
123
- f"Missing event_ids : {missing_event_ids}",
124
- RuntimeWarning,
125
- stacklevel=2,
126
- )
127
-
128
- output = {}
129
-
130
- # first, register singletons (event-level, others)
131
- for name in {*self.event_ids, *self.singletons}:
132
- if name in missing_event_ids:
133
- continue
134
- output[name] = branch_forms[name]
135
-
136
- # next, go through and start grouping up collections
137
- for name in collections:
138
- mixin = self.mixins.get(name, "NanoCollection")
139
- content = {}
140
- used = set()
141
-
142
- for subname in subcollections:
143
- prefix = f"{name}_{subname}_"
144
- used.update({k for k in branch_forms if k.startswith(prefix)})
145
- subcontent = {
146
- k[len(prefix) :]: branch_forms[k]
147
- for k in branch_forms
148
- if k.startswith(prefix)
149
- }
150
- if subcontent:
151
- # create the nominal version
152
- content[subname] = branch_forms[f"{prefix}NOSYS"]
153
- # create a collection of the systematic variations for the given variable
154
- content[f"{subname}_syst"] = zip_forms(
155
- subcontent, f"{name}_syst", record_name="NanoCollection"
156
- )
157
-
158
- content.update(
159
- {
160
- k[len(name) + 1 :]: branch_forms[k]
161
- for k in branch_forms
162
- if k.startswith(name + "_") and k not in used
163
- }
164
- )
165
-
166
- if not used and not content:
167
- warnings.warn(
168
- f"I identified a branch that likely does not have any leaves: '{name}'. I will treat this as a 'singleton'. To suppress this warning next time, please define your singletons explicitly.",
169
- RuntimeWarning,
170
- stacklevel=2,
171
- )
172
- self.singletons.add(name)
173
- output[name] = branch_forms[name]
174
-
175
- else:
176
- output[name] = zip_forms(content, name, record_name=mixin)
177
-
178
- output[name].setdefault("parameters", {})
179
- output[name]["parameters"].update({"collection_name": name})
180
-
181
- if output[name]["class"] == "ListOffsetArray":
182
- parameters = output[name]["content"]["fields"]
183
- contents = output[name]["content"]["contents"]
184
- elif output[name]["class"] == "RecordArray":
185
- parameters = output[name]["fields"]
186
- contents = output[name]["contents"]
187
- elif output[name]["class"] == "NumpyArray":
188
- # these are singletons that we just pass through
189
- continue
190
- else:
191
- msg = f"Unhandled class {output[name]['class']}"
192
- raise RuntimeError(msg)
193
- # update docstrings as needed
194
- # NB: must be before flattening for easier logic
195
- for index, parameter in enumerate(parameters):
196
- if "parameters" not in contents[index]:
197
- continue
198
-
199
- parsed_name = parameter.replace("_NOSYS", "")
200
- contents[index]["parameters"]["__doc__"] = self.docstrings.get(
201
- parsed_name,
202
- contents[index]["parameters"].get(
203
- "__doc__", "no docstring available"
204
- ),
205
- )
206
-
207
- return output.keys(), output.values()
208
-
209
- @classmethod
210
- def behavior(cls) -> Behavior:
211
- """Behaviors necessary to implement this schema"""
212
- from atlas_schema.methods import behavior as roaster
213
-
214
- return roaster
File without changes
File without changes