def test_issue_334():
a = awkward1.Array([1, 2, 3, 4])
b = awkward1.Array([-1])
c = awkward1.Array([True, False, True, True])
assert awkward1.where(c, a, b).tolist() == [1, -1, 3, 4]
assert awkward1.where(*awkward1.broadcast_arrays(c, a, b)).tolist() == [1, -1, 3, 4]
assert awkward1.where(c, a, -1).tolist() == [1, -1, 3, 4]
assert awkward1.where(*awkward1.broadcast_arrays(c, a, -1)).tolist() == [1, -1, 3, 4]
def test_broadcast_arrays():
a = awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True)
b = awkward1.Array([100, 200, 300], check_valid=True)
out = awkward1.broadcast_arrays(a, b)
assert awkward1.to_list(out[0]) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
assert awkward1.to_list(out[1]) == [[100, 100, 100], [], [300, 300]]
def apply_global_index(index, target):
def flat_take(layout):
idx = awkward1.Array(layout)
return target._content()[idx.mask[idx >= 0]]
def descend(layout, depth):
if layout.purelist_depth == 1:
return lambda: flat_take(layout)
(index, ) = awkward1.broadcast_arrays(index)
out = awkward1._util.recursively_apply(index.layout, descend)
return awkward1._util.wrap(out, target.behavior)
def expr_to_vals(expr):
vals = eval(expr, dict(), loc)
# if varexp is a simple constant, broadcast it to an array
if _array_ndim(vals) == 0:
vals = vals * np.ones(len(df))
if sel:
if _array_ndim(vals) < _array_ndim(globalmask):
vals, _ = awkward1.broadcast_arrays(vals, globalmask)
vals = vals[globalmask]
if _array_ndim(vals) > 1:
vals = awkward1.flatten(vals)
vals = awkward1.to_numpy(vals)
return vals
def _apply_global_index(self, index):
"""Internal method to take from a collection using a flat index
This is often necessary to be able to still resolve cross-references on
reduced arrays or single records.
"""
if isinstance(index, int):
out = self._content()[index]
else:
def flat_take(layout):
idx = awkward1.Array(layout)
return self._content()[idx.mask[idx >= 0]]
def descend(layout, depth):
if layout.purelist_depth == 1:
return lambda: flat_take(layout)
(index, ) = awkward1.broadcast_arrays(index)
out = awkward1._util.recursively_apply(index.layout, descend)
return awkward1._util.wrap(out, self.behavior, cache=self.cache)
def mass(self):
return awkward1.broadcast_arrays(self.pt, 0.0)[1]
j1j2mass = j1.mass + j2.mass
j2j3mass = j2.mass + j3.mass
j1j3mass = j1.mass + j3.mass
j1pt = j1.pt
j1phi = j1.phi
j1eta = abs(j1.eta)
j1mass = j1.mass
j2pt = j2.pt
j2phi = j2.phi
j2eta = abs(j2.eta)
j2mass = j2.mass
j3pt = j3.pt
j3phi = j3.phi
j3eta = abs(j3.eta)
j3mass = j3.mass
event = ak.flatten(ak.broadcast_arrays(final.event, combs['0'].pt)[0])
processedMLdata = ak.zip({
"j1pt": j1.pt,
"j1phi": j1.phi,
"j1eta": j1.eta,
"j1mass": j1.mass,
"j2pt": j2.pt,
"j2phi": j2.phi,
"j2eta": j2.eta,
"j2mass": j2.mass,
"j3pt": j3.pt,
"j3phi": j3.phi,
"j3eta": j3.eta,
"j3mass": j3.mass,
"dR12": dR1_2,
"dR13": dR1_3,
def awkwardReshape(akArray, npArray):
return ak.broadcast_arrays(akArray.pt, 1.0)[1] * npArray