def test_zip():
x = awkward1.Array([1.1, 2.2, 3.3])
y = awkward1.Array(["one", "two", "three"])
assert awkward1.zip({
"x": x,
"y": y
}).tolist() == [{
"x": 1.1,
"y": "one"
}, {
"x": 2.2,
"y": "two"
}, {
"x": 3.3,
"y": "three"
}]
y = awkward1.to_categorical(y)
assert awkward1.zip({
"x": x,
"y": y
}).tolist() == [{
"x": 1.1,
"y": "one"
}, {
"x": 2.2,
"y": "two"
}, {
"x": 3.3,
"y": "three"
}]
def test_virtual_slice_numba():
numba = pytest.importorskip("numba")
materialize_count = defaultdict(int)
def gen(x, name):
materialize_count[name] += 1
return x
x1 = awkward1.Array([1, 2, 3, 4, 5])
x2 = awkward1.Array([1, 2, 3, 4, 5])
x = awkward1.zip({"x1": x1, "x2": x2}, with_name="xthing")
xv = awkward1.virtual(lambda: gen(x, "x"),
length=len(x),
form=x.layout.form)
y = x1 * 10.
yv = awkward1.virtual(lambda: gen(y, "y"),
length=len(y),
form=y.layout.form)
array = awkward1.zip({"x": xv, "y": yv}, with_name="Point", depth_limit=1)
virtual = awkward1.virtual(lambda: gen(array, "array"),
length=len(array),
form=array.layout.form)
@numba.njit
def dostuff(array):
x = 0
for item in array:
x += item
return x
assert dostuff(virtual.x.x1) == 15
assert dict(materialize_count) == {"x": 3, "array": 3}
materialize_count.clear()
@numba.njit
def dostuff(array):
x = 0
for item in array:
x += item.x.x1
return x
assert dostuff(virtual[["x"]]) == 15
assert dict(materialize_count) == {"x": 1, "array": 2}
materialize_count.clear()
assert dostuff(virtual[::2]) == 9
assert dict(materialize_count) == {"x": 1, "array": 2}
materialize_count.clear()
assert dostuff(virtual[:3]) == 6
assert dict(materialize_count) == {"x": 1, "array": 1}
materialize_count.clear()
slicedvirtual = virtual[awkward1.Array([True, False, False, True, False])]
assert dostuff(slicedvirtual) == 5
assert dict(materialize_count) == {"x": 1, "array": 2}
materialize_count.clear()
def get_vec(var, mass):
dict1 = {}
dict2 = {}
for key, val in key_dict.items():
dict1[key] = evtDict["HH4V"]['genHiggs1' + var + val]
dict2[key] = evtDict["HH4V"]['genHiggs2' + var + val]
dict1['mass'] = np.ones(len(evtDict["HH4V"]['genHiggs1' + var + val])) * mass
dict2['mass'] = np.ones(len(evtDict["HH4V"]['genHiggs2' + var + val])) * mass
return (ak.zip(dict1, with_name="PtEtaPhiMLorentzVector"), ak.zip(dict2, with_name="PtEtaPhiMLorentzVector"))
def test_virtual_record():
materialize_count = defaultdict(int)
def gen(x, name):
materialize_count[name] += 1
return x
x1 = awkward1.Array([1, 2, 3, 4, 5])
x2 = awkward1.Array([1, 2, 3, 4, 5])
x = awkward1.zip({"x1": x1, "x2": x2}, with_name="xthing")
assert x.layout.purelist_parameter("__record__") == "xthing"
xv = awkward1.virtual(lambda: gen(x, "x"),
length=len(x),
form=x.layout.form)
assert xv.layout.purelist_parameter("__record__") == "xthing"
y = x1 * 10.
yv = awkward1.virtual(lambda: gen(y, "y"),
length=len(y),
form=y.layout.form)
array = awkward1.zip({"x": xv, "y": yv}, with_name="Point", depth_limit=1)
assert array.layout.purelist_parameter("__record__") == "Point"
virtual = awkward1.virtual(lambda: gen(array, "array"),
length=len(array),
form=array.layout.form)
assert virtual.layout.purelist_parameter("__record__") == "Point"
assert len(materialize_count) == 0
slicedvirtual = virtual.x
assert len(materialize_count) == 0
assert slicedvirtual.layout.purelist_parameter("__record__") == "xthing"
assert len(materialize_count) == 0
slicedvirtual = virtual[["x", "y"]]
assert len(materialize_count) == 0
assert slicedvirtual.layout.purelist_parameter("__record__") == None
assert len(materialize_count) == 0
slicedvirtual = virtual[::2]
assert len(materialize_count) == 0
assert slicedvirtual.layout.purelist_parameter("__record__") == "Point"
assert len(materialize_count) == 0
slicedvirtual = virtual[:3]
assert len(materialize_count) == 0
assert slicedvirtual.layout.purelist_parameter("__record__") == "Point"
assert len(materialize_count) == 0
slicedvirtual = virtual[awkward1.Array([True, False, False, True, False])]
assert len(materialize_count) == 0
assert slicedvirtual.layout.purelist_parameter("__record__") == "Point"
assert len(materialize_count) == 0
def test_cache_chain():
cache1 = {}
one = awkward1.virtual(lambda: [[1.1, 2.2, 3.3], [], [4.4, 5.5]],
cache=cache1,
length=3)
two = awkward1.virtual(lambda: [100, 200, 300], cache=cache1, length=3)
array1 = awkward1.zip({"x": one, "y": two}, depth_limit=1)
assert len(cache1) == 0
cache2 = {}
array2 = awkward1.with_cache(array1, cache2)
assert awkward1.to_list(array2["x"]) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
assert len(cache1) == 0
assert len(cache2) == 1
array3 = awkward1.with_cache(array2, cache1, chain="first")
assert awkward1.to_list(array3) == [{
"x": [1.1, 2.2, 3.3],
"y": 100
}, {
"x": [],
"y": 200
}, {
"x": [4.4, 5.5],
"y": 300
}]
assert len(cache1) == 1
assert len(cache2) == 1
def test():
array = awkward1.Array([[[0.0, 1.1, 2.2], []], [[3.3, 4.4]], [],
[[5.5], [], [6.6, 7.7, 8.8, 9.9]]])
assert awkward1.to_list(awkward1.local_index(array,
axis=0)) == [0, 1, 2, 3]
assert awkward1.to_list(awkward1.local_index(array, axis=1)) == [[0, 1],
[0], [],
[0, 1, 2]]
assert awkward1.to_list(awkward1.local_index(array,
axis=2)) == [[[0, 1, 2], []],
[[0, 1]], [],
[[0], [],
[0, 1, 2, 3]]]
assert awkward1.to_list(awkward1.local_index(array,
axis=-1)) == [[[0, 1, 2], []],
[[0, 1]], [],
[[0], [],
[0, 1, 2, 3]]]
assert awkward1.to_list(awkward1.local_index(array,
axis=-2)) == [[0, 1], [0], [],
[0, 1, 2]]
assert awkward1.to_list(awkward1.local_index(array,
axis=-3)) == [0, 1, 2, 3]
assert awkward1.to_list(
awkward1.zip([
awkward1.local_index(array, axis=0),
awkward1.local_index(array, axis=1),
awkward1.local_index(array, axis=2)
])) == [[[(0, 0, 0), (0, 0, 1), (0, 0, 2)], []],
[[(1, 0, 0), (1, 0, 1)]], [],
[[(3, 0, 0)], [], [(3, 2, 0), (3, 2, 1), (3, 2, 2),
(3, 2, 3)]]]
def test_broken():
ex = awkward1.Array([[1, 2, 3], [], [4, 5]])
p4 = awkward1.zip({"x": ex})
p4c = awkward1.cartesian({"a": p4, "b": p4})
df = awkward1.to_pandas(p4c)
assert df["a", "x"].values.tolist() == [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 5, 5]
assert df["b", "x"].values.tolist() == [1, 2, 3, 1, 2, 3, 1, 2, 3, 4, 5, 4, 5]
def get_vec(var, mass):
dict = {}
for key, val in key_dict.items():
dict[key] = ak.flatten(ak.Array([evtDict["HH4V"]['genHiggs1' + var + val][h1W], evtDict["HH4V"]['genHiggs2' + var + val][h2W]]))
dict['mass'] = np.ones(totW) * mass
return ak.zip(dict, with_name="PtEtaPhiMLorentzVector")
def zip_rle(output, dataset):
return ak.to_numpy(
ak.zip([
output['%s_run' % dataset].value.astype(int),
output['%s_lumi' % dataset].value.astype(int),
output['%s_event' % dataset].value.astype(int),
]))
def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata['dataset']
muons = ak.zip(
{
"pt": events.Muon_pt,
"eta": events.Muon_eta,
"phi": events.Muon_phi,
"mass": events.Muon_mass,
"charge": events.Muon_charge,
},
with_name="PtEtaPhiMCandidate")
cut = (ak.num(muons) == 2) & (ak.sum(muons.charge) == 0)
# add first and second muon in every event together
dimuon = muons[cut][:, 0] + muons[cut][:, 1]
output["sumw"][dataset] += len(events)
output["mass"].fill(
dataset=dataset,
mass=dimuon.mass,
)
return output
def prod(self, other):
return awkward1.zip(
{
"x": self.x * other,
"y": self.y * other
},
with_name="TwoVector",
)
def add(self, other):
return awkward1.zip(
{
"x": self.x + other.x,
"y": self.y + other.y
},
with_name="TwoVector",
)
def sum(self, axis=-1):
return awkward1.zip(
{
"x": awkward1.sum(self.x, axis=axis),
"y": awkward1.sum(self.y, axis=axis),
},
with_name="TwoVector",
)
def point_add(self, other):
print("hi")
return awkward1.zip(
{
"x": self.x + other.x,
"y": self.y + other.y
},
with_name="Point",
)
def add(self, other):
"""Add two vectors together elementwise using `x` and `y` components"""
return awkward1.zip(
{
"x": self.x + other.x,
"y": self.y + other.y
},
with_name="TwoVector",
)
def multiply(self, other):
"""Multiply this vector by a scalar elementwise using `x` and `y` components"""
return awkward1.zip(
{
"x": self.x * other,
"y": self.y * other
},
with_name="TwoVector",
)
def add(self, other):
return awkward1.zip(
{
"x": self.x + other.x,
"y": self.y + other.y,
"z": self.z + other.z
},
with_name="ThreeVector",
)
def test_0230():
rec = awkward1.zip(
{
"x": awkward1.virtual(lambda: awkward1.Array([1, 2, 3, 4]),
length=4),
},
depth_limit=1)
assert awkward1.to_list(rec.x[1:]) == [2, 3, 4]
assert awkward1.to_list(rec.x[1:] * 2) == [4, 6, 8]
def add(self, other):
"""Add two vectors together elementwise using `x`, `y`, and `z` components"""
return awkward1.zip(
{
"x": self.x + other.x,
"y": self.y + other.y,
"z": self.z + other.z
},
with_name="ThreeVector",
)
def prod(self, other):
return awkward1.zip(
{
"pt": self.pt * other,
"eta": self.eta,
"phi": self.phi,
"energy": self.energy * other,
},
with_name="PtEtaPhiELorentzVector",
)
def full_events_to_only_necessary_columns_pt_eta_phi(
arrays: ak.Array) -> ak.Array:
return ak.zip(
{
"particle_ID": arrays["particle_ID"],
"status": arrays["status"],
"pt": np.sqrt(arrays["px"]**2 + arrays["py"]**2),
"eta": arrays["eta"],
"phi": arrays["phi"],
}, )
def getPtEtaPhi(coll, pt_var='pt', eta_var='eta', phi_var='phi'):
return ak.zip({
'pt': getattr(coll, pt_var),
'eta': getattr(coll, eta_var),
'phi': getattr(coll, phi_var),
'p': coll.p, # this is uncorrected....
'btagDeepFlavB': coll.btagDeepFlavB,
'jetId': coll.jetId,
'puId': coll.puId,
})
def prod(self, other):
return awkward1.zip(
{
"x": self.x * other,
"y": self.y * other,
"z": self.z * other,
"t": self.t * other,
},
with_name="LorentzVector",
)
def sum(self, axis=-1):
return awkward1.zip(
{
"x": awkward1.sum(self.x, axis=axis),
"y": awkward1.sum(self.y, axis=axis),
"z": awkward1.sum(self.z, axis=axis),
"t": awkward1.sum(self.t, axis=axis),
},
with_name="LorentzVector",
)
def add(self, other):
return awkward1.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 multiply(self, other):
"""Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components"""
return awkward1.zip(
{
"x": self.x * other,
"y": self.y * other,
"z": self.z * other,
"t": self.t * other,
},
with_name="LorentzVector",
)
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x` and `y` components"""
out = awkward1.zip(
{
"x": awkward1.sum(self.x, axis=axis),
"y": awkward1.sum(self.y, axis=axis),
},
with_name="TwoVector",
highlevel=False,
)
return awkward1._util.wrap(out, cache=self.cache, behavior=self.behavior)
def full_events_to_only_necessary_columns_E_px_py_pz(
arrays: ak.Array) -> ak.Array:
return ak.zip(
{
"particle_ID": arrays["particle_ID"],
"status": arrays["status"],
"E": arrays["E"],
"px": arrays["px"],
"py": arrays["py"],
"pz": arrays["pz"],
}, )
def add(self, other):
return awkward1.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 sum(self, axis=-1):
return awkward1.zip(
{
"x": awkward1.sum(self.x, axis=axis),
"y": awkward1.sum(self.y, axis=axis),
"z": awkward1.sum(self.z, axis=axis),
"t": awkward1.sum(self.t, axis=axis),
"charge": awkward1.sum(self.charge, axis=axis),
},
with_name="Candidate",
)
