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",
)
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",
)
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x` and `y` components"""
out = awkward.zip(
{
"x": awkward.sum(self.x, axis=axis),
"y": awkward.sum(self.y, axis=axis),
},
with_name="TwoVector",
highlevel=False,
)
return awkward.Array(out, behavior=self.behavior)
def test():
x = np.random.randint(0, 64, 128)
y = np.random.randint(0, 64, 128)
point_1030 = ak.zip({
"x": x,
"y": y
},
behavior=behavior,
with_name="1030_point")
assert isinstance(point_1030.mag2, ak.Array)
point = ak.zip({
"x": x,
"y": y
},
behavior=behavior,
with_name="AVeryVeryVeryLargeClassName")
with pytest.raises(AttributeError):
assert isinstance(point.mag2, ak.Array)
def TLorentz_vector(vec):
vec = ak.zip(
{
"x": vec.x,
"y": vec.y,
"z": vec.z,
"t": vec.t
},
with_name="LorentzVector",
)
return vec
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",
)
def p4(self):
"""4-momentum vector of tracks associated to this SV"""
return awkward.zip(
{
"pt": self["pt"],
"eta": self["eta"],
"phi": self["phi"],
"mass": self["mass"],
},
with_name="PtEtaPhiMLorentzVector",
)
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",
)
def add_systematic(
self,
name: str,
kind: str,
what: Union[str, List[str], Tuple[str]],
varying_function: Callable,
):
"""
name: str, name of the systematic variation / uncertainty source
kind: str, the name of the kind of systematic variation
what: Union[str, List[str], Tuple[str]], name what gets varied,
this could be a list or tuple of column names
varying_function: Union[function, bound method], a function that describes how 'what' is varied, it must close over all non-event-data arguments.
"""
self._ensure_systematics()
if name in awkward.fields(self["__systematics__"]):
raise ValueError(
f"{name} already exists as a systematic for this object!")
if kind not in self._systematic_kinds:
raise ValueError(
f"{kind} is not an available systematics type, please add it and try again!"
)
wrap = partial(awkward_rewrap,
like_what=self["__systematics__"],
gfunc=rewrap_recordarray)
flat = (self
if isinstance(self, coffea.nanoevents.methods.base.NanoEvents)
else awkward.flatten(self))
if what == "weight" and "__ones__" not in awkward.fields(
flat["__systematics__"]):
flat["__systematics__",
"__ones__"] = numpy.ones(len(flat), dtype=numpy.float32)
rendered_type = flat.layout.parameters["__record__"]
as_syst_type = awkward.with_parameter(flat, "__record__", kind)
as_syst_type._build_variations(name, what, varying_function)
variations = as_syst_type.describe_variations()
flat["__systematics__", name] = awkward.zip(
{
v: getattr(as_syst_type, v)(name, what, rendered_type)
for v in variations
},
depth_limit=1,
with_name=f"{name}Systematics",
)
self["__systematics__"] = wrap(flat["__systematics__"])
self.behavior[("__typestr__", f"{name}Systematics")] = f"{kind}"
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` 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),
},
with_name="LorentzVector",
behavior=self.behavior,
)
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,
)
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",
)
def weighted_add(self, other):
sumw = self.weight + other.weight
return ak.zip(
{
"x":
(self.x * self.weight + other.x * other.weight) / sumw,
"y":
(self.y * self.weight + other.y * other.weight) / sumw,
"weight": sumw,
},
with_name="WeightedPoint",
)
def negative(self):
"""Returns the negative of the vector"""
return awkward.zip(
{
"pt": self.pt,
"eta": -self.eta,
"phi": self.phi % (2 * numpy.pi) - numpy.pi,
"energy": -self.energy,
},
with_name="PtEtaPhiELorentzVector",
behavior=self.behavior,
)
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",
)
def Lazy(file_list,branches):
cache = uproot.ArrayCache("2 GB")
tree = uproot.lazyarrays(
file_list,
"Delphes",
branches,
cache=cache,
)
print(len(tree))
Electron = ak.zip(
{
"PT": tree["Electron.PT"],
"Eta": tree["Electron.Eta"],
"Phi": tree["Electron.Phi"],
"Charge": tree["Electron.Charge"],
})
Muon = ak.zip(
{
"PT": tree["MuonLoose.PT"],
"Eta": tree["MuonLoose.Eta"],
"Phi": tree["MuonLoose.Phi"],
"Charge": tree["MuonLoose.Charge"],
})
Photon = ak.zip(
{
"PT": tree["PhotonLoose.PT"],
"Eta": tree["PhotonLoose.Eta"],
"Phi": tree["PhotonLoose.Phi"],
})
MET = ak.zip(
{
"PT": tree["PuppiMissingET.MET"],
"Phi": tree["PuppiMissingET.Phi"],
})
def Loop(file_list):
# define array
histo={}
# --Start File Loop
for arrays in uproot.iterate(flist,branches): # for Uproot4
#for arrays in uproot.iterate(flist,'Delphes',branches): # for Uproot3
print(len(arrays))
Electron = ak.zip(
{
"PT": arrays[b"Electron.PT"],
"Eta": arrays[b"Electron.Eta"],
"Phi": arrays[b"Electron.Phi"],
"Charge": arrays[b"Electron.Charge"],
})
Muon = ak.zip(
{
"PT": arrays[b"MuonLoose.PT"],
"Eta": arrays[b"MuonLoose.Eta"],
"Phi": arrays[b"MuonLoose.Phi"],
"Charge": arrays[b"MuonLoose.Charge"],
})
Photon = ak.zip(
{
"PT": arrays[b"PhotonLoose.PT"],
"Eta": arrays[b"PhotonLoose.Eta"],
"Phi": arrays[b"PhotonLoose.Phi"],
})
MET = ak.zip(
{
"PT": arrays[b"PuppiMissingET.MET"],
"Phi": arrays[b"PuppiMissingET.Phi"],
})
print(Muon.PT)
def test_lorentz_vector_numba(a_dtype, b_dtype):
a = ak.zip(
{
"x": np.array([1, 2, 3, 4], dtype=a_dtype),
"y": np.array([5, 6, 7, 8], dtype=a_dtype),
"z": np.array([9, 10, 11, 12], dtype=a_dtype),
"t": np.array([50, 51, 52, 53], dtype=b_dtype), # b on purpose
},
with_name="LorentzVector",
behavior=vector.behavior,
)
b = ak.zip(
{
"x": np.array([4, 1, 10, 11], dtype=b_dtype),
"y": np.array([17, 7, 11, 6], dtype=b_dtype),
"z": np.array([9, 11, 5, 16], dtype=b_dtype),
"t": np.array([60, 61, 62, 63], dtype=b_dtype),
},
with_name="LorentzVector",
behavior=vector.behavior,
)
assert pytest.approx(a.mass) == [
48.91829923454004,
49.60846701924985,
50.24937810560445,
50.84289527554464,
]
assert pytest.approx((a + b).mass) == [
106.14612569472331,
109.20164833920778,
110.66616465749593,
110.68423555321688,
]
assert pytest.approx(a.delta_phi(b), abs=1e-7) == [
0.03369510734601633,
-0.1798534997924781,
0.33292327383538156,
0.6078019961139605,
]
def corrected_polar_met(met_pt, met_phi, jet_pt, jet_phi, jet_pt_orig, deltas=None):
sj, cj = numpy.sin(jet_phi), numpy.cos(jet_phi)
x = met_pt * numpy.cos(met_phi) + awkward.sum(
jet_pt * cj - jet_pt_orig * cj, axis=1
)
y = met_pt * numpy.sin(met_phi) + awkward.sum(
jet_pt * sj - jet_pt_orig * sj, axis=1
)
if deltas:
positive, dx, dy = deltas
x = x + dx if positive else x - dx
y = y + dy if positive else y - dy
return awkward.zip({"pt": numpy.hypot(x, y), "phi": numpy.arctan2(y, x)})
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",
)
def TLorentz_vector_cylinder(vec):
vec = ak.zip(
{
"pt": vec.pt,
"eta": vec.eta,
"phi": vec.phi,
"mass": vec.mass,
},
with_name="PtEtaPhiMLorentzVector",
)
return vec
def test_zip():
x = ak._v2.Array([1.1, 2.2, 3.3])
y = ak._v2.Array(["one", "two", "three"])
assert ak.zip({
"x": x,
"y": y
}).tolist() == [
{
"x": 1.1,
"y": "one"
},
{
"x": 2.2,
"y": "two"
},
{
"x": 3.3,
"y": "three"
},
]
y = ak._v2.behaviors.categorical.to_categorical(y)
assert ak.zip({
"x": x,
"y": y
}).tolist() == [
{
"x": 1.1,
"y": "one"
},
{
"x": 2.2,
"y": "two"
},
{
"x": 3.3,
"y": "three"
},
]
def multiply(self, other):
"""Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components
In reality, this multiplies `pt` and `energy` by the scalar quanitity for performance
"""
return awkward.zip(
{
"pt": self.pt * other,
"eta": self.eta,
"phi": self.phi,
"energy": self.energy * other,
},
with_name="PtEtaPhiELorentzVector",
)
def multiply(self, other):
"""Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components
In reality, this directly adjusts `pt`, `eta`, `phi` and `energy` for performance
"""
return awkward.zip(
{
"pt": self.pt * abs(other),
"eta": self.eta * numpy.sign(other),
"phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)),
"energy": self.energy * other,
},
with_name="PtEtaPhiELorentzVector",
)
def getData(fname="", procName="Truth"):
dq_dict = uproot.open(fname)["Truth"]
dq_events = ak.zip(
{
"Electrons":
ak.zip({
"ge": dq_dict["ge"].array(),
"gvx": dq_dict["gvx"].array(),
"gvy": dq_dict["gvy"].array(),
"gvz": dq_dict["gvz"].array(),
"gpt": dq_dict["gpt"].array(),
}),
"Showers":
ak.zip({
"nshowers": dq_dict["n_showers"].array(),
"sedep_ecal": dq_dict["sedep_ecal"].array(),
"sx": dq_dict["sx_ecal"].array(),
"sy": dq_dict["sy_ecal"].array(),
"sz": dq_dict["sz_ecal"].array(),
}),
},
depth_limit=1)
return dq_events
def fill_lepton_kinematics():
import uproot
import awkward as ak
from coffea.nanoevents.methods import candidate
ak.behavior.update(candidate.behavior)
# histogram creation and manipulation
from coffea import hist
fin = uproot.open("HZZ.root")
tree = fin["events"]
arrays = {k.replace('Electron_', '') \
.strip('P') \
.replace('E','t').lower(): v for k, v in tree.arrays(filter_name="Electron_*",
how=dict).items()}
electrons = ak.zip(arrays, with_name='Candidate')
arrays = {k.replace('Muon_', '') \
.strip('P') \
.replace('E','t').lower(): v for k, v in tree.arrays(filter_name="Muon_*",
how=dict).items()}
muons = ak.zip(arrays, with_name='Candidate')
# Two types of axes exist presently: bins and categories
lepton_kinematics = hist.Hist("Events",
hist.Cat("flavor", "Lepton flavor"),
hist.Bin("pt", "$p_{T}$", 19, 10, 100),
hist.Bin("eta", r"$\eta$", [-2.5, -1.4, 0, 1.4, 2.5]),
)
# Pass keyword arguments to fill, all arrays must be flat numpy arrays
# User is responsible for ensuring all arrays have same jagged structure!
lepton_kinematics.fill(flavor="electron", pt=ak.flatten(electrons.pt), eta=ak.flatten(electrons.eta))
lepton_kinematics.fill(flavor="muon", pt=ak.flatten(muons.pt), eta=ak.flatten(muons.eta))
return lepton_kinematics
def jet_features(jets, mask_bool=False, mask=None):
vecs = ak.zip(
{
"pt": jets[:, :, 2:3],
"eta": jets[:, :, 0:1],
"phi": jets[:, :, 1:2],
"mass": ak.full_like(jets[:, :, 2:3], 0),
},
with_name="PtEtaPhiMLorentzVector")
sum_vecs = vecs.sum(axis=1)
jf = np.nan_to_num(
np.array(np.concatenate((sum_vecs.mass, sum_vecs.pt), axis=1)))
return jf
def test_spherical_three_vector():
a = ak.zip(
{
"rho": [[1.0, 2.0], [], [3.0], [4.0]],
"theta": [[1.2, 0.7], [], [1.8], [1.9]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]],
},
with_name="SphericalThreeVector",
highlevel=False,
)
a = ak.Array(a, behavior=vector.behavior)
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)
def lazy_variant(unc, metpt, metphi, jetpt, jetphi, jetptraw):
return awkward.zip(
{
"up": make_variant(
MET[metpt],
MET[metphi],
corrected_jets[unc].up[jetpt],
corrected_jets[unc].up[jetphi],
corrected_jets[unc].up[jetptraw],
),
"down": make_variant(
MET[metpt],
MET[metphi],
corrected_jets[unc].down[jetpt],
corrected_jets[unc].down[jetphi],
corrected_jets[unc].down[jetptraw],
),
},
depth_limit=1,
with_name="METSystematic",
)
def get_variation(self, name, what, astype, updown):
"""Calculate and up or down variation."""
fields = awkward.fields(self)
fields.remove("__systematics__")
varied = self["__systematics__", f"__{name}__", :, self._udmap[updown]]
params = copy(self.layout.parameters)
params["variation"] = f"{name}-{what}-{updown}"
out = {field: self[field] for field in fields}
if what == "weight":
out[f"weight_{name}"] = varied
else:
out[what] = varied
return awkward.zip(
out,
depth_limit=1,
parameters=params,
behavior=self.behavior,
with_name=astype,
)