atlas-schema 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

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/__init__.py CHANGED
@@ -6,7 +6,12 @@ atlas_schema: Collection of utilities and helper functions for HEP ATLAS analyse
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
+ import warnings
10
+
9
11
  from atlas_schema._version import version as __version__
10
12
  from atlas_schema.enums import ParticleOrigin, PhotonID
13
+ from atlas_schema.utils import isin
14
+
15
+ warnings.filterwarnings("ignore", module="coffea.*")
11
16
 
12
- __all__ = ["__version__", "ParticleOrigin", "PhotonID"]
17
+ __all__ = ["ParticleOrigin", "PhotonID", "__version__", "isin"]
atlas_schema/_version.py CHANGED
@@ -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.0'
16
- __version_tuple__ = version_tuple = (0, 2, 0)
15
+ __version__ = version = '0.2.2'
16
+ __version_tuple__ = version_tuple = (0, 2, 2)
atlas_schema/enums.py CHANGED
@@ -1,116 +1,145 @@
1
1
  from __future__ import annotations
2
2
 
3
- from enum import IntEnum
4
-
5
- from atlas_schema.typing_compat import Annotated
6
-
7
-
8
- # https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103
9
- class ParticleType(IntEnum):
10
- Unknown: Annotated[int, "Unknown"] = 0
11
- UnknownElectron: Annotated[int, "UnknownElectron"] = 1
12
- IsoElectron: Annotated[int, "IsoElectron"] = 2
13
- NonIsoElectron: Annotated[int, "NonIsoElectron"] = 3
14
- BkgElectron: Annotated[int, "BkgElectron"] = 4
15
- UnknownMuon: Annotated[int, "UnknownMuon"] = 5
16
- IsoMuon: Annotated[int, "IsoMuon"] = 6
17
- NonIsoMuon: Annotated[int, "NonIsoMuon"] = 7
18
- BkgMuon: Annotated[int, "BkgMuon"] = 8
19
- UnknownTau: Annotated[int, "UnknownTau"] = 9
20
- IsoTau: Annotated[int, "IsoTau"] = 10
21
- NonIsoTau: Annotated[int, "NonIsoTau"] = 11
22
- BkgTau: Annotated[int, "BkgTau"] = 12
23
- UnknownPhoton: Annotated[int, "UnknownPhoton"] = 13
24
- IsoPhoton: Annotated[int, "IsoPhoton"] = 14
25
- NonIsoPhoton: Annotated[int, "NonIsoPhoton"] = 15
26
- BkgPhoton: Annotated[int, "BkgPhoton"] = 16
27
- Hadron: Annotated[int, "Hadron"] = 17
28
- Neutrino: Annotated[int, "Neutrino"] = 18
29
- NuclFrag: Annotated[int, "NuclFrag"] = 19
30
- NonPrimary: Annotated[int, "NonPrimary"] = 20
31
- GenParticle: Annotated[int, "GenParticle"] = 21
32
- SUSYParticle: Annotated[int, "SUSYParticle"] = 22
33
- OtherBSMParticle: Annotated[int, "OtherBSMParticle"] = 39
34
- BBbarMesonPart: Annotated[int, "BBbarMesonPart"] = 23
35
- BottomMesonPart: Annotated[int, "BottomMesonPart"] = 24
36
- CCbarMesonPart: Annotated[int, "CCbarMesonPart"] = 25
37
- CharmedMesonPart: Annotated[int, "CharmedMesonPart"] = 26
38
- BottomBaryonPart: Annotated[int, "BottomBaryonPart"] = 27
39
- CharmedBaryonPart: Annotated[int, "CharmedBaryonPart"] = 28
40
- StrangeBaryonPart: Annotated[int, "StrangeBaryonPart"] = 29
41
- LightBaryonPart: Annotated[int, "LightBaryonPart"] = 30
42
- StrangeMesonPart: Annotated[int, "StrangeMesonPart"] = 31
43
- LightMesonPart: Annotated[int, "LightMesonPart"] = 32
44
- BJet: Annotated[int, "BJet"] = 33
45
- CJet: Annotated[int, "CJet"] = 34
46
- LJet: Annotated[int, "LJet"] = 35
47
- GJet: Annotated[int, "GJet"] = 36
48
- TauJet: Annotated[int, "TauJet"] = 37
49
- UnknownJet: Annotated[int, "UnknownJet"] = 38
50
-
51
-
52
- # https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103
53
- class ParticleOrigin(IntEnum):
54
- NonDefined: Annotated[int, "NonDefined"] = 0
55
- SingleElec: Annotated[int, "SingleElec"] = 1
56
- SingleMuon: Annotated[int, "SingleMuon"] = 2
57
- SinglePhot: Annotated[int, "SinglePhot"] = 3
58
- SingleTau: Annotated[int, "SingleTau"] = 4
59
- PhotonConv: Annotated[int, "PhotonConv"] = 5
60
- DalitzDec: Annotated[int, "DalitzDec"] = 6
61
- ElMagProc: Annotated[int, "ElMagProc"] = 7
62
- Mu: Annotated[int, "Mu"] = 8
63
- TauLep: Annotated[int, "TauLep"] = 9
64
- top: Annotated[int, "top"] = 10
65
- QuarkWeakDec: Annotated[int, "QuarkWeakDec"] = 11
66
- WBoson: Annotated[int, "WBoson"] = 12
67
- ZBoson: Annotated[int, "ZBoson"] = 13
68
- Higgs: Annotated[int, "Higgs"] = 14
69
- HiggsMSSM: Annotated[int, "HiggsMSSM"] = 15
70
- HeavyBoson: Annotated[int, "HeavyBoson"] = 16
71
- WBosonLRSM: Annotated[int, "WBosonLRSM"] = 17
72
- NuREle: Annotated[int, "NuREle"] = 18
73
- NuRMu: Annotated[int, "NuRMu"] = 19
74
- NuRTau: Annotated[int, "NuRTau"] = 20
75
- LQ: Annotated[int, "LQ"] = 21
76
- SUSY: Annotated[int, "SUSY"] = 22
77
- OtherBSM: Annotated[int, "OtherBSM"] = 46
78
- LightMeson: Annotated[int, "LightMeson"] = 23
79
- StrangeMeson: Annotated[int, "StrangeMeson"] = 24
80
- CharmedMeson: Annotated[int, "CharmedMeson"] = 25
81
- BottomMeson: Annotated[int, "BottomMeson"] = 26
82
- CCbarMeson: Annotated[int, "CCbarMeson"] = 27
83
- JPsi: Annotated[int, "JPsi"] = 28
84
- BBbarMeson: Annotated[int, "BBbarMeson"] = 29
85
- LightBaryon: Annotated[int, "LightBaryon"] = 30
86
- StrangeBaryon: Annotated[int, "StrangeBaryon"] = 31
87
- CharmedBaryon: Annotated[int, "CharmedBaryon"] = 32
88
- BottomBaryon: Annotated[int, "BottomBaryon"] = 33
89
- PionDecay: Annotated[int, "PionDecay"] = 34
90
- KaonDecay: Annotated[int, "KaonDecay"] = 35
91
- BremPhot: Annotated[int, "BremPhot"] = 36
92
- PromptPhot: Annotated[int, "PromptPhot"] = 37
93
- UndrPhot: Annotated[int, "UndrPhot"] = 38
94
- ISRPhot: Annotated[int, "ISRPhot"] = 39
95
- FSRPhot: Annotated[int, "FSRPhot"] = 40
96
- NucReact: Annotated[int, "NucReact"] = 41
97
- PiZero: Annotated[int, "PiZero"] = 42
98
- DiBoson: Annotated[int, "DiBoson"] = 43
99
- ZorHeavyBoson: Annotated[int, "ZorHeavyBoson"] = 44
100
- MultiBoson: Annotated[int, "MultiBoson"] = 47
101
- QCD: Annotated[int, "QCD"] = 45
102
-
103
-
104
- # https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/EGammaIdentificationRun2#Photon_isEM_word
105
- class PhotonID(IntEnum):
106
- Rhad: Annotated[int, "ClusterHadronicLeakage_Photon"] = 10
107
- E277: Annotated[int, "ClusterMiddleEnergy_Photon"] = 11
108
- Reta: Annotated[int, "ClusterMiddleEratio37_Photon"] = 12
109
- Rphi: Annotated[int, "ClusterMiddleEratio33_Photon"] = 13
110
- Weta2: Annotated[int, "ClusterMiddleWidth_Photon"] = 14
111
- f1: Annotated[int, "ClusterStripsEratio_Photon"] = 15
112
- DeltaE: Annotated[int, "ClusterStripsDeltaE_Photon"] = 17
113
- Wstot: Annotated[int, "ClusterStripsWtot_Photon"] = 18
114
- fside: Annotated[int, "ClusterStripsFracm_Photon"] = 19
115
- Ws3: Annotated[int, "ClusterStripsWeta1c_Photon"] = 20
116
- ERatio: Annotated[int, "ClusterStripsDEmaxs1_Photon"] = 21
3
+ import sys
4
+ from enum import Enum, IntEnum
5
+
6
+ if sys.version_info >= (3, 11):
7
+ from enum import EnumType
8
+ else:
9
+ from enum import EnumMeta as EnumType
10
+
11
+ from typing import Callable, TypeVar, cast
12
+
13
+ _E = TypeVar("_E", bound=Enum)
14
+
15
+
16
+ class MultipleEnumAccessMeta(EnumType):
17
+ """
18
+ Enum Metaclass to provide a way to access multiple values all at once.
19
+ """
20
+
21
+ def __getitem__(self: type[_E], key: str | tuple[str]) -> _E | list[_E]: # type:ignore[misc,override]
22
+ getitem = cast(Callable[[str], _E], super().__getitem__) # type:ignore[misc]
23
+ if isinstance(key, tuple):
24
+ return [getitem(name) for name in key]
25
+ return getitem(key)
26
+
27
+
28
+ class ParticleType(IntEnum, metaclass=MultipleEnumAccessMeta):
29
+ """
30
+ Taken from `ATLAS Truth Utilities for ParticleType <https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L8-49>`_.
31
+ """
32
+
33
+ Unknown = 0
34
+ UnknownElectron = 1
35
+ IsoElectron = 2
36
+ NonIsoElectron = 3
37
+ BkgElectron = 4
38
+ UnknownMuon = 5
39
+ IsoMuon = 6
40
+ NonIsoMuon = 7
41
+ BkgMuon = 8
42
+ UnknownTau = 9
43
+ IsoTau = 10
44
+ NonIsoTau = 11
45
+ BkgTau = 12
46
+ UnknownPhoton = 13
47
+ IsoPhoton = 14
48
+ NonIsoPhoton = 15
49
+ BkgPhoton = 16
50
+ Hadron = 17
51
+ Neutrino = 18
52
+ NuclFrag = 19
53
+ NonPrimary = 20
54
+ GenParticle = 21
55
+ SUSYParticle = 22
56
+ OtherBSMParticle = 39
57
+ BBbarMesonPart = 23
58
+ BottomMesonPart = 24
59
+ CCbarMesonPart = 25
60
+ CharmedMesonPart = 26
61
+ BottomBaryonPart = 27
62
+ CharmedBaryonPart = 28
63
+ StrangeBaryonPart = 29
64
+ LightBaryonPart = 30
65
+ StrangeMesonPart = 31
66
+ LightMesonPart = 32
67
+ BJet = 33
68
+ CJet = 34
69
+ LJet = 35
70
+ GJet = 36
71
+ TauJet = 37
72
+ UnknownJet = 38
73
+
74
+
75
+ class ParticleOrigin(IntEnum, metaclass=MultipleEnumAccessMeta):
76
+ """
77
+ Taken from `ATLAS Truth Utilities for ParticleOrigin <https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103>`_.
78
+ """
79
+
80
+ NonDefined = 0
81
+ SingleElec = 1
82
+ SingleMuon = 2
83
+ SinglePhot = 3
84
+ SingleTau = 4
85
+ PhotonConv = 5
86
+ DalitzDec = 6
87
+ ElMagProc = 7
88
+ Mu = 8
89
+ TauLep = 9
90
+ top = 10
91
+ QuarkWeakDec = 11
92
+ WBoson = 12
93
+ ZBoson = 13
94
+ Higgs = 14
95
+ HiggsMSSM = 15
96
+ HeavyBoson = 16
97
+ WBosonLRSM = 17
98
+ NuREle = 18
99
+ NuRMu = 19
100
+ NuRTau = 20
101
+ LQ = 21
102
+ SUSY = 22
103
+ OtherBSM = 46
104
+ LightMeson = 23
105
+ StrangeMeson = 24
106
+ CharmedMeson = 25
107
+ BottomMeson = 26
108
+ CCbarMeson = 27
109
+ JPsi = 28
110
+ BBbarMeson = 29
111
+ LightBaryon = 30
112
+ StrangeBaryon = 31
113
+ CharmedBaryon = 32
114
+ BottomBaryon = 33
115
+ PionDecay = 34
116
+ KaonDecay = 35
117
+ BremPhot = 36
118
+ PromptPhot = 37
119
+ UndrPhot = 38
120
+ ISRPhot = 39
121
+ FSRPhot = 40
122
+ NucReact = 41
123
+ PiZero = 42
124
+ DiBoson = 43
125
+ ZorHeavyBoson = 44
126
+ MultiBoson = 47
127
+ QCD = 45
128
+
129
+
130
+ class PhotonID(IntEnum, metaclass=MultipleEnumAccessMeta):
131
+ """
132
+ Taken from the `EGamma Identification CP group's twiki <https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/EGammaIdentificationRun2#Photon_isEM_word>`_.
133
+ """
134
+
135
+ Rhad = 10 # ClusterHadronicLeakage_Photon
136
+ E277 = 11 # ClusterMiddleEnergy_Photon
137
+ Reta = 12 # ClusterMiddleEratio37_Photon
138
+ Rphi = 13 # ClusterMiddleEratio33_Photon
139
+ Weta2 = 14 # ClusterMiddleWidth_Photon
140
+ f1 = 15 # ClusterStripsEratio_Photon
141
+ DeltaE = 17 # ClusterStripsDeltaE_Photon
142
+ Wstot = 18 # ClusterStripsWtot_Photon
143
+ fside = 19 # ClusterStripsFracm_Photon
144
+ Ws3 = 20 # ClusterStripsWeta1c_Photon
145
+ ERatio = 21 # ClusterStripsDEmaxs1_Photon
atlas_schema/methods.py CHANGED
@@ -6,6 +6,7 @@ from functools import reduce
6
6
  from operator import ior
7
7
 
8
8
  import awkward
9
+ import particle
9
10
  from coffea.nanoevents.methods import base, candidate, vector
10
11
  from dask_awkward import dask_method
11
12
 
@@ -20,9 +21,9 @@ behavior.update(candidate.behavior)
20
21
 
21
22
  class NtupleEvents(behavior["NanoEvents"]): # type: ignore[misc, valid-type, name-defined]
22
23
  def __repr__(self):
23
- return f"<event {getattr(self,'runNumber','??')}:\
24
- {getattr(self,'eventNumber','??')}:\
25
- {getattr(self,'mcChannelNumber','??')}>"
24
+ return f"<event {getattr(self, 'runNumber', '??')}:\
25
+ {getattr(self, 'eventNumber', '??')}:\
26
+ {getattr(self, 'mcChannelNumber', '??')}>"
26
27
 
27
28
 
28
29
  behavior["NanoEvents"] = NtupleEvents
@@ -49,6 +50,10 @@ class Pass(base.NanoCollection, base.Systematic): ...
49
50
 
50
51
  _set_repr_name("Pass")
51
52
 
53
+ behavior.update(
54
+ awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior)
55
+ )
56
+
52
57
 
53
58
  @awkward.mixin_class(behavior)
54
59
  class Particle(vector.PtEtaPhiMLorentzVector):
@@ -97,42 +102,45 @@ class Particle(vector.PtEtaPhiMLorentzVector):
97
102
 
98
103
  _set_repr_name("Particle")
99
104
 
105
+ ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
106
+ ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
107
+ ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821
108
+ ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
100
109
 
101
- @awkward.mixin_class(behavior)
102
- class MasslessParticle(Particle, base.NanoCollection):
103
- @property
104
- def mass(self):
105
- r"""Invariant mass (+, -, -, -)
106
110
 
107
- :math:`\sqrt{t^2-x^2-y^2-z^2}`
108
- """
109
- return 0.0 * self.pt
110
-
111
-
112
- _set_repr_name("MasslessParticle")
111
+ behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior))
113
112
 
114
113
 
115
114
  @awkward.mixin_class(behavior)
116
- class MissingET(MasslessParticle, base.NanoCollection, base.Systematic):
115
+ class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic):
117
116
  @property
118
- def pt(self):
119
- """Alias for `r`"""
120
- return self["met"] / 1.0e3
117
+ def r(self):
118
+ """Distance from origin in XY plane"""
119
+ return self["met"]
121
120
 
122
- @property
123
- def eta(self):
124
- r"""Pseudorapidity
125
121
 
126
- :math:`-\ln\tan(\theta/2) = \text{arcsinh}(z/r)`
127
- """
128
- return 0.0 * self.pt
122
+ _set_repr_name("MissingET")
129
123
 
124
+ MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821
125
+ MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821
126
+ MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821
127
+ MissingETArray.MomentumClass = MissingETArray # noqa: F821
130
128
 
131
- _set_repr_name("MissingET")
129
+ behavior.update(awkward._util.copy_behaviors("Particle", "Photon", behavior))
132
130
 
133
131
 
134
132
  @awkward.mixin_class(behavior)
135
- class Photon(MasslessParticle, base.NanoCollection, base.Systematic):
133
+ class Photon(Particle, base.NanoCollection, base.Systematic):
134
+ @property
135
+ def mass(self):
136
+ """Return zero mass for photon."""
137
+ return awkward.zeros_like(self.pt)
138
+
139
+ @property
140
+ def charge(self):
141
+ """Return zero charge for photon."""
142
+ return awkward.zeros_like(self.pt)
143
+
136
144
  @property
137
145
  def isEM(self):
138
146
  return self.isEM_syst.NOSYS == 0
@@ -146,20 +154,68 @@ class Photon(MasslessParticle, base.NanoCollection, base.Systematic):
146
154
 
147
155
  _set_repr_name("Photon")
148
156
 
157
+ PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
158
+ PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
159
+ PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821
160
+ PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
161
+
162
+ behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior))
163
+
149
164
 
150
165
  @awkward.mixin_class(behavior)
151
- class Electron(MasslessParticle, base.NanoCollection, base.Systematic): ...
166
+ class Electron(Particle, base.NanoCollection, base.Systematic):
167
+ @property
168
+ def mass(self):
169
+ """Electron mass in GeV"""
170
+ return particle.literals.e_minus.mass / 1.0e3
152
171
 
153
172
 
154
173
  _set_repr_name("Electron")
155
174
 
175
+ ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
176
+ ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
177
+ ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821
178
+ ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
179
+
180
+ behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior))
181
+
156
182
 
157
183
  @awkward.mixin_class(behavior)
158
- class Muon(MasslessParticle, base.NanoCollection, base.Systematic): ...
184
+ class Muon(Particle, base.NanoCollection, base.Systematic):
185
+ @property
186
+ def mass(self):
187
+ """Muon mass in GeV"""
188
+ return particle.literals.mu_minus.mass / 1.0e3
159
189
 
160
190
 
161
191
  _set_repr_name("Muon")
162
192
 
193
+ MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
194
+ MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
195
+ MuonArray.ProjectionClass4D = MuonArray # noqa: F821
196
+ MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
197
+
198
+ behavior.update(awkward._util.copy_behaviors("Particle", "Tau", behavior))
199
+
200
+
201
+ @awkward.mixin_class(behavior)
202
+ class Tau(Particle, base.NanoCollection, base.Systematic):
203
+ @property
204
+ def mass(self):
205
+ """Tau mass in GeV"""
206
+ return particle.literals.tau_minus.mass / 1.0e3
207
+
208
+
209
+ _set_repr_name("Tau")
210
+
211
+ TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
212
+ TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
213
+ TauArray.ProjectionClass4D = TauArray # noqa: F821
214
+ TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
215
+
216
+
217
+ behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior))
218
+
163
219
 
164
220
  @awkward.mixin_class(behavior)
165
221
  class Jet(Particle, base.NanoCollection, base.Systematic): ...
@@ -167,15 +223,19 @@ class Jet(Particle, base.NanoCollection, base.Systematic): ...
167
223
 
168
224
  _set_repr_name("Jet")
169
225
 
226
+ JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
227
+ JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
228
+ JetArray.ProjectionClass4D = JetArray # noqa: F821
229
+ JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
170
230
 
171
231
  __all__ = [
172
- "NtupleEvents",
173
- "Weight",
174
- "Pass",
232
+ "Electron",
233
+ "Jet",
175
234
  "MissingET",
235
+ "Muon",
236
+ "NtupleEvents",
176
237
  "Particle",
238
+ "Pass",
177
239
  "Photon",
178
- "Electron",
179
- "Muon",
180
- "Jet",
240
+ "Weight",
181
241
  ]
atlas_schema/utils.py ADDED
@@ -0,0 +1,39 @@
1
+ from __future__ import annotations
2
+
3
+ from enum import Enum
4
+ from typing import TypeVar, Union, cast
5
+
6
+ import awkward as ak
7
+ import dask_awkward as dak
8
+
9
+ Array = TypeVar("Array", bound=Union[dak.Array, ak.Array])
10
+ _E = TypeVar("_E", bound=Enum)
11
+
12
+
13
+ def isin(haystack: Array, needles: dak.Array | ak.Array, axis: int = -1) -> Array:
14
+ """
15
+ Find needles in haystack.
16
+
17
+ This works by first transforming needles to an array with one more
18
+ dimension than the haystack, placing the needles at axis, and then doing a
19
+ comparison.
20
+
21
+ Args:
22
+ haystack (dak.Array or ak.Array): haystack of values.
23
+ needles (dak.Array or ak.Array): one-dimensional set of needles to find in haystack.
24
+ axis (int): the axis along which the comparison is performed
25
+
26
+ Returns:
27
+ dak.Array or ak.Array: result of comparison for needles in haystack
28
+ """
29
+ assert needles.ndim == 1, "Needles must be one-dimensional"
30
+ assert axis >= -1, "axis must be -1 or positive-valued"
31
+ assert axis < haystack.ndim + 1, "axis too large for the haystack"
32
+
33
+ # First, build up the transformation, with slice(None) indicating where to stick the needles
34
+ reshaper: list[None | slice] = [None] * haystack.ndim
35
+ axis = haystack.ndim if axis == -1 else axis
36
+ reshaper.insert(axis, slice(None))
37
+
38
+ # Note: reshaper needs to be a tuple for indexing purposes
39
+ return cast(Array, ak.any(haystack == needles[tuple(reshaper)], axis=-1))
@@ -1,11 +1,13 @@
1
- Metadata-Version: 2.3
1
+ Metadata-Version: 2.4
2
2
  Name: atlas-schema
3
- Version: 0.2.0
3
+ Version: 0.2.2
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: Changelog, https://github.com/scipp-atlas/atlas-schema/releases
8
+ Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.2.2/
9
+ Project-URL: Releases, https://github.com/scipp-atlas/atlas-schema/releases
10
+ Project-URL: Release Notes, https://atlas-schema.readthedocs.io/en/latest/history.html
9
11
  Author-email: Giordon Stark <kratsg@gmail.com>
10
12
  License:
11
13
  Apache License
@@ -209,6 +211,7 @@ License:
209
211
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
210
212
  See the License for the specific language governing permissions and
211
213
  limitations under the License.
214
+ License-File: LICENSE
212
215
  Classifier: Development Status :: 1 - Planning
213
216
  Classifier: Intended Audience :: Developers
214
217
  Classifier: Intended Audience :: Science/Research
@@ -225,14 +228,19 @@ Classifier: Topic :: Scientific/Engineering
225
228
  Classifier: Typing :: Typed
226
229
  Requires-Python: >=3.9
227
230
  Requires-Dist: coffea[dask]>=2024.4.1
231
+ Requires-Dist: particle>=0.25.0
228
232
  Provides-Extra: dev
229
233
  Requires-Dist: pytest-cov>=3; extra == 'dev'
230
234
  Requires-Dist: pytest>=6; extra == 'dev'
231
235
  Provides-Extra: docs
232
236
  Requires-Dist: furo>=2023.08.17; extra == 'docs'
237
+ Requires-Dist: intersphinx-registry>=0.2411.17; extra == 'docs'
238
+ Requires-Dist: ipywidgets; extra == 'docs'
233
239
  Requires-Dist: myst-parser>=0.13; extra == 'docs'
234
240
  Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
235
- Requires-Dist: sphinx-copybutton; extra == 'docs'
241
+ Requires-Dist: sphinx-click; extra == 'docs'
242
+ Requires-Dist: sphinx-copybutton!=0.5.1,>=0.3.2; extra == 'docs'
243
+ Requires-Dist: sphinx-issues; extra == 'docs'
236
244
  Requires-Dist: sphinx>=7.0; extra == 'docs'
237
245
  Provides-Extra: test
238
246
  Requires-Dist: build; extra == 'test'
@@ -243,7 +251,7 @@ Requires-Dist: tbump>=6.7.0; extra == 'test'
243
251
  Requires-Dist: twine; extra == 'test'
244
252
  Description-Content-Type: text/markdown
245
253
 
246
- # atlas-schema v0.2.0
254
+ # atlas-schema v0.2.2
247
255
 
248
256
  [![Actions Status][actions-badge]][actions-link]
249
257
  [![Documentation Status][rtd-badge]][rtd-link]
@@ -0,0 +1,13 @@
1
+ atlas_schema/__init__.py,sha256=ebY-rTiwSGnfvt1yWATze2GE7K3fVgJj6fT64Sl4sH8,469
2
+ atlas_schema/_version.py,sha256=RrHB9KG1O3GPm--rbTedqmZbdDrbgeRLXBmT4OBUqqI,411
3
+ atlas_schema/_version.pyi,sha256=j5kbzfm6lOn8BzASXWjGIA1yT0OlHTWqlbyZ8Si_o0E,118
4
+ atlas_schema/enums.py,sha256=hwgOvFBmITNxL0MQkrNpbiPv9VMezFoE-eyGgjzem8E,3688
5
+ atlas_schema/methods.py,sha256=K7u6HGKXrtpMg7jjCjKPwIEnknOShUH4HQ1ibKBzkZ0,6832
6
+ atlas_schema/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
7
+ atlas_schema/schema.py,sha256=YRVaiDa5Evl2HZ9CzH23d0-TLkvxqyvFQhn0ixyWCcw,7668
8
+ atlas_schema/typing_compat.py,sha256=RwkxiiYbXO9yxkeaL8CdRaOHH7wq6vO_epg1YD7RbRs,439
9
+ atlas_schema/utils.py,sha256=Oe2G3pe009Uhawsdk9e0MuqOHbAa5vZ8F2F9pOmz_Ok,1442
10
+ atlas_schema-0.2.2.dist-info/METADATA,sha256=QeHezHbhZY-hA2xdVlrQNeZN2OSCA8hn24jzoMUZDX8,16823
11
+ atlas_schema-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
12
+ atlas_schema-0.2.2.dist-info/licenses/LICENSE,sha256=snem82NV8fgAi4DKaaUIfReaM5RqIWbH5OOXOvy40_w,11344
13
+ atlas_schema-0.2.2.dist-info/RECORD,,
@@ -1,4 +1,4 @@
1
1
  Wheel-Version: 1.0
2
- Generator: hatchling 1.26.3
2
+ Generator: hatchling 1.27.0
3
3
  Root-Is-Purelib: true
4
4
  Tag: py3-none-any
@@ -1,12 +0,0 @@
1
- atlas_schema/__init__.py,sha256=eba1N4_cWS5YzEOgDdCJGiPKqUtO-Vn7t8dYq0Q6gk8,354
2
- atlas_schema/_version.py,sha256=H-qsvrxCpdhaQzyddR-yajEqI71hPxLa4KxzpP3uS1g,411
3
- atlas_schema/_version.pyi,sha256=j5kbzfm6lOn8BzASXWjGIA1yT0OlHTWqlbyZ8Si_o0E,118
4
- atlas_schema/enums.py,sha256=RktHMdqNcjcqHOtXzRP5pnakhrfo0VfRl60Mqsa7i74,5522
5
- atlas_schema/methods.py,sha256=Vo9pQ52ZVm13TBpJ6KiiDlN2kS1MMGBMLjwRPB62ces,4111
6
- atlas_schema/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
7
- atlas_schema/schema.py,sha256=YRVaiDa5Evl2HZ9CzH23d0-TLkvxqyvFQhn0ixyWCcw,7668
8
- atlas_schema/typing_compat.py,sha256=RwkxiiYbXO9yxkeaL8CdRaOHH7wq6vO_epg1YD7RbRs,439
9
- atlas_schema-0.2.0.dist-info/METADATA,sha256=TAZNZkLcm-Cy-6s5WEmwQo7YyejBUy0uOlpV9_MNKZI,16395
10
- atlas_schema-0.2.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
11
- atlas_schema-0.2.0.dist-info/licenses/LICENSE,sha256=snem82NV8fgAi4DKaaUIfReaM5RqIWbH5OOXOvy40_w,11344
12
- atlas_schema-0.2.0.dist-info/RECORD,,