/
candidate.py
73 lines (54 loc) · 2.23 KB
/
candidate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
"""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
from coffea.nanoevents.methods import vector
behavior = dict(vector.behavior)
@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`.
"""
@awkward.mixin_class_method(numpy.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,
)
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,
)
@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`.
"""
pass
@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`.
"""
pass
__all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"]