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", )