def test_NumpyArray(): a = ak._v2.contents.RegularArray( v1_to_v2(ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout), 3) assert ak.to_list(a[1]) == [ [15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29], ] assert ak.to_list(a[1, -2]) == [20, 21, 22, 23, 24] assert a[1, -2, 2] == 22 with pytest.raises(IndexError): a[1, -2, 2, 0] assert ak.to_list(a[1, -2, 2:]) == [22, 23, 24] with pytest.raises(IndexError): a[1, -2, 2:, 0] with pytest.raises(IndexError): a[1, -2, "hello"] with pytest.raises(IndexError): a[1, -2, ["hello", "there"]] assert ak.to_list(a[1, -2, np.newaxis, 2]) == [22] assert ak.to_list(a[1, -2, np.newaxis, np.newaxis, 2]) == [[22]] assert ak.to_list(a[1, -2, ...]) == [20, 21, 22, 23, 24] assert a[1, -2, ..., 2] == 22 with pytest.raises(IndexError): a[1, -2, ..., 2, 2] b = ak.layout.RegularArray( ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout, 3) assert ak.to_list(b[1, -2, [3, 1, 1, 2]]) == [23, 21, 21, 22] assert ak.to_list(a[1, -2, [3, 1, 1, 2]]) == [23, 21, 21, 22] with pytest.raises(IndexError): a[1, -2, [3, 1, 1, 2], 2]
def set_genZ(events, genBranches, selection_options, debug): if genBranches is None: events["genZ_decayMode"] = ak.from_numpy(-1 * numpy.ones(len(events))) else: electron_idxs = abs(genBranches.pdgId) == 11 muon_idxs = abs(genBranches.pdgId) == 13 tau_idxs = abs(genBranches.pdgId) == 15 motherOfElectrons = genBranches.genPartIdxMother[electron_idxs] motherOfMuons = genBranches.genPartIdxMother[muon_idxs] motherOfTaus = genBranches.genPartIdxMother[tau_idxs] ZToEleEvents = ak.sum( (genBranches.pdgId[motherOfElectrons] == 23), axis=1) > 0 ZToMuEvents = ak.sum( (genBranches.pdgId[motherOfMuons] == 23), axis=1) > 0 ZToTauEvents = ak.sum( (genBranches.pdgId[motherOfTaus] == 23), axis=1) > 0 # W decay dudes WToEEvents = ak.sum( (abs(genBranches.pdgId[motherOfElectrons]) == 24), axis=1) > 0 WToMuEvents = ak.sum( (abs(genBranches.pdgId[motherOfMuons]) == 24), axis=1) > 0 WToTauEvents = ak.sum( (abs(genBranches.pdgId[motherOfTaus]) == 24), axis=1) > 0 events["genZ_decayMode"] = ak.from_numpy( numpy.zeros(len(events)) ) + 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents + 4 * WToEEvents + 5 * WToMuEvents + 6 * WToTauEvents events["tau_motherID"] = ak.fill_none( ak.firsts(genBranches.pdgId[motherOfTaus]), 0) return events
def test_ArrayBuilder_append_2(): # issue #415 A = ak.from_numpy(np.array([0, 1, 2], dtype=np.float32)) B = ak.from_numpy(np.array([0, 1], dtype=np.float32)) builder = ak.ArrayBuilder() with builder.list(): builder.append(A.tolist()) with builder.list(): builder.append(A.tolist()) with builder.list(): pass with builder.list(): builder.append(B.tolist()) assert builder.snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [], [[0, 1]]] assert str(ak.type(builder.snapshot())) == "4 * var * var * float64" builder = ak.ArrayBuilder() with builder.list(): builder.append(A) with builder.list(): builder.append(A) with builder.list(): pass with builder.list(): builder.append(B) assert builder.snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [], [[0, 1]]] assert str(ak.type(builder.snapshot())) == "4 * var * var * float32" @numba.njit def f1(builder, A, B): builder.begin_list() builder.append(A) builder.end_list() builder.begin_list() builder.append(A) builder.end_list() builder.begin_list() builder.end_list() builder.begin_list() builder.append(B) builder.end_list() return builder assert f1(ak.ArrayBuilder(), A, B).snapshot().tolist() == [ [[0, 1, 2]], [[0, 1, 2]], [], [[0, 1]], ]
def test_fromnumpy(): assert ak.to_list(ak.from_numpy(np.array(["uno", "dos", "tres", "quatro"]))) == [ "uno", "dos", "tres", "quatro", ] assert ak.to_list( ak.from_numpy(np.array([["uno", "dos"], ["tres", "quatro"]]))) == [["uno", "dos"], ["tres", "quatro"]] assert ak.to_list( ak.from_numpy(np.array([["uno", "dos"], ["tres", "quatro"]]), regulararray=True)) == [["uno", "dos"], ["tres", "quatro"]]
def test_localindex(): array = ak.from_iter( [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False) assert ak.to_list(array.localindex(0)) == [0, 1, 2, 3, 4] assert ak.to_list(array.localindex(1)) == [[0, 1, 2], [], [0, 1], [0], [0, 1, 2, 3]] array = ak.from_iter( [[[0.0, 1.1, 2.2], [], [3.3, 4.4]], [], [[5.5]], [[6.6, 7.7, 8.8, 9.9]] ], highlevel=False, ) assert ak.to_list(array.localindex(0)) == [0, 1, 2, 3] assert ak.to_list(array.localindex(1)) == [[0, 1, 2], [], [0], [0]] assert ak.to_list(array.localindex(2)) == [ [[0, 1, 2], [], [0, 1]], [], [[0]], [[0, 1, 2, 3]], ] array = ak.from_numpy(np.arange(2 * 3 * 5).reshape(2, 3, 5), regulararray=True, highlevel=False) assert ak.to_list(array.localindex(0)) == [0, 1] assert ak.to_list(array.localindex(1)) == [[0, 1, 2], [0, 1, 2]] assert ak.to_list(array.localindex(2)) == [ [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], ]
def to_awkward( self, featurization_key="last", y_dtype="float32", clear_after=False, ): """ Creates an awkward array out of the featurized systems and the associated measurements. Returns ------- awkward array Notes ----- Awkward Array is a library for nested, variable-sized data, including arbitrary-length lists, records, mixed types, and missing data, using NumPy-like idioms. Arrays are dynamically typed, but operations on them are compiled and fast. Their behavior coincides with NumPy when array dimensions are regular and generalizes when they’re not. """ features = self.featurized_systems(key=featurization_key, clear_after=clear_after) # Features is a list of systems (s) and their features (f): [(s0f0, s0f1), (s1f0, s1f1)...] # We are going to iterate over columns: (s0f0, s1f0... snf0), (s0f1, s1f1, ..., snf1) # The result is that X will contain the an array of f0, then an array for f1... etc. X = [ak.from_iter(subX) for subX in zip(*features)] y = ak.from_numpy(self.measurements_as_array(dtype=y_dtype)) return X, y
def test_regular(): np_data = np.array([[1, 3, 5, 4, 2], [3, 7, 8, 2, 4], [2, 3, 1, 7, 7], [5, 1, 9, 10, 2]]) ak_data = ak.from_numpy(np_data) assert ak.ptp(ak_data, axis=1).tolist() == np.ptp(np_data, axis=1).tolist() assert ak.ptp(ak_data, axis=0).tolist() == np.ptp(np_data, axis=0).tolist() assert ak.ptp(ak_data) == np.ptp(np_data)
def test(): first = ak.from_numpy(np.array([1, 2, 3])) deltas = ak.Array([[1, 2], [1, 2], [1, 2, 3]]) assert np.hstack( (first[:, np.newaxis], ak.fill_none(ak.pad_none(deltas, 3, axis=-1), 999))).tolist() == [[1, 1, 2, 999], [2, 1, 2, 999], [3, 1, 2, 3]]
def test(): parameters = ak.from_numpy(np.arange(3 * 3).reshape(-1, 3), regulararray=True) array = ak.zip((parameters, parameters)) assert array.layout.purelist_isregular assert ak.to_list(array) == [ [(0, 0), (1, 1), (2, 2)], [(3, 3), (4, 4), (5, 5)], [(6, 6), (7, 7), (8, 8)], ]
def ak_num(array, **kwargs): """ awkward.num() wrapper also for work in empty array :param array: Data containing nested lists to count. :param kwargs: keywords arguments for awkward.num(). :return: an array of integers specifying the number of elements at a particular level. If array is empty, return empty. """ if len(array) == 0: return ak.from_numpy(np.empty(0, dtype='int64')) return ak.num(array, **kwargs)
def cluster(inter, classify_by_energy=False): """ Function which clusters the found clusters together. To cluster events a weighted mean is computed for time and position. The individual interactions are weighted by their energy. The energy of clustered interaction is given by the total sum. Events can be classified either by the first interaction in time in the cluster or by the highest energy deposition. Args: inter (awkward.Array): Array containing at least the following fields: x,y,z,t,ed,cluster_ids, type, parenttype, creaproc, edproc. Kwargs: classify_by_energy (bool): If true events are classified according to the properties of the highest energy deposit within the cluster. If false cluster is classified according to first interaction. Returns: awkward.Array: Clustered events with nest conform classification. """ if len(inter) == 0: result_cluster_dtype = [('x', 'float64'), ('y', 'float64'), ('z', 'float64'), ('t', 'float64'), ('ed', 'float64'), ('nestid', 'int64'), ('A', 'int64'), ('Z', 'int64')] return ak.from_numpy(np.empty(0, dtype=result_cluster_dtype)) # Sort interactions by cluster_ids to simplify looping inds = ak.argsort(inter['cluster_ids']) inter = inter[inds] # TODO: Better way to do this with awkward? x = inter['x'] y = inter['y'] z = inter['z'] ed = inter['ed'] time = inter['t'] ci = inter['cluster_ids'] types = inter['type'] parenttype = inter['parenttype'] creaproc = inter['creaproc'] edproc = inter['edproc'] # Init result and cluster: res = ak.ArrayBuilder() _cluster(x, y, z, ed, time, ci, types, parenttype, creaproc, edproc, classify_by_energy, res) return res.snapshot()
def test_IndexedArray(): old = ak.layout.IndexedArray64( ak.layout.Index64(np.array([1, 0], np.int64)), ak.layout.RegularArray( ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout, 3), ) new = v1_to_v2(old) assert v1v2_equal(old[1, 1:], new[1, 1:]) assert ak.to_list(old[1, 1:]) == [[5, 6, 7, 8, 9], [10, 11, 12, 13, 14]] assert ak.to_list(new[1, 1:]) == [[5, 6, 7, 8, 9], [10, 11, 12, 13, 14]] with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list(new[0, np.newaxis]) == [[[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29]]] assert ak.to_list(old[0, np.newaxis]) == [[[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29]]] assert old.minmax_depth == (3, 3) assert new.minmax_depth == (3, 3) assert v1v2_equal(old[1, ...], new[1, ...]) assert ak.to_list(old[1, ...]) == [ [0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14], ] assert ak.to_list(new[1, ...]) == [ [0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14], ] expectation = [ [[0, 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]], ] assert (ak.to_list(old[[1, 0], ]) == expectation) assert (ak.to_list(new[[1, 0], ]) == expectation) assert ak.to_list(old[[1, 0]]) == expectation assert ak.to_list(new[[1, 0]]) == expectation assert ak.to_list(old[1, [1, 0]]) == [[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]] assert ak.to_list(new[1, [1, 0]]) == [[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]] assert v1v2_equal(old[1, [1, 0]], new[1, [1, 0]])
def in_sensitive_volume(events, sensitive_volumes): """ Function which identifies which events are inside sensitive volumes. Further, tests if sensitive volumes overlap. Assigns volume id, and xenon density to interactions. Args: events (ak.records): Awkward record of the interactions. sensitive_volumes (dict): Dictionary of the different volumes defined via the SensitiveVolume class. Returns: ak.array: Awkward array containing the event ids. """ if len(events) == 0: res_det_dtype = [('xe_density', 'float64'), ('vol_id', 'int64'), ('create_S2', 'bool'), ] return ak.from_numpy(np.empty(0, dtype=res_det_dtype)) for ind, vol in enumerate(sensitive_volumes): res = ak.ArrayBuilder() res = _inside_sens_vol(events['x'], events['y'], events['z'], vol.roi, vol.volume_id, vol.xe_density, vol.create_S2, res) if ind: # Now we add the other results, but first test if # volumes overlap. Only possible if we explicitly loop # over everything. This reduces performance but adds layer of # safety. m = (result['vol_id'] > 0) & (res['vol_id'] == vol.volume_id) if ak.any(m): overlapping_id = result[m][0] # Get volume name: name = [vol.name for vol in sensitive_volumes if vol.volume_id == overlapping_id][0] raise ValueError(f'The volume {vol.name} is overlapping with' f' volume {name}!') new_results = res.snapshot() for field in result.fields: # Workaround since we cannot sum up records-arrays anymore result[field] = result[field] + new_results[field] else: # First result initiates the array result = res.snapshot() return result
def test_lists(): one = ak.Array([[1, 2, 3], [], [4, 5]]).layout two = ak.Array([[1.1, 2.2], [3.3, 4.4]]).layout three = ak.layout.EmptyArray() four = ak.from_numpy(np.array([[10], [20]]), regulararray=True, highlevel=False) assert ak.to_list(one.mergemany([two, three, four])) == [ [1.0, 2.0, 3.0], [], [4.0, 5.0], [1.1, 2.2], [3.3, 4.4], [10.0], [20.0], ] assert ak.to_list(four.mergemany([three, two, one])) == [ [10.0], [20.0], [1.1, 2.2], [3.3, 4.4], [1.0, 2.0, 3.0], [], [4.0, 5.0], ] one = ak.layout.ListArray64(one.starts, one.stops, one.content) two = ak.layout.ListArray64(two.starts, two.stops, two.content) assert ak.to_list(one.mergemany([two, three, four])) == [ [1.0, 2.0, 3.0], [], [4.0, 5.0], [1.1, 2.2], [3.3, 4.4], [10.0], [20.0], ] assert ak.to_list(four.mergemany([three, two, one])) == [ [10.0], [20.0], [1.1, 2.2], [3.3, 4.4], [1.0, 2.0, 3.0], [], [4.0, 5.0], ]
def test_fromnumpy(): a = np.arange(2 * 3 * 5).reshape((2, 3, 5)) b = ak.from_numpy(a) assert ak.to_list(a) == ak.to_list(b)
def test(): np_data = np.random.random(size=(4, 100 * 1024 * 1024 // 8 // 4)) array = ak.from_numpy(np_data, regulararray=False) assert np_data.nbytes == array.nbytes
def test_toawkward0(): array = ak.from_iter([1.1, 2.2, 3.3, 4.4], highlevel=False) assert isinstance(ak.to_awkward0(array), np.ndarray) assert ak.to_awkward0(array).tolist() == [1.1, 2.2, 3.3, 4.4] array = ak.from_numpy(np.arange(2 * 3 * 5).reshape(2, 3, 5), highlevel=False).toRegularArray() assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray) assert ak.to_awkward0(array).tolist() == [ [[0, 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]], ] array = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray) assert ak.to_awkward0(array).tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]] array = ak.layout.ListArray64( ak.layout.Index64(np.array([4, 999, 1], dtype=np.int64)), ak.layout.Index64(np.array([7, 999, 3], dtype=np.int64)), ak.layout.NumpyArray( np.array([3.14, 4.4, 5.5, 123, 1.1, 2.2, 3.3, 321])), ) assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray) assert ak.to_awkward0(array).tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]] array = ak.from_iter( [ { "x": 0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }, ], highlevel=False, ) assert isinstance(ak.to_awkward0(array[2]), dict) assert ak.to_awkward0(array[2])["x"] == 2.2 assert isinstance(ak.to_awkward0(array[2])["y"], np.ndarray) assert ak.to_awkward0(array[2])["y"].tolist() == [2, 2] assert isinstance(ak.to_awkward0(array), awkward0.Table) assert ak.to_awkward0(array).tolist() == [ { "x": 0, "y": [] }, { "x": 1.1, "y": [1] }, { "x": 2.2, "y": [2, 2] }, { "x": 3.3, "y": [3, 3, 3] }, ] array = ak.from_iter([(0, []), (1.1, [1]), (2.2, [2, 2]), (3.3, [3, 3, 3])], highlevel=False) assert isinstance(ak.to_awkward0(array), awkward0.Table) assert ak.to_awkward0(array).tolist() == [ (0, []), (1.1, [1]), (2.2, [2, 2]), (3.3, [3, 3, 3]), ] assert isinstance(ak.to_awkward0(array[2]), tuple) assert ak.to_awkward0(array[2])[0] == 2.2 assert ak.to_awkward0(array[2])[1].tolist() == [2, 2] array = ak.from_iter([0.0, [], 1.1, [1], 2.2, [2, 2], 3.3, [3, 3, 3]], highlevel=False) assert isinstance(ak.to_awkward0(array), awkward0.UnionArray) assert ak.to_awkward0(array).tolist() == [ 0.0, [], 1.1, [1], 2.2, [2, 2], 3.3, [3, 3, 3], ] array = ak.from_iter([1.1, 2.2, None, None, 3.3, None, 4.4], highlevel=False) assert isinstance(ak.to_awkward0(array), awkward0.IndexedMaskedArray) assert ak.to_awkward0(array).tolist() == [ 1.1, 2.2, None, None, 3.3, None, 4.4 ] content = ak.layout.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) index = ak.layout.Index64(np.array([3, 2, 2, 5, 0], dtype=np.int64)) array = ak.layout.IndexedArray64(index, content) assert isinstance(ak.to_awkward0(array), awkward0.IndexedArray) assert ak.to_awkward0(array).tolist() == [3.3, 2.2, 2.2, 5.5, 0.0]
def test_regulararray(): layout = ak.from_numpy(np.array([[1, 2, 3], [4, 5, 6]]), regulararray=True, highlevel=False) numbatype = ak._connect._numba.arrayview.tonumbatype(layout.form) assert ak_connect_numba_layout.typeof(layout).name == numbatype.name lookup1 = ak_connect_numba_arrayview.Lookup(layout) lookup2 = ak_connect_numba_arrayview.Lookup(layout.form) numbatype.form_fill(0, layout, lookup2) assert np.array_equal(lookup1.arrayptrs, lookup2.arrayptrs) assert np.array_equal(lookup1.sharedptrs == -1, lookup2.sharedptrs == -1) counter = [0] def materialize(): counter[0] += 1 return layout generator = ak.layout.ArrayGenerator(materialize, form=layout.form, length=len(layout)) virtualarray = ak.layout.VirtualArray(generator) lookup3 = ak_connect_numba_arrayview.Lookup(virtualarray) assert len(lookup1.arrayptrs) + 3 == len(lookup3.arrayptrs) array = ak.Array(virtualarray) array.numba_type assert counter[0] == 0 @numba.njit def f3(x): return x assert isinstance(f3(array).layout, ak.layout.VirtualArray) assert counter[0] == 0 @numba.njit def f1(x): return x[1][1] assert f1(array) == 5 assert counter[0] == 1 assert f1(array) == 5 assert counter[0] == 1 @numba.njit def f2(x): return x[1] assert ak.to_list(f2(array)) == [4, 5, 6] assert counter[0] == 1 assert ak.to_list(f2(array)) == [4, 5, 6] assert counter[0] == 1 assert ak.to_list(f3(array)) == [[1, 2, 3], [4, 5, 6]]
def test(): array = ak.from_numpy(np.zeros((3, 0), dtype=np.int32)) buffs = ak.to_buffers(array) new_array = ak.from_buffers(*buffs) assert ak.to_list(new_array) == [[], [], []]
def test_awkward_from_numpy_array(): x = np.arange(12) y = ak.from_numpy(x) z = ak.is_none(y) assert z.tolist() == [False] * 12
def test(): array = ak.from_numpy(np.zeros((3, 3, 5))) flattened = ak.flatten(array, axis=None) assert flattened.ndim == 1
def ggTauTau_inclusive_preselection(events, photons, electrons, muons, taus, jets, dR, genPart, Category_pairsLoose, options, debug): """ Performs inclusive ggTauTau preselection, requiring >=1 (leptons + tau_h). Assumes diphoton preselection has already been applied. Also calculates relevant event-level variables. """ cut_diagnostics = utils.CutDiagnostics( events=events, debug=debug, cut_set="[analysis_selections.py : ggTauTau_inclusive_preselection]") # Get number of electrons, muons, taus selected_electrons = electrons[lepton_selections.select_electrons( events, photons, electrons, options, debug)] selected_muons = muons[lepton_selections.select_muons( events, photons, muons, options, debug)] selected_taus = taus[tau_selections.select_taus(events, photons, selected_muons, selected_electrons, taus, options, debug)] n_electrons = awkward.num(selected_electrons) n_muons = awkward.num(selected_muons) n_taus = awkward.num(selected_taus) # Require >= 1 lep/tau n_leptons_and_taus = n_electrons + n_muons + n_taus # only events with hadronic taus (no leptonic taus!!!!!!!!!!) atleast_one_had_tau_cut = (n_taus >= 1) # Require OS leptons/taus for events with 2 leptons/taus sum_charge = awkward.sum(selected_electrons.charge, axis=1) + awkward.sum( selected_muons.charge, axis=1) + awkward.sum(selected_taus.charge, axis=1) charge_cut = sum_charge == 0 two_leptons = n_leptons_and_taus == 2 not_two_leptons = n_leptons_and_taus != 2 os_cut = ( two_leptons & charge_cut ) | not_two_leptons # only require 2 OS leptons if there are ==2 leptons in the event # Select jets (don't cut on jet quantities for selection, but they will be useful for BDT training) selected_jets = jets[jet_selections.select_jets(events, photons, selected_electrons, selected_muons, selected_taus, jets, options, debug)] all_cuts = os_cut & atleast_one_had_tau_cut cut_diagnostics.add_cuts([atleast_one_had_tau_cut, os_cut, all_cuts], ["N_taus >= 1", "OS dileptons", "all"]) # Keep only selected events selected_events = events[all_cuts] selected_photons = photons[all_cuts] selected_electrons = selected_electrons[all_cuts] selected_muons = selected_muons[all_cuts] selected_taus = selected_taus[all_cuts] selected_jets = selected_jets[all_cuts] dR = dR[all_cuts] # Calculate event-level variables selected_events = lepton_selections.set_electrons(selected_events, selected_electrons, debug) selected_events = lepton_selections.set_muons(selected_events, selected_muons, debug) selected_events = tau_selections.set_taus(selected_events, selected_taus, debug) selected_events = jet_selections.set_jets(selected_events, selected_jets, options, debug) if genPart is not None: genPart = genPart[all_cuts] selected_events = gen_selections.set_genZ(selected_events, genPart, options, debug) else: selected_events["genZ_decayMode"] = awkward.from_numpy( numpy.ones(len(selected_events)) * -1) selected_events["tau_motherID"] = awkward.from_numpy( numpy.ones(len(selected_events)) * -1) selected_events = compound_selections.compound_selections( selected_events, options, debug) return selected_events
def test(): array = ak.from_numpy(np.zeros((1, 0), dtype=np.int32), regulararray=True) packed = ak.packed(array) assert ak.to_list(packed) == [[]]
def test_toRegularArray(): assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 5, 7, 11), np.float64)).layout.toRegularArray()))) == "2 * 3 * 5 * 7 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 5, 7, 11, 0), np.float64)).layout.toRegularArray()))) == "2 * 3 * 5 * 7 * 11 * 0 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 5, 7, 0, 11), np.float64)).layout.toRegularArray()))) == "2 * 3 * 5 * 7 * 0 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 5, 0, 7, 11), np.float64)).layout.toRegularArray()))) == "2 * 3 * 5 * 0 * 7 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 0, 5, 7, 11), np.float64)).layout.toRegularArray()))) == "2 * 3 * 0 * 5 * 7 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 0, 3, 5, 7, 11), np.float64)).layout.toRegularArray()))) == "2 * 0 * 3 * 5 * 7 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((0, 2, 3, 5, 7, 11), np.float64)).layout.toRegularArray()))) == "0 * 2 * 3 * 5 * 7 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 5, 7, 0, 11, 0), np.float64)).layout.toRegularArray()))) == "2 * 3 * 5 * 7 * 0 * 11 * 0 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 5, 0, 7, 11, 0), np.float64)).layout.toRegularArray()))) == "2 * 3 * 5 * 0 * 7 * 11 * 0 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 3, 0, 5, 7, 0, 11), np.float64)).layout.toRegularArray()))) == "2 * 3 * 0 * 5 * 7 * 0 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((2, 0, 3, 5, 7, 0, 11), np.float64)).layout.toRegularArray()))) == "2 * 0 * 3 * 5 * 7 * 0 * 11 * float64") assert (str( ak.type( ak.Array( ak.from_numpy(np.empty((0, 2, 3, 5, 7, 0, 11), np.float64)).layout.toRegularArray()))) == "0 * 2 * 3 * 5 * 7 * 0 * 11 * float64")
def test_actual(): x = ak.from_numpy(np.arange(2 * 3 * 4, dtype=np.int64).reshape(2, 3, 4)) s = x[..., :0] result = ak.zip({"q": s, "t": s}) assert str(ak.type(result)) == '2 * 3 * 0 * {"q": int64, "t": int64}'
def test(): array = np.array([1, 2, 9, 0]) nplike = ak.nplike.of(array) ak_array = ak.from_numpy(array) assert ak.to_list(nplike.to_rectilinear(array)) == ak.to_list(ak_array)