refactor!: back nanoevents.methods.vector with scikit-hep vector

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",

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

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"]

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,41 +199,83 @@ 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): ...
 
 
 _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): ...
 
 
 _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): ...
 
 
 _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): ...
 
 
 _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): ...
 
 
 _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",

src/coffea/nanoevents/methods/nanoaod.py

@@ -31,13 +31,31 @@ 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"""
 
     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):