Skip to content

Commit 90a50ce

Browse files
ikrommydikrommyd
authored
feat: ensure the right fields are present at the schema level to satisfy 4-vector behaviors (#1481)
* add CorrT1METJet to nanoaod and add zero trigobj mass * refactor * fix comments * make tracing work with form changes? * do not use dask-awkward utils * better implementation using awkward internals * move report def * remove commented out code * remove commented out code * add key assertion in tracing * remove behaviors for what's added at the schema level * move behaviors to the schema for dune * physlite and pdune * implement full-like form and fix physlite and pdune * fix missinget in delphes * copy behaviors too * add comment about MET eta * we don't want these here * this feels better * stuck ci * show value in full-like warning
1 parent 70262f5 commit 90a50ce

File tree

9 files changed

+190
-239
lines changed

9 files changed

+190
-239
lines changed

‎src/coffea/nanoevents/methods/delphes.py‎

Lines changed: 17 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
"""
55

66
import awkward
7-
import numpy
87

98
from coffea.nanoevents.methods import base, candidate, vector
109

@@ -86,26 +85,27 @@ class ScalarHT(base.NanoCollection): ...
8685
_set_repr_name("ScalarHT")
8786

8887

88+
behavior.update(
89+
awkward._util.copy_behaviors("SphericalThreeVector", "MissingET", behavior)
90+
)
91+
92+
8993
@awkward.mixin_class(behavior)
90-
class MissingET(vector.SphericalThreeVector, base.NanoCollection):
91-
@property
92-
def rho(self):
93-
return self["MET"] * numpy.cosh(self.eta)
94+
class MissingET(vector.SphericalThreeVector, base.NanoCollection): ...
9495

95-
@property
96-
def theta(self):
97-
return 2 * numpy.arctan(numpy.exp(-self.eta))
9896

99-
@property
100-
def phi(self):
101-
return self["Phi"]
97+
_set_repr_name("MissingET")
10298

103-
@property
104-
def eta(self):
105-
return self["Eta"]
10699

100+
MissingETArray.ProjectionClass2D = vector.PolarTwoVectorArray # noqa: F821
101+
MissingETArray.ProjectionClass3D = MissingETArray # noqa: F821
102+
MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821
103+
MissingETArray.MomentumClass = MissingETArray # noqa: F821
104+
MissingETRecord.ProjectionClass2D = vector.PolarTwoVectorRecord # noqa: F821
105+
MissingETRecord.ProjectionClass3D = MissingETRecord # noqa: F821
106+
MissingETRecord.ProjectionClass4D = vector.LorentzVectorRecord # noqa: F821
107+
MissingETRecord.MomentumClass = MissingETRecord # noqa: F821
107108

108-
_set_repr_name("MissingET")
109109

110110
behavior.update(awkward._util.copy_behaviors("LorentzVector", "Vertex", behavior))
111111

@@ -114,22 +114,6 @@ def eta(self):
114114
class Vertex(vector.LorentzVector):
115115
"""Generic vertex collection that has Lorentz vector properties"""
116116

117-
@property
118-
def t(self):
119-
return self["T"]
120-
121-
@property
122-
def x(self):
123-
return self["X"]
124-
125-
@property
126-
def y(self):
127-
return self["Y"]
128-
129-
@property
130-
def z(self):
131-
return self["Z"]
132-
133117

134118
_set_repr_name("Vertex")
135119

@@ -167,22 +151,6 @@ class Particle(vector.PtEtaPhiMLorentzVector):
167151
- Z: particle vertex position (z component)
168152
"""
169153

170-
@property
171-
def pt(self):
172-
return self["PT"]
173-
174-
@property
175-
def eta(self):
176-
return self["Eta"]
177-
178-
@property
179-
def phi(self):
180-
return self["Phi"]
181-
182-
@property
183-
def mass(self):
184-
return self["Mass"]
185-
186154

187155
_set_repr_name("Particle")
188156

@@ -199,10 +167,7 @@ def mass(self):
199167

200168

201169
@awkward.mixin_class(behavior)
202-
class MasslessParticle(Particle, base.NanoCollection):
203-
@property
204-
def mass(self):
205-
return awkward.zeros_like(self.pt)
170+
class MasslessParticle(Particle, base.NanoCollection): ...
206171

207172

208173
_set_repr_name("MasslessParticle")
@@ -310,10 +275,7 @@ class Track(Particle, base.NanoCollection): ...
310275

311276

312277
@awkward.mixin_class(behavior)
313-
class Tower(MasslessParticle, base.NanoCollection):
314-
@property
315-
def pt(self):
316-
return self["ET"]
278+
class Tower(MasslessParticle, base.NanoCollection): ...
317279

318280

319281
_set_repr_name("Tower")

‎src/coffea/nanoevents/methods/nanoaod.py‎

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -459,14 +459,6 @@ class Photon(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic)
459459
TIGHT = 3
460460
"cutBased selection minimum value"
461461

462-
@property
463-
def mass(self):
464-
return awkward.zeros_like(self.pt)
465-
466-
@property
467-
def charge(self):
468-
return awkward.zeros_like(self.pt)
469-
470462
@property
471463
def isLoose(self):
472464
"""Returns a boolean array marking loose cut-based photons"""
@@ -576,10 +568,6 @@ class Jet(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic):
576568
TIGHTLEPVETO = 2
577569
"jetId bit position"
578570

579-
@property
580-
def charge(self):
581-
return awkward.zeros_like(self.pt)
582-
583571
@property
584572
def isLoose(self):
585573
"""Returns a boolean array marking loose jets according to jetId index"""
@@ -674,10 +662,6 @@ class FatJet(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic)
674662
TIGHTLEPVETO = 2
675663
"jetId bit position"
676664

677-
@property
678-
def charge(self):
679-
return awkward.zeros_like(self.pt)
680-
681665
@property
682666
def isLoose(self):
683667
"""Returns a boolean array marking loose jets according to jetId index"""
@@ -745,11 +729,6 @@ def constituents(self, dask_array):
745729
class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic):
746730
"""NanoAOD Missing transverse energy object"""
747731

748-
@property
749-
def r(self):
750-
"""Distance from origin in XY plane"""
751-
return self["pt"]
752-
753732

754733
_set_repr_name("MissingET")
755734

‎src/coffea/nanoevents/methods/pdune.py‎

Lines changed: 0 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -89,10 +89,6 @@ def _get_global_index(target, eventindex, index):
8989
class Particle(vector.LorentzVector, base.NanoCollection):
9090
"""Generic particle collection that has Lorentz vector properties"""
9191

92-
@property
93-
def mass(self):
94-
return self.m
95-
9692

9793
_set_repr_name("Particle")
9894

@@ -116,34 +112,6 @@ class TrackParticle(vector.LorentzVector, base.NanoCollection):
116112
<https://gitlab.cern.ch/atlas/athena/-/blob/21.2/Event/xAOD/xAODTracking/Root/TrackParticle_v1.cxx#L82>`_.
117113
"""
118114

119-
@property
120-
def theta(self):
121-
return self["theta"]
122-
123-
@property
124-
def phi(self):
125-
return self["phi"]
126-
127-
@property
128-
def p(self):
129-
return 1.0 / numpy.abs(self.qOverP)
130-
131-
@property
132-
def x(self):
133-
return self.p * numpy.sin(self.theta) * numpy.cos(self.phi)
134-
135-
@property
136-
def y(self):
137-
return self.p * numpy.sin(self.theta) * numpy.sin(self.phi)
138-
139-
@property
140-
def z(self):
141-
return self.p * numpy.cos(self.theta)
142-
143-
@property
144-
def t(self):
145-
return numpy.sqrt(139.570**2 + self.x**2 + self.y**2 + self.z**2)
146-
147115

148116
_set_repr_name("TrackParticle")
149117

@@ -235,26 +203,6 @@ class TruthParticle(vector.LorentzVector, base.NanoCollection):
235203
<https://gitlab.cern.ch/atlas/athena/-/blob/21.2/Event/xAOD/xAODTruth/Root/TruthParticle_v1.cxx>`_.
236204
"""
237205

238-
@property
239-
def x(self):
240-
return self["px"]
241-
242-
@property
243-
def y(self):
244-
return self["py"]
245-
246-
@property
247-
def z(self):
248-
return self["pz"]
249-
250-
@property
251-
def t(self):
252-
return self["e"]
253-
254-
@property
255-
def mass(self):
256-
return self["m"]
257-
258206
@property
259207
def children(self):
260208
return _element_link_multiple(

‎src/coffea/nanoevents/methods/physlite.py‎

Lines changed: 0 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -150,10 +150,6 @@ def _get_global_index(target, eventindex, index):
150150
class Particle(vector.PtEtaPhiMLorentzVector, base.NanoCollection):
151151
"""Generic particle collection that has Lorentz vector properties"""
152152

153-
@property
154-
def mass(self):
155-
return self.m
156-
157153

158154
_set_repr_name("Particle")
159155

@@ -177,34 +173,6 @@ class TrackParticle(vector.LorentzVector, base.NanoCollection):
177173
<https://gitlab.cern.ch/atlas/athena/-/blob/21.2/Event/xAOD/xAODTracking/Root/TrackParticle_v1.cxx#L82>`_.
178174
"""
179175

180-
@property
181-
def theta(self):
182-
return self["theta"]
183-
184-
@property
185-
def phi(self):
186-
return self["phi"]
187-
188-
@property
189-
def p(self):
190-
return 1.0 / numpy.abs(self.qOverP)
191-
192-
@property
193-
def x(self):
194-
return self.p * numpy.sin(self.theta) * numpy.cos(self.phi)
195-
196-
@property
197-
def y(self):
198-
return self.p * numpy.sin(self.theta) * numpy.sin(self.phi)
199-
200-
@property
201-
def z(self):
202-
return self.p * numpy.cos(self.theta)
203-
204-
@property
205-
def t(self):
206-
return numpy.sqrt(139.570**2 + self.x**2 + self.y**2 + self.z**2)
207-
208176

209177
_set_repr_name("TrackParticle")
210178

@@ -328,26 +296,6 @@ class TruthParticle(vector.LorentzVector, base.NanoCollection):
328296
<https://gitlab.cern.ch/atlas/athena/-/blob/21.2/Event/xAOD/xAODTruth/Root/TruthParticle_v1.cxx>`_.
329297
"""
330298

331-
@property
332-
def x(self):
333-
return self["px"]
334-
335-
@property
336-
def y(self):
337-
return self["py"]
338-
339-
@property
340-
def z(self):
341-
return self["pz"]
342-
343-
@property
344-
def t(self):
345-
return self["e"]
346-
347-
@property
348-
def mass(self):
349-
return self["m"]
350-
351299
@property
352300
def children(self):
353301
return _element_link_multiple(

‎src/coffea/nanoevents/schemas/delphes.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,41 @@ def _preprocess_branch_form(objname, form):
283283
for k in branch_forms
284284
if k.startswith(name + "/" + name)
285285
}
286+
# add appropriate aliases expected from scikit-hep/vector
287+
if mixin == "MissingET":
288+
# see the discussion in https://github.com/scikit-hep/coffea/pull/1481
289+
# regarding why MET has an eta
290+
content["rho"] = transforms.met_to_rho_form(
291+
content["MET"], content["Eta"]
292+
)
293+
content["eta"] = content["Eta"]
294+
content["phi"] = content["Phi"]
295+
elif mixin == "Vertex":
296+
content["t"] = content["T"]
297+
content["x"] = content["X"]
298+
content["y"] = content["Y"]
299+
content["z"] = content["Z"]
300+
elif mixin == "Particle" or mixin == "Jet" or mixin == "Track":
301+
content["pt"] = content["PT"]
302+
content["eta"] = content["Eta"]
303+
content["phi"] = content["Phi"]
304+
content["mass"] = content["Mass"]
305+
elif (
306+
mixin == "MasslessParticle"
307+
or mixin == "Photon"
308+
or mixin == "Electron"
309+
or mixin == "Muon"
310+
or mixin == "Tower"
311+
):
312+
if "PT" not in content and "ET" in content:
313+
content["PT"] = content["ET"]
314+
content["pt"] = content["PT"]
315+
content["eta"] = content["Eta"]
316+
content["phi"] = content["Phi"]
317+
content["mass"] = transforms.full_like_from_offsets_form(
318+
branch_forms[f"o{name}"], 0.0
319+
)
320+
286321
output[name] = zip_forms(content, name, record_name=mixin, offsets=offsets)
287322

288323
# update docstrings as needed

0 commit comments

Comments
 (0)