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