diff --git a/pyproject.toml b/pyproject.toml index 6543c555a..2529ae967 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ "dask[array]>=2023.4.0;python_version<'3.9'", "dask-awkward>=2024.3.0", "dask-histogram>=2024.3.0", + "vector>=1.3.1", "correctionlib>=2.6.0", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", diff --git a/src/coffea/nanoevents/methods/base.py b/src/coffea/nanoevents/methods/base.py index ec9ae2d6b..896f53a25 100644 --- a/src/coffea/nanoevents/methods/base.py +++ b/src/coffea/nanoevents/methods/base.py @@ -154,6 +154,9 @@ def add_systematic( Systematic.add_kind(kind) +behavior[("__typestr__", "NanoEvents")] = "event" + + @awkward.mixin_class(behavior) class NanoEvents(Systematic): """NanoEvents mixin class @@ -167,9 +170,6 @@ def metadata(self): return self.layout.purelist_parameter("metadata") -behavior[("__typestr__", "NanoEvents")] = "event" - - @awkward.mixin_class(behavior) class NanoCollection: """A NanoEvents collection diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 2291a744a..31a6ac0e5 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -12,6 +12,8 @@ behavior = dict(vector.behavior) +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Candidate", behavior)) + @awkward.mixin_class(behavior) class Candidate(vector.LorentzVector): @@ -70,4 +72,19 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass +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 1f22e814e..6b061d369 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -107,6 +107,8 @@ def eta(self): _set_repr_name("MissingET") +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Vertex", behavior)) + @awkward.mixin_class(behavior) class Vertex(vector.LorentzVector): @@ -131,6 +133,15 @@ def z(self): _set_repr_name("Vertex") +VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +VertexArray.ProjectionClass4D = VertexArray # noqa: F821 +VertexArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) +) + @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector): @@ -171,6 +182,13 @@ def mass(self): _set_repr_name("Particle") +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "MasslessParticle", behavior)) + @awkward.mixin_class(behavior) class MasslessParticle(Particle, base.NanoCollection): @@ -181,6 +199,13 @@ def mass(self): _set_repr_name("MasslessParticle") +MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass4D = MasslessParticleArray # noqa: F821 +MasslessParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Photon", behavior)) + @awkward.mixin_class(behavior) class Photon(MasslessParticle, base.NanoCollection): ... @@ -188,6 +213,13 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Electron", behavior)) + @awkward.mixin_class(behavior) class Electron(MasslessParticle, base.NanoCollection): ... @@ -195,6 +227,13 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(MasslessParticle, base.NanoCollection): ... @@ -202,6 +241,13 @@ class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior)) + @awkward.mixin_class(behavior) class Jet(Particle, base.NanoCollection): ... @@ -209,6 +255,13 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "Track", behavior)) + @awkward.mixin_class(behavior) class Track(Particle, base.NanoCollection): ... @@ -216,6 +269,13 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") +TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackArray.ProjectionClass4D = TrackArray # noqa: F821 +TrackArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Tower", behavior)) + @awkward.mixin_class(behavior) class Tower(MasslessParticle, base.NanoCollection): @@ -226,6 +286,11 @@ def pt(self): _set_repr_name("Tower") +TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TowerArray.ProjectionClass4D = TowerArray # noqa: F821 +TowerArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + __all__ = [ "DelphesEvents", diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index d947d33da..f6517a599 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -31,6 +31,13 @@ def namefcn(self): behavior[classname].__repr__ = namefcn +behavior.update( + awkward._util.copy_behaviors( + "PtEtaPhiMLorentzVector", "PtEtaPhiMCollection", behavior + ) +) + + @awkward.mixin_class(behavior) class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """Generic collection that has Lorentz vector properties""" @@ -38,6 +45,17 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass +PtEtaPhiMCollectionArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass4D = PtEtaPhiMCollectionArray # noqa: F821 +PtEtaPhiMCollectionArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenParticle", behavior) +) + + @awkward.mixin_class(behavior) class GenParticle(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """NanoAOD generator-level particle object, including parent and child self-references @@ -145,6 +163,15 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") +GenParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenParticleArray.ProjectionClass4D = GenParticleArray # noqa: F821 +GenParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenVisTau", behavior) +) + @awkward.mixin_class(behavior) class GenVisTau(candidate.PtEtaPhiMCandidate, base.NanoCollection): @@ -165,6 +192,15 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") +GenVisTauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass4D = GenVisTauArray # noqa: F821 +GenVisTauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Electron", behavior) +) + @awkward.mixin_class(behavior) class Electron(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -229,6 +265,15 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "LowPtElectron", behavior) +) + @awkward.mixin_class(behavior) class LowPtElectron(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -263,6 +308,13 @@ def matched_photon(self, dask_array): _set_repr_name("LowPtElectron") +LowPtElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +LowPtElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +LowPtElectronArray.ProjectionClass4D = LowPtElectronArray # noqa: F821 +LowPtElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -297,6 +349,13 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Tau", behavior)) + @awkward.mixin_class(behavior) class Tau(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -321,6 +380,13 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") +TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TauArray.ProjectionClass4D = TauArray # noqa: F821 +TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Photon", behavior)) + @awkward.mixin_class(behavior) class Photon(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -396,6 +462,15 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "FsrPhoton", behavior) +) + @awkward.mixin_class(behavior) class FsrPhoton(candidate.PtEtaPhiMCandidate, base.NanoCollection): @@ -412,6 +487,13 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") +FsrPhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass4D = FsrPhotonArray # noqa: F821 +FsrPhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Jet", behavior)) + @awkward.mixin_class(behavior) class Jet(vector.PtEtaPhiMLorentzVector, base.NanoCollection, base.Systematic): @@ -482,6 +564,15 @@ def constituents(self, dask_array): _set_repr_name("Jet") +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "FatJet", behavior) +) + @awkward.mixin_class(behavior) class FatJet(vector.PtEtaPhiMLorentzVector, base.NanoCollection, base.Systematic): @@ -544,6 +635,13 @@ def constituents(self, dask_array): _set_repr_name("FatJet") +FatJetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FatJetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FatJetArray.ProjectionClass4D = FatJetArray # noqa: F821 +FatJetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior)) + @awkward.mixin_class(behavior) class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic): @@ -556,6 +654,11 @@ def r(self): _set_repr_name("MissingET") +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) class Vertex(base.NanoCollection): diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index aff8e7730..a32e6c0e0 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -82,6 +82,9 @@ def _get_global_index(target, eventindex, index): return target_offsets + index +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Particle", behavior)) + + @awkward.mixin_class(behavior) class Particle(vector.LorentzVector, base.NanoCollection): """Generic particle collection that has Lorentz vector properties""" @@ -93,6 +96,15 @@ def mass(self): _set_repr_name("Particle") +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) +) + @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -131,6 +143,13 @@ def t(self): _set_repr_name("TrackParticle") +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(Particle): @@ -150,6 +169,13 @@ def trackParticle(self): _set_repr_name("Muon") +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) + @awkward.mixin_class(behavior) class Electron(Particle): @@ -177,6 +203,15 @@ def trackParticle(self): _set_repr_name("Electron") +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) +) + @awkward.mixin_class(behavior) class TruthParticle(vector.LorentzVector, base.NanoCollection): @@ -218,3 +253,8 @@ def parents(self): _set_repr_name("TruthParticle") + +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 a20b7e8ec..61165ee56 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -133,6 +133,11 @@ def _get_global_index(target, eventindex, index): return target_offsets + index +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) +) + + @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """Generic particle collection that has Lorentz vector properties""" @@ -144,6 +149,15 @@ def mass(self): _set_repr_name("Particle") +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) +) + @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -182,6 +196,13 @@ def t(self): _set_repr_name("TrackParticle") +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(Particle): @@ -210,6 +231,13 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) + @awkward.mixin_class(behavior) class Electron(Particle): @@ -262,6 +290,15 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) +) + @awkward.mixin_class(behavior) class TruthParticle(vector.LorentzVector, base.NanoCollection): @@ -303,3 +340,8 @@ def parents(self): _set_repr_name("TruthParticle") + +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 515c5ef7c..2f111373b 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -50,19 +50,25 @@ import numba import numpy import pytz +import vector from dask_awkward import dask_method +from vector.backends.awkward import ( + MomentumAwkward2D, + MomentumAwkward3D, + MomentumAwkward4D, +) from coffea.util import deprecate _cst = pytz.timezone("US/Central") -_depttime = _cst.localize(datetime(2024, 6, 30, 11, 59, 59)) +_depttime = _cst.localize(datetime(2024, 12, 31, 11, 59, 59)) deprecate( ( "coffea.nanoevents.methods.vector will be removed and replaced with scikit-hep vector. " "Nanoevents schemas internal to coffea will be migrated. " "Otherwise please consider using that package!" ), - version="2024.7.0", + version="2025.1.0", date=str(_depttime), category=FutureWarning, ) @@ -109,10 +115,11 @@ def delta_r(eta1, phi1, eta2, phi2): behavior = {} +behavior.update(vector.backends.awkward.behavior) @awkward.mixin_class(behavior) -class TwoVector: +class TwoVector(MomentumAwkward2D): """A cartesian 2-dimensional vector A heavy emphasis towards a momentum vector interpretation is assumed, hence @@ -127,40 +134,12 @@ def r(self): :math:`\sqrt{x^2+y^2}` """ - return numpy.sqrt(self.r2) - - @property - def phi(self): - r"""Polar angle relative to X axis - - :math:`\text{arctan2}(y, x)` - """ - return numpy.arctan2(self.y, self.x) - - @property - def px(self): - """Alias for `x`""" - return self.x - - @property - def py(self): - """Alias for `y`""" - return self.y + return self.rho @property def r2(self): """Squared `r`""" - return self.x * self.x + self.y * self.y - - @property - def pt2(self): - """Alias for `r2`""" - return self.r2 - - @property - def pt(self): - """Alias for `r`""" - return self.r + return self.rho2 @awkward.mixin_class_method(numpy.absolute) def absolute(self): @@ -173,40 +152,7 @@ def absolute(self): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y}, - with_name="TwoVector", - behavior=self.behavior, - ) - - @awkward.mixin_class_method(numpy.add, {"TwoVector"}) - def add(self, other): - """Add two vectors together elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x + other.x, "y": self.y + other.y}, - with_name="TwoVector", - behavior=self.behavior, - ) - - @awkward.mixin_class_method( - numpy.subtract, - { - "TwoVector", - "ThreeVector", - "SphericalThreeVector", - "LorentzVector", - "PtEtaPhiMLorentzVector", - "PtEtaPhiELorentzVector", - }, - transpose=False, - ) - def subtract(self, other): - """Subtract a vector from another elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x - other.x, "y": self.y - other.y}, - with_name="TwoVector", - behavior=self.behavior, - ) + return self.scale(-1) def sum(self, axis=-1): """Sum an array of vectors elementwise using `x` and `y` components""" @@ -222,29 +168,21 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x * other, "y": self.y * other}, - with_name="TwoVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def delta_phi(self, other): """Compute difference in angle between two vectors Returns a value within [-pi, pi) """ - return delta_phi(self.phi, other.phi) - - def dot(self, other): - """Compute the dot product of two vectors""" - return self.x * other.x + self.y * other.y + return self.deltaphi(other) @property def unit(self): @@ -256,74 +194,26 @@ def unit(self): class PolarTwoVector(TwoVector): """A polar coordinate 2-dimensional vector - This mixin class requires the parent class to provide items `r` and `phi`. + This mixin class requires the parent class to provide items `rho` and `phi`. Some additional properties are overridden for performance """ - @property - def x(self): - r"""Cartesian x value - - :math:`r \cos{\phi}` - """ - return self.r * numpy.cos(self.phi) - - @property - def y(self): - r"""Cartesian y value - - :math:`r \sin{\phi}` - """ - return self.r * numpy.sin(self.phi) - - @property - def r(self): - r"""Distance from origin in XY plane - - :math:`\sqrt{x^2+y^2}` - """ - return self["r"] - - @property - def phi(self): - r"""Azimuthal angle relative to X axis in XY plane - - :math:`\text{arctan2}(y, x)` - """ - return self["phi"] - - @property - def r2(self): - """Squared `r`""" - return self.r * self.r - @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using using `x` and `y` components In reality, this directly adjusts `r` and `phi` for performance """ - return awkward.zip( - { - "r": self.r * abs(other), - "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - }, - with_name="PolarTwoVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"r": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, - with_name="PolarTwoVector", - behavior=self.behavior, - ) + return self.scale(-1) @awkward.mixin_class(behavior) -class ThreeVector(TwoVector): +class ThreeVector(MomentumAwkward3D): """A cartesian 3-dimensional vector A heavy emphasis towards a momentum vector interpretation is assumed. @@ -331,41 +221,18 @@ class ThreeVector(TwoVector): """ @property - def pz(self): - """Alias for `z`""" - return self.z - - @property - def rho2(self): - """Squared `rho`""" - return self.r2 + self.z * self.z - - @property - def rho(self): - r"""Distance from origin in 3D - - :math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}` - """ - return numpy.sqrt(self.rho2) - - @property - def theta(self): - r"""Inclination angle from XY plane + def r(self): + r"""Distance from origin in XY plane - :math:`\text{arctan2}(r, z)` + :math:`\sqrt{x^2+y^2}` """ - return numpy.arctan2(self.r, self.z) + return self.rho @property - def p2(self): - """Squared `p`""" + def r2(self): + """Squared `r`""" return self.rho2 - @property - def p(self): - """Alias for `rho`""" - return self.rho - @awkward.mixin_class_method(numpy.absolute) def absolute(self): """Returns magnitude of the 3D vector @@ -377,39 +244,13 @@ def absolute(self): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y, "z": -self.z}, - with_name="ThreeVector", - behavior=self.behavior, - ) - - @awkward.mixin_class_method(numpy.add, {"ThreeVector"}) - def add(self, other): - """Add two vectors together elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - {"x": self.x + other.x, "y": self.y + other.y, "z": self.z + other.z}, - with_name="ThreeVector", - behavior=self.behavior, - ) + return self.scale(-1) - @awkward.mixin_class_method( - numpy.subtract, - { - "ThreeVector", - "SphericalThreeVector", - "LorentzVector", - "PtEtaPhiMLorentzVector", - "PtEtaPhiELorentzVector", - }, - transpose=False, - ) - def subtract(self, other): - """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - {"x": self.x - other.x, "y": self.y - other.y, "z": self.z - other.z}, - with_name="ThreeVector", - behavior=self.behavior, - ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.scale(1 / other) def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, and `z` components""" @@ -426,27 +267,14 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - {"x": self.x * other, "y": self.y * other, "z": self.z * other}, - with_name="ThreeVector", - behavior=self.behavior, - ) + return self.scale(other) - def dot(self, other): - """Compute the dot product of two vectors""" - return self.x * other.x + self.y * other.y + self.z * other.z + def delta_phi(self, other): + """Compute difference in angle between two vectors - def cross(self, other): - """Compute the cross product of two vectors""" - return awkward.zip( - { - "x": self.y * other.z - self.z * other.y, - "y": self.z * other.x - self.x * other.z, - "z": self.x * other.y - self.y * other.x, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) + Returns a value within [-pi, pi) + """ + return self.deltaphi(other) @property def unit(self): @@ -455,7 +283,7 @@ def unit(self): @awkward.mixin_class(behavior) -class SphericalThreeVector(ThreeVector, PolarTwoVector): +class SphericalThreeVector(ThreeVector): """A spherical coordinate 3-dimensional vector This mixin class requires the parent class to provide items `rho`, `theta`, and `phi`. @@ -468,70 +296,20 @@ def r(self): :math:`\sqrt{x^2+y^2} = \rho \sin(\theta)` """ - return self.rho * numpy.sin(self.theta) - - @property - def z(self): - r"""Cartesian z value - - :math:`\rho \cos(\theta)` - """ - return self.rho * numpy.cos(self.theta) - - @property - def rho(self): - r"""Distance from origin in 3D - - :math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}` - """ - return self["rho"] - - @property - def theta(self): - r"""Inclination angle from XY plane - - :math:`\text{arctan2}(r, z)` - """ - return self["theta"] - - @property - def p(self): - """Alias for `rho`""" return self.rho - @property - def p2(self): - """Squared `p`""" - return self.rho * self.rho - @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, and `z` components In reality, this directly adjusts `r`, `theta` and `phi` for performance """ - return awkward.zip( - { - "rho": self.rho * abs(other), - "theta": (numpy.sign(other) * self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi * (other < 0), - }, - with_name="SphericalThreeVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - { - "rho": self.rho, - "theta": (-self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi, - }, - with_name="SphericalThreeVector", - behavior=self.behavior, - ) + return self.scale(-1) def _metric_table_core(a, b, axis, metric, return_combinations): @@ -561,7 +339,7 @@ def _nearest_core(x, y, axis, metric, return_metric, threshold): @awkward.mixin_class(behavior) -class LorentzVector(ThreeVector): +class LorentzVector(MomentumAwkward4D): """A cartesian Lorentz vector A heavy emphasis towards a momentum vector interpretation is assumed. @@ -603,34 +381,6 @@ def absolute(self): """ return self.mass - @awkward.mixin_class_method(numpy.add, {"LorentzVector"}) - def add(self, other): - """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - "t": self.t + other.t, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) - - @awkward.mixin_class_method(numpy.subtract, {"LorentzVector"}, transpose=False) - def subtract(self, other): - """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - "t": self.t - other.t, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) - def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` components""" return awkward.zip( @@ -647,36 +397,36 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x * other, - "y": self.y * other, - "z": self.z * other, - "t": self.t * other, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.scale(other) + + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.scale(1 / other) def delta_r2(self, other): """Squared `delta_r`""" - return delta_r(self.eta, self.phi, other.eta, other.phi) ** 2 + return self.deltaR2(other) def delta_r(self, other): r"""Distance between two Lorentz vectors in (eta,phi) plane :math:`\sqrt{\Delta\eta^2 + \Delta\phi^2}` """ - return delta_r(self.eta, self.phi, other.eta, other.phi) + return self.deltaR(other) + + def delta_phi(self, other): + """Compute difference in angle between two vectors + + Returns a value within [-pi, pi) + """ + return self.deltaphi(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y, "z": -self.z, "t": -self.t}, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.scale(-1) @property def pvec(self): @@ -687,53 +437,14 @@ def pvec(self): behavior=self.behavior, ) - @property - def rapidity(self): - pz = self.z - e = self.energy - return 0.5 * (numpy.log(e + pz) - numpy.log(e - pz)) - @property def boostvec(self): """The `x`, `y` and `z` components divided by `t` as a `ThreeVector` - This can be used for boosting. For cases where `|t| <= rho`, this + This can be used for boosting. For cases where `|t| <= r`, this returns the unit vector. """ - rho = self.rho - t = self.t - with numpy.errstate(divide="ignore"): - out = self.pvec * awkward.where( - rho == 0, 0, awkward.where(abs(t) <= rho, 1 / rho, 1 / t) - ) - return out - - def boost(self, other): - """Apply a Lorentz boost given by the `ThreeVector` `other` and return it - - Note that this follows the convention that, for example in order to boost - a vector into its own rest frame, one needs to use the negative of its `boostvec` - """ - b2 = other.rho2 - gamma = (1 - b2) ** (-0.5) - mask = b2 == 0 - b2 = awkward.where(mask, 1, b2) - gamma2 = awkward.where(mask, 0, (gamma - 1) / b2) - - bp = self.dot(other) - t = self.t - v = gamma2 * bp * other + t * gamma * other - - return awkward.zip( - { - "x": self.x + v.x, - "y": self.y + v.y, - "z": self.z + v.z, - "t": gamma * (t + bp), - }, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.to_beta3() @dask_method def metric_table( @@ -821,7 +532,7 @@ def nearest( @awkward.mixin_class(behavior) -class PtEtaPhiMLorentzVector(LorentzVector, SphericalThreeVector): +class PtEtaPhiMLorentzVector(LorentzVector): """A Lorentz vector using pseudorapidity and mass This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, and `mass`. @@ -944,9 +655,15 @@ def negative(self): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + @awkward.mixin_class(behavior) -class PtEtaPhiELorentzVector(LorentzVector, SphericalThreeVector): +class PtEtaPhiELorentzVector(LorentzVector): """A Lorentz vector using pseudorapidity and energy This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, and `energy`. @@ -1060,6 +777,65 @@ def negative(self): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + + +_binary_dispatch_cls = { + "TwoVector": TwoVector, + "PolarTwoVector": TwoVector, + "ThreeVector": ThreeVector, + "SphericalThreeVector": ThreeVector, + "LorentzVector": LorentzVector, + "PtEtaPhiMLorentzVector": LorentzVector, + "PtEtaPhiELorentzVector": LorentzVector, +} +_rank = [TwoVector, ThreeVector, LorentzVector] + +for lhs, lhs_to in _binary_dispatch_cls.items(): + for rhs, rhs_to in _binary_dispatch_cls.items(): + out_to = min(lhs_to, rhs_to, key=_rank.index) + behavior[(numpy.add, lhs, rhs)] = out_to.add + behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract + + +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/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 7fcb7177c..665970ab1 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -1,16 +1,25 @@ import awkward as ak import numpy as np import pytest +from numpy.testing import assert_allclose from coffea.nanoevents.methods import vector ATOL = 1e-8 -def record_arrays_equal(a, b): - return (ak.fields(a) == ak.fields(b)) and all( - ak.all(a[f] == b[f]) for f in ak.fields(a) - ) +def assert_record_arrays_equal(a, b, check_type=False): + if check_type: + assert type(a) is type(b) + assert ak.fields(a) == ak.fields(b) + assert all(ak.all(ak.isclose(a[f], b[f])) for f in ak.fields(a)) + + +def assert_awkward_allclose(actual, desired): + flat_actual = ak.flatten(actual, axis=None) + flat_desired = ak.flatten(desired, axis=None) + # we should check None values, but not used in these tests + assert_allclose(flat_actual, flat_desired) def test_two_vector(): @@ -25,31 +34,31 @@ def test_two_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip({"x": [[-1, -2], [], [-3], [-4]], "y": [[-5, -6], [], [-7], [-8]]}) ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip({"x": [[12, 14], [], [16], [18]], "y": [[20, 22], [], [24], [26]]}), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( {"x": [[-10, -10], [], [-10], [-10]], "y": [[-10, -10], [], [-10], [-10]]} ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip({"x": [[2, 4], [], [6], [8]], "y": [[10, 12], [], [14], [16]]}) ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip({"x": [[0.5, 1], [], [1.5], [2]], "y": [[2.5, 3], [], [3.5], [4]]}), ) - assert record_arrays_equal(a.dot(b), ak.Array([[86, 120], [], [158], [200]])) - assert record_arrays_equal(b.dot(a), ak.Array([[86, 120], [], [158], [200]])) + assert_awkward_allclose(a.dot(b), ak.Array([[86, 120], [], [158], [200]])) + assert_awkward_allclose(b.dot(a), ak.Array([[86, 120], [], [158], [200]])) assert ak.all(abs(a.unit.r - 1) < ATOL) assert ak.all(abs(a.unit.phi - a.phi) < ATOL) @@ -58,18 +67,18 @@ def test_two_vector(): def test_polar_two_vector(): a = ak.zip( { - "r": [[1, 2], [], [3], [4]], + "rho": [[1, 2], [], [3], [4]], "phi": [[0.3, 0.4], [], [0.5], [0.6]], }, with_name="PolarTwoVector", behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, - ak.zip({"r": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), + ak.zip({"rho": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), ) - assert ak.all((a * (-2)).r == [[2, 4], [], [6], [8]]) + assert ak.all((a * (-2)).rho == [[2, 4], [], [6], [8]]) assert ak.all( (a * (-2)).phi - ak.Array( @@ -77,18 +86,18 @@ def test_polar_two_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( - {"r": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} + {"rho": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} ), ) assert ak.all(abs((-a).x + a.x) < ATOL) assert ak.all(abs((-a).y + a.y) < ATOL) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a) - assert ak.all(a.unit.phi == a.phi) + assert ak.all(ak.isclose(a.unit.phi, a.phi)) def test_three_vector(): @@ -111,7 +120,7 @@ def test_three_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip( { @@ -122,7 +131,7 @@ def test_three_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip( { @@ -132,7 +141,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( { @@ -142,7 +151,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( b - a, ak.zip( { @@ -153,7 +162,7 @@ def test_three_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip( { @@ -163,7 +172,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -177,7 +186,7 @@ def test_three_vector(): assert ak.all(a.dot(b) == ak.Array([[170, 154], [], [162], [284]])) assert ak.all(b.dot(a) == ak.Array([[170, 154], [], [162], [284]])) - assert record_arrays_equal( + assert_record_arrays_equal( a.cross(b), ak.zip( { @@ -187,7 +196,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( b.cross(a), ak.zip( { @@ -216,7 +225,7 @@ def test_spherical_three_vector(): assert ak.all(abs((-a).x + a.x) < ATOL) assert ak.all(abs((-a).y + a.y) < ATOL) assert ak.all(abs((-a).z + a.z) < ATOL) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) def test_lorentz_vector(): @@ -241,7 +250,7 @@ def test_lorentz_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip( { @@ -253,7 +262,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip( { @@ -264,7 +273,7 @@ def test_lorentz_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( { @@ -276,7 +285,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip( { @@ -287,7 +296,7 @@ def test_lorentz_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -299,7 +308,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a.pvec, ak.zip( { @@ -344,7 +353,7 @@ def test_pt_eta_phi_m_lorentz_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -355,7 +364,7 @@ def test_pt_eta_phi_m_lorentz_vector(): } ), ) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) boosted = a.boost(-a.boostvec) assert ak.all(abs(boosted.x) < ATOL) @@ -390,7 +399,7 @@ def test_pt_eta_phi_e_lorentz_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -401,7 +410,7 @@ def test_pt_eta_phi_e_lorentz_vector(): } ), ) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) boosted = a.boost(-a.boostvec) assert ak.all(abs(boosted.x) < ATOL) @@ -444,12 +453,18 @@ def test_lorentz_vector_numba(a_dtype, b_dtype): 110.66616465749593, 110.68423555321688, ] - assert pytest.approx(a.delta_phi(b), abs=1e-7) == [ - 0.03369510734601633, - -0.1798534997924781, - 0.33292327383538156, - 0.6078019961139605, - ] + + computed_dphi = a.delta_phi(b).to_numpy() + + assert pytest.approx(computed_dphi, abs=1e-6) == np.array( + [ + 0.03369510734601633, + -0.1798534997924781, + 0.33292327383538156, + 0.6078019961139605, + ], + dtype=computed_dphi.dtype, + ) @pytest.mark.parametrize( @@ -519,22 +534,36 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): ) elif twocoord == "PolarTwoVector": c = ak.zip( - {"r": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]}, + {"rho": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]}, with_name=twocoord, behavior=vector.behavior, ) - assert record_arrays_equal(a + b, b + a) - assert record_arrays_equal(a + c, c + a) - assert record_arrays_equal(b + c, c + b) - - assert record_arrays_equal(a.delta_phi(b), b.delta_phi(a)) - assert record_arrays_equal(a.delta_phi(c), c.delta_phi(a)) - assert record_arrays_equal(b.delta_phi(c), c.delta_phi(b)) - - assert record_arrays_equal(a - b, -(b - a)) - assert record_arrays_equal(a - c, -(c - a)) - assert record_arrays_equal(b - c, -(c - b)) + assert_record_arrays_equal(a.like(b) + b, b + a.like(b), check_type=True) + 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)) + + assert_record_arrays_equal((a.like(b) - b), -(b - a.like(b)), check_type=True) + 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])