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 match_with_pt(first, second, deltaRCut=0.4, ptCut=0.5): ''' match based on deltaR between first and second, and impose that second.pt > first.pt*ptCut ''' drCut2 = deltaRCut**2 combs = ak.cartesian([first, second], nested=True) return ak.any((delta_r2(combs['0'], combs['1']) < drCut2) & (combs['1'].pt > ptCut * combs['0'].pt), axis=2)
def nearest(self, other, metric=lambda a, b: a.delta_r(b), return_metric=False): """Return nearest object to this one Only works for first axis (i.e. top-level ListArrays) """ a, b = awkward1.unzip(awkward1.cartesian([self, other], nested=True)) mval = metric(a, b) mmin = awkward1.argmin(mval, axis=-1) if return_metric: return b[mmin], mval[mmin] return b[mmin]
def test_axis2(): one = awkward1.Array([[[0, 1, 2], [], [3, 4]], [[0, 1, 2], [], [3, 4]]]) two = awkward1.Array([[[100, 200], [300], [400, 500]], [[100, 200], [300], [400, 500]]]) assert awkward1.to_list(awkward1.cartesian( [one, two], axis=2)) == [[[(0, 100), (0, 200), (1, 100), (1, 200), (2, 100), (2, 200)], [], [(3, 400), (3, 500), (4, 400), (4, 500)]], [[(0, 100), (0, 200), (1, 100), (1, 200), (2, 100), (2, 200)], [], [(3, 400), (3, 500), (4, 400), (4, 500)]]]
def nearest(self, other, metric=lambda a, b: a.delta_r(b), return_metric=False): """Return nearest object to this one The default metric is `delta_r`. """ a, b = awkward1.unzip(awkward1.cartesian([self, other], nested=True)) mval = metric(a, b) mmin = awkward1.argmin(mval, axis=-1) if return_metric: return b[mmin], mval[mmin] return b[mmin]
def matchJets(self, obj, jet, deltaRCut=0.4): combs = ak.cartesian([obj, jet], nested=True) jet_index = ak.local_index(delta_r( combs['0'], combs['1']))[delta_r(combs['0'], combs['1']) < 0.4] jet_index_pad = ak.flatten(ak.fill_none( ak.pad_none(jet_index, target=1, clip=True, axis=2), 0), axis=2) mask = ak.num(jet_index, axis=2) > 0 # a mask for obj with a matched jet mask_match = mask * 1 + ~mask * 0 mask_nomatch = mask * 0 + ~mask * 1 return jet_index_pad, mask_match, mask_nomatch
def nearest( self, other, axis=1, metric=lambda a, b: a.delta_r(b), return_metric=False, threshold=None, ): """Return nearest object to this one Finds item in ``other`` satisfying ``min(metric(self, other))``. The two arrays should be broadcast-compatible on all axes other than the specified axis, which will be used to form a cartesian product. If axis=None, broadcast arrays directly. The return shape will be that of ``self``. Parameters ---------- other : awkward1.Array Another array with same shape in all but ``axis`` axis : int, optional The axis to form the cartesian product (default 1). If None, the metric is directly evaluated on the input arrays (i.e. they should broadcast) metric : callable A function of two arguments, returning a scalar. The default metric is `delta_r`. return_metric : bool, optional If true, return both the closest object and its metric (default false) threshold : Number, optional If set, any objects with ``metric > threshold`` will be masked from the result """ if axis is None: a, b = self, other # NotImplementedError: ak.firsts with axis=-1 axis = other.layout.purelist_depth - 2 else: a, b = awkward1.unzip( awkward1.cartesian([self, other], axis=axis, nested=True) ) mval = metric(a, b) # prefer keepdims=True: awkward-1.0 #434 mmin = awkward1.singletons(awkward1.argmin(mval, axis=axis + 1)) out = awkward1.firsts(b[mmin], axis=axis + 1) metric = awkward1.firsts(mval[mmin], axis=axis + 1) if threshold is not None: out = out.mask[metric <= threshold] if return_metric: return out, metric return out
def cross(first, second): tmp = ak.cartesian([first, second]) combs = (tmp['0'] + tmp['1']) combs['0'] = tmp['0'] combs['1'] = tmp['1'] return combs
def match2(first, second, deltaRCut=0.4): drCut2 = deltaRCut**2 combs = ak.cartesian([first, second], nested=True) return ak.any((combs['0'].delta_r2(combs['1']) < drCut2), axis=2)
def delta_r_v2(first, second): combs = ak.cartesian([first, second], nested=True) return np.sqrt(delta_r2(combs['0'], combs['1']))
def test_axis1(): one = awkward1.Array([[0, 1, 2], [], [3, 4]]) two = awkward1.Array([[100, 200], [300], [400, 500]]) three = awkward1.Array([["a", "b"], ["c", "d"], ["e"]]) assert awkward1.to_list(awkward1.cartesian([one])) == [[(0,), (1,), (2,)], [], [(3,), (4,)]] assert awkward1.to_list(awkward1.cartesian({"x": one})) == [[{"x": 0}, {"x": 1}, {"x": 2}], [], [{"x": 3}, {"x": 4}]] assert awkward1.to_list(awkward1.cartesian([one, two])) == [[(0, 100), (0, 200), (1, 100), (1, 200), (2, 100), (2, 200)], [], [(3, 400), (3, 500), (4, 400), (4, 500)]] if not py27 and not py35: assert awkward1.to_list(awkward1.cartesian({"x": one, "y": two})) == [[{"x": 0, "y": 100}, {"x": 0, "y": 200}, {"x": 1, "y": 100}, {"x": 1, "y": 200}, {"x": 2, "y": 100}, {"x": 2, "y": 200}], [], [{"x": 3, "y": 400}, {"x": 3, "y": 500}, {"x": 4, "y": 400}, {"x": 4, "y": 500}]] assert awkward1.to_list(awkward1.cartesian([one, two, three])) == [[(0, 100, "a"), (0, 100, "b"), (0, 200, "a"), (0, 200, "b"), (1, 100, "a"), (1, 100, "b"), (1, 200, "a"), (1, 200, "b"), (2, 100, "a"), (2, 100, "b"), (2, 200, "a"), (2, 200, "b")], [], [(3, 400, "e"), (3, 500, "e"), (4, 400, "e"), (4, 500, "e")]] assert awkward1.to_list(awkward1.cartesian([one, two, three], nested=[0])) == [[[(0, 100, "a"), (0, 100, "b"), (0, 200, "a"), (0, 200, "b")], [(1, 100, "a"), (1, 100, "b"), (1, 200, "a"), (1, 200, "b")], [(2, 100, "a"), (2, 100, "b"), (2, 200, "a"), (2, 200, "b")]], [], [[(3, 400, "e"), (3, 500, "e")], [(4, 400, "e"), (4, 500, "e")]]] assert awkward1.to_list(awkward1.cartesian([one, two, three], nested=[1])) == [[[(0, 100, "a"), (0, 100, "b")], [(0, 200, "a"), (0, 200, "b")], [(1, 100, "a"), (1, 100, "b")], [(1, 200, "a"), (1, 200, "b")], [(2, 100, "a"), (2, 100, "b")], [(2, 200, "a"), (2, 200, "b")]], [], [[(3, 400, "e")], [(3, 500, "e")], [(4, 400, "e")], [(4, 500, "e")]]] assert awkward1.to_list(awkward1.cartesian([one, two, three], nested=[0, 1])) == [[[[(0, 100, "a"), (0, 100, "b")], [(0, 200, "a"), (0, 200, "b")]], [[(1, 100, "a"), (1, 100, "b")], [(1, 200, "a"), (1, 200, "b")]], [[(2, 100, "a"), (2, 100, "b")], [(2, 200, "a"), (2, 200, "b")]]], [], [[[(3, 400, "e")], [(3, 500, "e")]], [[(4, 400, "e")], [(4, 500, "e")]]]] assert awkward1.to_list(awkward1.cartesian([one, two, three], nested=[])) == awkward1.to_list(awkward1.cartesian([one, two, three], nested=False)) assert awkward1.to_list(awkward1.cartesian([one, two, three], nested=[])) == awkward1.to_list(awkward1.cartesian([one, two, three], nested=None)) assert awkward1.to_list(awkward1.cartesian([one, two, three], nested=[0, 1])) == awkward1.to_list(awkward1.cartesian([one, two, three], nested=True))
def test_axis0(): one = awkward1.Array([1.1, 2.2, 3.3]) two = awkward1.Array([100, 200, 300, 400, 500]) three = awkward1.Array(["a", "b"]) assert awkward1.to_list(awkward1.cartesian([one], axis=0)) == [(1.1,), (2.2,), (3.3,)] assert awkward1.to_list(awkward1.cartesian({"x": one}, axis=0)) == [{"x": 1.1}, {"x": 2.2}, {"x": 3.3}] assert awkward1.to_list(awkward1.cartesian([one, two], axis=0)) == [(1.1, 100), (1.1, 200), (1.1, 300), (1.1, 400), (1.1, 500), (2.2, 100), (2.2, 200), (2.2, 300), (2.2, 400), (2.2, 500), (3.3, 100), (3.3, 200), (3.3, 300), (3.3, 400), (3.3, 500)] if not py27 and not py35: assert awkward1.to_list(awkward1.cartesian({"x": one, "y": two}, axis=0)) == [{"x": 1.1, "y": 100}, {"x": 1.1, "y": 200}, {"x": 1.1, "y": 300}, {"x": 1.1, "y": 400}, {"x": 1.1, "y": 500}, {"x": 2.2, "y": 100}, {"x": 2.2, "y": 200}, {"x": 2.2, "y": 300}, {"x": 2.2, "y": 400}, {"x": 2.2, "y": 500}, {"x": 3.3, "y": 100}, {"x": 3.3, "y": 200}, {"x": 3.3, "y": 300}, {"x": 3.3, "y": 400}, {"x": 3.3, "y": 500}] assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0)) == [(1.1, 100, "a"), (1.1, 100, "b"), (1.1, 200, "a"), (1.1, 200, "b"), (1.1, 300, "a"), (1.1, 300, "b"), (1.1, 400, "a"), (1.1, 400, "b"), (1.1, 500, "a"), (1.1, 500, "b"), (2.2, 100, "a"), (2.2, 100, "b"), (2.2, 200, "a"), (2.2, 200, "b"), (2.2, 300, "a"), (2.2, 300, "b"), (2.2, 400, "a"), (2.2, 400, "b"), (2.2, 500, "a"), (2.2, 500, "b"), (3.3, 100, "a"), (3.3, 100, "b"), (3.3, 200, "a"), (3.3, 200, "b"), (3.3, 300, "a"), (3.3, 300, "b"), (3.3, 400, "a"), (3.3, 400, "b"), (3.3, 500, "a"), (3.3, 500, "b")] assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=[0])) == [[(1.1, 100, "a"), (1.1, 100, "b"), (1.1, 200, "a"), (1.1, 200, "b"), (1.1, 300, "a")], [(1.1, 300, "b"), (1.1, 400, "a"), (1.1, 400, "b"), (1.1, 500, "a"), (1.1, 500, "b")], [(2.2, 100, "a"), (2.2, 100, "b"), (2.2, 200, "a"), (2.2, 200, "b"), (2.2, 300, "a")], [(2.2, 300, "b"), (2.2, 400, "a"), (2.2, 400, "b"), (2.2, 500, "a"), (2.2, 500, "b")], [(3.3, 100, "a"), (3.3, 100, "b"), (3.3, 200, "a"), (3.3, 200, "b"), (3.3, 300, "a")], [(3.3, 300, "b"), (3.3, 400, "a"), (3.3, 400, "b"), (3.3, 500, "a"), (3.3, 500, "b")]] assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=[1])) == [[(1.1, 100, "a"), (1.1, 100, "b")], [(1.1, 200, "a"), (1.1, 200, "b")], [(1.1, 300, "a"), (1.1, 300, "b")], [(1.1, 400, "a"), (1.1, 400, "b")], [(1.1, 500, "a"), (1.1, 500, "b")], [(2.2, 100, "a"), (2.2, 100, "b")], [(2.2, 200, "a"), (2.2, 200, "b")], [(2.2, 300, "a"), (2.2, 300, "b")], [(2.2, 400, "a"), (2.2, 400, "b")], [(2.2, 500, "a"), (2.2, 500, "b")], [(3.3, 100, "a"), (3.3, 100, "b")], [(3.3, 200, "a"), (3.3, 200, "b")], [(3.3, 300, "a"), (3.3, 300, "b")], [(3.3, 400, "a"), (3.3, 400, "b")], [(3.3, 500, "a"), (3.3, 500, "b")]] assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=[0, 1])) == [[[(1.1, 100, "a"), (1.1, 100, "b")], [(1.1, 200, "a"), (1.1, 200, "b")], [(1.1, 300, "a"), (1.1, 300, "b")], [(1.1, 400, "a"), (1.1, 400, "b")], [(1.1, 500, "a"), (1.1, 500, "b")]], [[(2.2, 100, "a"), (2.2, 100, "b")], [(2.2, 200, "a"), (2.2, 200, "b")], [(2.2, 300, "a"), (2.2, 300, "b")], [(2.2, 400, "a"), (2.2, 400, "b")], [(2.2, 500, "a"), (2.2, 500, "b")]], [[(3.3, 100, "a"), (3.3, 100, "b")], [(3.3, 200, "a"), (3.3, 200, "b")], [(3.3, 300, "a"), (3.3, 300, "b")], [(3.3, 400, "a"), (3.3, 400, "b")], [(3.3, 500, "a"), (3.3, 500, "b")]]] assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=[])) == awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=False)) assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=[])) == awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=None)) assert awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=[0, 1])) == awkward1.to_list(awkward1.cartesian([one, two, three], axis=0, nested=True))
def run(self, decay, ncpu=5): nevents_real = 0 import time start_time = time.time() processEvents = {} for pr in self.processes: fin = self.baseDir + pr + '.root' #input file if not os.path.isfile(fin): print('file ', fin, ' does not exist. exit') exit(3) tfin = ROOT.TFile.Open(fin) tfin.cd() found = False for key in tfin.GetListOfKeys(): if 'eventsProcessed' == key.GetName(): events = tfin.eventsProcessed.GetVal() processEvents[pr] = events found = True if not found: processEvents[pr] = 1 tfin.Close() for pr in self.processes: print(' running over process : ', pr) fin = self.baseDir + pr + '.root' #input file fout = self.baseDir + pr + '.root' #output file for tree fhisto = self.baseDir + pr + '_histo.root' #output file for histograms file = uproot.open(fin) tree = file['events'] #Number of events to keep and analyse n_events = 1000 #Container for the reco particles p_c = 'RP' events = tree.arrays(library="ak", how="zip", filter_name=f"{p_c}*")[:n_events] print('events loaded ', events) p = events[p_c] print('fin ', fin) p_c = 'RP' #Number of events to keep and analyse n_events = 10000 file = uproot.open(fin) tree = file[self.treename] print('here-1.2') events = tree.arrays(library="ak", how="zip", filter_name=f"{p_c}*")[:n_events] print('here0') #Container for the reco particles p = events[p_c][:n_events] print('here1') p["p"] = kinematics_flat.calc_p(p) p_cut = p["p"] > 1. p = p[p_cut] print('here2') pi_cut = abs(p["mass"] - lp.pi_plus.mass / 1000.) < 1e-4 pi = p[pi_cut] print('here3') k_cut = abs(p["mass"] - lp.K_plus.mass / 1000.) < 1e-4 k = p[k_cut] print('here4') D = ak.cartesian({"k": k, "pi": pi}) D_cut = np.sign(D["k", "charge"]) != np.sign(D["pi", "charge"]) D = D[D_cut] print('here5') PDG_K_m = lp.K_plus.mass / 1000. PDG_pi_m = lp.pi_plus.mass / 1000. D["mass"] = kinematics_flat.mass([D["k"], D["pi"]], [PDG_K_m, PDG_pi_m]) print('here6') PDG_D_m = lp.D_0.mass / 1000. D_window = 0.05 D_cut = abs(D["mass"] - PDG_D_m) < D_window D = D[D_cut] print('here7') B = ak.cartesian({"D_k": D["k"], "D_pi": D["pi"], "pi": pi}) B_cut = np.sign(B["D_k", "charge"]) == np.sign(B["pi", "charge"]) B = B[B_cut] B["mass"] = kinematics_flat.mass([B["D_k"], B["D_pi"], B["pi"]], [PDG_K_m, PDG_pi_m, PDG_pi_m]) print('here8') data_np = ak.to_numpy(ak.flatten(B["mass"])) print(data_np) validfile = self.testfile(fout) if not validfile: continue nevents_real += df_cut.Count().GetValue() tf = ROOT.TFile.Open(fhisto, 'RECREATE') for v in self.variables: model = ROOT.RDF.TH1DModel( v, ";{};".format(self.variables[v]["title"]), self.variables[v]["bin"], self.variables[v]["xmin"], self.variables[v]["xmax"]) h = snapshot_tdf.Histo1D(model, self.variables[v]["name"]) try: h.Scale(1. * self.procDict[pr]["crossSection"] * self.procDict[pr]["kfactor"] * self.procDict[pr]["matchingEfficiency"] / processEvents[pr]) except KeyError: h.Scale(1. / h.Integral(0, -1)) h.Write() tf.Close() elapsed_time = time.time() - start_time print( '==============================SUMMARY==============================' ) print('Elapsed time (H:M:S) : ', time.strftime("%H:%M:%S", time.gmtime(elapsed_time))) print('Events Processed/Second : ', int(nevents_real / elapsed_time)) print('Total Events Processed : ', nevents_real) print( '===================================================================' )
def cartesian(self, other, **kwargs): if isinstance(other, AwkwardSeries): other = other.data return AwkwardSeries(ak.cartesian([self.arr.data, other], **kwargs))