From d4b4d6cc68ce80cc8642ca47cf7e9f78e616e2d1 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 19 Jul 2024 11:32:24 +0200 Subject: [PATCH 1/4] Revert "fix: updates from review" --- src/coffea/nanoevents/methods/candidate.py | 39 +++--- src/coffea/nanoevents/methods/delphes.py | 134 ++++++++----------- src/coffea/nanoevents/methods/nanoaod.py | 144 ++++++++------------- src/coffea/nanoevents/methods/pdune.py | 72 +++++------ src/coffea/nanoevents/methods/physlite.py | 70 ++++------ src/coffea/nanoevents/methods/vector.py | 86 +++++------- src/coffea/util.py | 7 - tests/test_nanoevents_vector.py | 14 -- 8 files changed, 219 insertions(+), 347 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 43eca9ea7..e3c427276 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -9,7 +9,6 @@ import numpy from coffea.nanoevents.methods import vector -from coffea.util import register_projection_classes behavior = dict(vector.behavior) @@ -71,29 +70,23 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior)) - -register_projection_classes( - CandidateArray, # noqa: F821 - vector.PolarTwoVectorArray, - vector.SphericalThreeVectorArray, - vector.LorentzVectorArray, - CandidateArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiMCandidateArray, # noqa: F821 - vector.PolarTwoVectorArray, - vector.SphericalThreeVectorArray, - vector.LorentzVectorArray, - PtEtaPhiMCandidateArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiECandidateArray, # noqa: F821 - vector.PolarTwoVectorArray, - vector.SphericalThreeVectorArray, - vector.LorentzVectorArray, - PtEtaPhiECandidateArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior) ) +CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +CandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +CandidateArray.MomentumClass = CandidateArray # noqa: F821 + +PtEtaPhiMCandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.MomentumClass = PtEtaPhiMCandidateArray # noqa: F821 + +PtEtaPhiECandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiECandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiECandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +PtEtaPhiECandidateArray.MomentumClass = PtEtaPhiECandidateArray # noqa: F821 __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index cc65121f6..c232fedbe 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -7,7 +7,6 @@ import numpy from coffea.nanoevents.methods import base, candidate, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -131,16 +130,15 @@ def z(self): _set_repr_name("Vertex") -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior)) - -register_projection_classes( - VertexArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - VertexArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) ) +VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +VertexArray.ProjectionClass4D = VertexArray # noqa: F821 +VertexArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector): @@ -180,17 +178,14 @@ def mass(self): _set_repr_name("Particle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) -register_projection_classes( - ParticleArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - ParticleArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -201,63 +196,57 @@ def mass(self): _set_repr_name("MasslessParticle") -behavior.update(awkward._util.copy_behaviors(Particle, MasslessParticle, behavior)) - -register_projection_classes( - MasslessParticleArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - MasslessParticleArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) ) +MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass4D = MasslessParticleArray # noqa: F821 +MasslessParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Photon, behavior)) - -register_projection_classes( - PhotonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - PhotonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) ) +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Electron, behavior)) - -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) ) +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) - -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - MuonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -265,15 +254,12 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) - -register_projection_classes( - JetArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - JetArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) + +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -281,15 +267,12 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) - -register_projection_classes( - TrackArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - TrackArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) + +TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackArray.ProjectionClass4D = TrackArray # noqa: F821 +TrackArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -300,15 +283,12 @@ def pt(self): _set_repr_name("Tower") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) - -register_projection_classes( - TowerArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - TowerArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) + +TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TowerArray.ProjectionClass4D = TowerArray # noqa: F821 +TowerArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 __all__ = [ diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 141d329d0..8fa587410 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -6,7 +6,6 @@ from dask_awkward import dask_property from coffea.nanoevents.methods import base, candidate, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -39,19 +38,16 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors( vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior ) ) -register_projection_classes( - PtEtaPhiMCollectionArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - PtEtaPhiMCollectionArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +PtEtaPhiMCollectionArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass4D = PtEtaPhiMCollectionArray # noqa: F821 +PtEtaPhiMCollectionArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -160,17 +156,14 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) ) -register_projection_classes( - GenParticleArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - GenParticleArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +GenParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenParticleArray.ProjectionClass4D = GenParticleArray # noqa: F821 +GenParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -191,17 +184,14 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) ) -register_projection_classes( - GenVisTauArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - GenVisTauArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +GenVisTauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass4D = GenVisTauArray # noqa: F821 +GenVisTauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -266,17 +256,14 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) ) -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -353,17 +340,14 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) ) -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - MuonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -388,17 +372,14 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) ) -register_projection_classes( - TauArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - TauArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TauArray.ProjectionClass4D = TauArray # noqa: F821 +TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -474,17 +455,14 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) ) -register_projection_classes( - PhotonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - PhotonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -501,17 +479,14 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) ) -register_projection_classes( - FsrPhotonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - FsrPhotonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +FsrPhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass4D = FsrPhotonArray # noqa: F821 +FsrPhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -582,17 +557,14 @@ def constituents(self, dask_array): _set_repr_name("Jet") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) ) -register_projection_classes( - JetArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - JetArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -655,17 +627,14 @@ def constituents(self, dask_array): _set_repr_name("FatJet") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) ) -register_projection_classes( - FatJetArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - FatJetArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +FatJetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FatJetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FatJetArray.ProjectionClass4D = FatJetArray # noqa: F821 +FatJetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -678,17 +647,14 @@ def r(self): _set_repr_name("MissingET") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) ) -register_projection_classes( - MissingETArray, # noqa: F821 - MissingETArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 - MissingETArray, # noqa: F821 -) +MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821 +MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821 +MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +MissingETArray.MomentumClass = MissingETArray # noqa: F821 @awkward.mixin_class(behavior) diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 088855850..4c4296008 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -4,7 +4,6 @@ import numpy from coffea.nanoevents.methods import base, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -93,16 +92,15 @@ def mass(self): _set_repr_name("Particle") -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior)) - -register_projection_classes( - ParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ParticleArray, # noqa: F821 - vector.LorentzVectorArray, +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior) ) +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -140,17 +138,14 @@ def t(self): _set_repr_name("TrackParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) -register_projection_classes( - TrackParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TrackParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -170,15 +165,12 @@ def trackParticle(self): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) - -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - MuonArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -206,15 +198,12 @@ def trackParticle(self): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) - -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -257,14 +246,11 @@ def parents(self): _set_repr_name("TruthParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) -register_projection_classes( - TruthParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TruthParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 +TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index d04787208..28554a6c9 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -8,7 +8,6 @@ from dask_awkward import dask_property from coffea.nanoevents.methods import base, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -144,17 +143,14 @@ def mass(self): _set_repr_name("Particle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) -register_projection_classes( - ParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -193,17 +189,14 @@ def t(self): _set_repr_name("TrackParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) -register_projection_classes( - TrackParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TrackParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -232,15 +225,12 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) - -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - MuonArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -293,15 +283,12 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) - -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -344,14 +331,11 @@ def parents(self): _set_repr_name("TruthParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) -register_projection_classes( - TruthParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TruthParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 +TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 79e334024..2944aa72d 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -58,7 +58,7 @@ MomentumAwkward4D, ) -from coffea.util import deprecate, register_projection_classes +from coffea.util import deprecate _cst = pytz.timezone("US/Central") _depttime = _cst.localize(datetime(2024, 6, 30, 11, 59, 59)) @@ -802,56 +802,40 @@ def divide(self, other): behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract -register_projection_classes( - TwoVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 -) -register_projection_classes( - PolarTwoVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 -) -register_projection_classes( - ThreeVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 -) -register_projection_classes( - SphericalThreeVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 -) -register_projection_classes( - LorentzVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiMLorentzVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiELorentzVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 -) - +TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +TwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 + +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PolarTwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 + +ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +ThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 + +SphericalThreeVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +SphericalThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 + +LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 + +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 + +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 __all__ = [ "TwoVector", diff --git a/src/coffea/util.py b/src/coffea/util.py index 662d28638..bfb0b5119 100644 --- a/src/coffea/util.py +++ b/src/coffea/util.py @@ -246,10 +246,3 @@ def _remove_not_interpretable(branch, emit_warning=True): return False else: return True - - -def register_projection_classes(vec_class, proj_2d, proj_3d, proj_4d, momentum_class): - vec_class.ProjectionClass2D = proj_2d - vec_class.ProjectionClass3D = proj_3d - vec_class.ProjectionClass4D = proj_4d - vec_class.MomentumClass = momentum_class diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 665970ab1..6aa13d164 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,13 +543,6 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) - with pytest.raises(TypeError): - a + b == b + a - with pytest.raises(TypeError): - a + c == c + a - with pytest.raises(TypeError): - b + c == c + b - assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) @@ -558,13 +551,6 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) - with pytest.raises(TypeError): - a - b == -(b - a) - with pytest.raises(TypeError): - a - c == -(c - a) - with pytest.raises(TypeError): - b - c == -(c - b) - @pytest.mark.parametrize("optimization_enabled", [True, False]) def test_dask_metric_table_and_nearest(optimization_enabled): From dc80e611278ea30b77fd64d675aaf25c7a553375 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 19 Jul 2024 11:44:18 +0200 Subject: [PATCH 2/4] Add back useful changes * More tests * Don't modify the global awkward behavior --- src/coffea/nanoevents/methods/candidate.py | 4 +--- src/coffea/nanoevents/methods/delphes.py | 20 +++++++++----------- src/coffea/nanoevents/methods/nanoaod.py | 22 +++++++++++----------- src/coffea/nanoevents/methods/pdune.py | 12 +++++------- src/coffea/nanoevents/methods/physlite.py | 10 +++++----- tests/test_nanoevents_vector.py | 14 ++++++++++++++ 6 files changed, 45 insertions(+), 37 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index e3c427276..f4f417b63 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,9 +70,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior) -) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior)) CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index c232fedbe..cc734824b 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -130,9 +130,7 @@ def z(self): _set_repr_name("Vertex") -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) -) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior)) VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -178,7 +176,7 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) @@ -196,7 +194,7 @@ def mass(self): _set_repr_name("MasslessParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) ) @@ -211,7 +209,7 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) ) @@ -226,7 +224,7 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) ) @@ -241,7 +239,7 @@ class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -254,7 +252,7 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -267,7 +265,7 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -283,7 +281,7 @@ def pt(self): _set_repr_name("Tower") -awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 8fa587410..e46b3c287 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -38,7 +38,7 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors( vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior ) @@ -156,7 +156,7 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) ) @@ -184,7 +184,7 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) ) @@ -256,7 +256,7 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) ) @@ -340,7 +340,7 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) ) @@ -372,7 +372,7 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) ) @@ -455,7 +455,7 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) ) @@ -479,7 +479,7 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) ) @@ -557,7 +557,7 @@ def constituents(self, dask_array): _set_repr_name("Jet") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) ) @@ -627,7 +627,7 @@ def constituents(self, dask_array): _set_repr_name("FatJet") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) ) @@ -647,7 +647,7 @@ def r(self): _set_repr_name("MissingET") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) ) diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 4c4296008..43e2649b5 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -92,9 +92,7 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior) -) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior)) ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -138,7 +136,7 @@ def t(self): _set_repr_name("TrackParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) @@ -165,7 +163,7 @@ def trackParticle(self): _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -198,7 +196,7 @@ def trackParticle(self): _set_repr_name("Electron") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -246,7 +244,7 @@ def parents(self): _set_repr_name("TruthParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index 28554a6c9..301ad3b08 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -143,7 +143,7 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) @@ -189,7 +189,7 @@ def t(self): _set_repr_name("TrackParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) @@ -225,7 +225,7 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -283,7 +283,7 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -331,7 +331,7 @@ def parents(self): _set_repr_name("TruthParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 6aa13d164..46b9cbd1e 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,6 +543,13 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) +with pytest.raises(TypeError): + a + b == b + a + with pytest.raises(TypeError): + a + c == c + a + with pytest.raises(TypeError): + b + c == c + b + assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) @@ -551,6 +558,13 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) + with pytest.raises(TypeError): + a - b == -(b - a) + with pytest.raises(TypeError): + a - c == -(c - a) + with pytest.raises(TypeError): + b - c == -(c - b) + @pytest.mark.parametrize("optimization_enabled", [True, False]) def test_dask_metric_table_and_nearest(optimization_enabled): From 8349641317f21ee0ab0784d55f27e52397a1a865 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 19 Jul 2024 09:45:44 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/coffea/nanoevents/methods/delphes.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index cc734824b..5a5b0979a 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -194,9 +194,7 @@ def mass(self): _set_repr_name("MasslessParticle") -behavior.update( - awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) -) +behavior.update(awkward._util.copy_behaviors(Particle, MasslessParticle, behavior)) MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -209,9 +207,7 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -behavior.update( - awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) -) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Photon, behavior)) PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -224,9 +220,7 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -behavior.update( - awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) -) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Electron, behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 From fed625b100e6e02e903e1a9cddef61d6c63a58f3 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 19 Jul 2024 11:46:57 +0200 Subject: [PATCH 4/4] Indentation --- tests/test_nanoevents_vector.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 46b9cbd1e..665970ab1 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,12 +543,12 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) -with pytest.raises(TypeError): - a + b == b + a - with pytest.raises(TypeError): - a + c == c + a - with pytest.raises(TypeError): - b + c == c + b + with pytest.raises(TypeError): + a + b == b + a + with pytest.raises(TypeError): + a + c == c + a + with pytest.raises(TypeError): + b + c == c + b assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) @@ -559,11 +559,11 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) with pytest.raises(TypeError): - a - b == -(b - a) - with pytest.raises(TypeError): - a - c == -(c - a) - with pytest.raises(TypeError): - b - c == -(c - b) + a - b == -(b - a) + with pytest.raises(TypeError): + a - c == -(c - a) + with pytest.raises(TypeError): + b - c == -(c - b) @pytest.mark.parametrize("optimization_enabled", [True, False])