"""Physics object candidate mixin
This provides just a Lorentz vector with charge, but maybe
in the future it will provide some sort of composite candidate building tool
that automatically resolves duplicates in the chain.
"""
import awkward
import numpy as np
from awkward_zipper.behaviors import vector
behavior = dict(vector.behavior)
[docs]
@awkward.mixin_class(behavior)
class Candidate(vector.LorentzVector):
"""A Lorentz vector with charge
This mixin class requires the parent class to provide items ``x``, ``y``, ``z``, ``t``, and ``charge``.
"""
[docs]
@awkward.mixin_class_method(np.add, {"Candidate"})
def add(self, other):
"""Add two candidates together elementwise using ``x``, ``y``, ``z``, ``t``, and ``charge`` components"""
return awkward.zip(
{
"x": self.x + other.x,
"y": self.y + other.y,
"z": self.z + other.z,
"t": self.t + other.t,
"charge": self.charge + other.charge,
},
with_name="Candidate",
behavior=self.behavior,
)
[docs]
@awkward.mixin_class_method(np.subtract, {"Candidate"})
def subtract(self, other):
"""Subtract a candidate from another elementwise using ``x``, ``y``, ``z``, ``t``, and ``charge`` 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", # subtraction only makes sense for raw Lorentz vectors
behavior=self.behavior,
)
[docs]
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using ``x``, ``y``, ``z``, ``t``, and ``charge`` components"""
return awkward.zip(
{
"x": awkward.sum(self.x, axis=axis),
"y": awkward.sum(self.y, axis=axis),
"z": awkward.sum(self.z, axis=axis),
"t": awkward.sum(self.t, axis=axis),
"charge": awkward.sum(self.charge, axis=axis),
},
with_name="Candidate",
behavior=self.behavior,
)
def __awkward_validation__(self):
if "charge" not in self.fields:
msg = f"{type(self).__name__} requires the 'charge' field"
raise ValueError(msg)
parent = super()
if hasattr(parent, "__awkward_validation__"):
parent.__awkward_validation__()
# Copy the cross-class LorentzVector behaviors (e.g. Candidate + TwoVector) onto
# Candidate, but only for keys the ``@mixin_class`` decorator has not already
# registered. This MUST run after the decorator: ``copy_behaviors`` would
# otherwise pre-register ``(add, Candidate, Candidate)`` -> LorentzVector.add via
# the ``setdefault`` used by ``mixin_class``, shadowing Candidate's own
# charge-propagating ``add`` (see scikit-hep/coffea#1578).
for _key, _value in awkward._util.copy_behaviors(
"LorentzVector", "Candidate", behavior
).items():
behavior.setdefault(_key, _value)
del _key, _value
[docs]
@awkward.mixin_class(behavior)
class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMLorentzVector):
"""A Lorentz vector in eta, mass coordinates with charge
This mixin class requires the parent class to provide items ``pt``, ``eta``, ``phi``, ``mass``, and ``charge``.
"""
[docs]
@awkward.mixin_class(behavior)
class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector):
"""A Lorentz vector in eta, energy coordinates with charge
This mixin class requires the parent class to provide items ``pt``, ``eta``, ``phi``, ``energy``, and ``charge``.
"""
CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
CandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821
CandidateArray.MomentumClass = CandidateArray # noqa: F821
CandidateRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
CandidateRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
CandidateRecord.ProjectionClass4D = vector.LorentzVectorRecord # noqa: F821
CandidateRecord.MomentumClass = CandidateRecord # 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
PtEtaPhiMCandidateRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
PtEtaPhiMCandidateRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
PtEtaPhiMCandidateRecord.ProjectionClass4D = vector.LorentzVectorRecord # noqa: F821
PtEtaPhiMCandidateRecord.MomentumClass = PtEtaPhiMCandidateRecord # 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
PtEtaPhiECandidateRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
PtEtaPhiECandidateRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
PtEtaPhiECandidateRecord.ProjectionClass4D = vector.LorentzVectorRecord # noqa: F821
PtEtaPhiECandidateRecord.MomentumClass = PtEtaPhiECandidateRecord # noqa: F821
__all__ = ["Candidate", "PtEtaPhiECandidate", "PtEtaPhiMCandidate"]