def test_jagged_pad(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
assert a.pad(4, clip=True).tolist() == [[1.1, 2.2, 3.3, None], [None, None, None, None], [4.4, 5.5, None, None]]
assert a.pad(4, numpy.ma.masked, clip=True).regular().tolist() == [[1.1, 2.2, 3.3, None], [None, None, None, None], [4.4, 5.5, None, None]]
assert a.pad(4).tolist() == [[1.1, 2.2, 3.3, None], [None, None, None, None], [4.4, 5.5, None, None]]
assert a.pad(4, numpy.ma.masked).regular().tolist() == [[1.1, 2.2, 3.3, None], [None, None, None, None], [4.4, 5.5, None, None]]
a = awkward0.fromiter([[1.1, 2.2, 3.3, 4.4, 5.5], [], [6.6, 7.7, 8.8], [9.9]])
assert a.pad(3).tolist() == [[1.1, 2.2, 3.3, 4.4, 5.5], [None, None, None], [6.6, 7.7, 8.8], [9.9, None, None]]
assert a.pad(3, clip=True).tolist() == [[1.1, 2.2, 3.3], [None, None, None], [6.6, 7.7, 8.8], [9.9, None, None]]
def test_physics_jetcleaning(self):
jet_m = awkward0.fromiter([[60.0, 70.0, 80.0], [], [90.0, 100.0]])
jet_pt = awkward0.fromiter([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
e_pt = awkward0.fromiter([[20.2, 50.5], [50.5], [50.5]])
jet_eta = awkward0.fromiter([[-3.0, -2.0, 2.0], [], [-1.0, 1.0]])
e_eta = awkward0.fromiter([[-2.2, 0.0], [0.0], [1.1]])
jet_phi = awkward0.fromiter([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
e_phi = awkward0.fromiter([[0.1, 0.78], [0.78], [-0.77]])
jets = uproot_methods.TLorentzVectorArray.from_ptetaphim(
jet_pt, jet_eta, jet_phi, jet_m)
electrons = uproot_methods.TLorentzVectorArray.from_ptetaphim(
e_pt, e_eta, e_phi, 0.000511)
combinations = jets.cross(electrons, nested=True)
def delta_r(one, two):
return one.delta_r(two)
assert (~(delta_r(combinations.i0, combinations.i1) < 0.5).any()
).tolist() == [[True, False, True], [], [True, False]]
(jets[~(delta_r(combinations.i0, combinations.i1) < 0.5).any()])
def test_issue367(self):
t = uproot3.open("tests/samples/issue367.root")["tree"]
assert awkward0.fromiter(
t.array("weights.second"))[0].counts.tolist() == [
1000, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 1000, 1000,
1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000,
100, 100, 100, 1
]
def test_jagged_cross_argnested(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
b = awkward0.fromiter([[100, 200], [300], [400]])
c = awkward0.fromiter([[999], [999], [999, 888]])
assert a.cross(b).tolist() == [[(1.1, 100), (1.1, 200), (2.2, 100), (2.2, 200), (3.3, 100), (3.3, 200)], [], [(4.4, 400), (5.5, 400)]]
assert a.argcross(b).tolist() == [[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)], [], [(0, 0), (1, 0)]]
assert a.cross(b, nested=True).tolist() == [[[(1.1, 100), (1.1, 200)], [(2.2, 100), (2.2, 200)], [(3.3, 100), (3.3, 200)]], [], [[(4.4, 400)], [(5.5, 400)]]]
assert a.argcross(b, nested=True).tolist() == [[[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)]], [], [[(0, 0)], [(1, 0)]]]
assert a.cross(b, nested=True).cross(c, nested=True).tolist()[0] == [[[(ai, bi, ci) for ci in c[0]] for bi in b[0]] for ai in a[0]]
assert a.cross(b, nested=True).cross(c, nested=True).tolist()[1] == [[[(ai, bi, ci) for ci in c[1]] for bi in b[1]] for ai in a[1]]
assert a.cross(b, nested=True).cross(c, nested=True).tolist()[2] == [[[(ai, bi, ci) for ci in c[2]] for bi in b[2]] for ai in a[2]]
assert a.cross(b).cross(c).tolist() == [[(1.1, 100, 999), (1.1, 200, 999), (2.2, 100, 999), (2.2, 200, 999), (3.3, 100, 999), (3.3, 200, 999)], [], [(4.4, 400, 999), (4.4, 400, 888), (5.5, 400, 999), (5.5, 400, 888)]]
assert a.cross(b, nested=True).cross(c).tolist() == [[[(1.1, 100, 999), (1.1, 200, 999)], [(2.2, 100, 999), (2.2, 200, 999)], [(3.3, 100, 999), (3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888)], [(5.5, 400, 999), (5.5, 400, 888)]]]
assert a.cross(b).cross(c, nested=True).tolist() == [[[(1.1, 100, 999)], [(1.1, 200, 999)], [(2.2, 100, 999)], [(2.2, 200, 999)], [(3.3, 100, 999)], [(3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888)], [(5.5, 400, 999), (5.5, 400, 888)]]]
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
b = awkward0.fromiter([[100, 200], [300], [400]])
c = awkward0.fromiter([[999], [999], [999, 888, 777]])
assert a.cross(b, nested=True).cross(c, nested=True).tolist()[0] == [[[(ai, bi, ci) for ci in c[0]] for bi in b[0]] for ai in a[0]]
assert a.cross(b, nested=True).cross(c, nested=True).tolist()[1] == [[[(ai, bi, ci) for ci in c[1]] for bi in b[1]] for ai in a[1]]
assert a.cross(b, nested=True).cross(c, nested=True).tolist()[2] == [[[(ai, bi, ci) for ci in c[2]] for bi in b[2]] for ai in a[2]]
assert a.cross(b).cross(c).tolist() == [[(1.1, 100, 999), (1.1, 200, 999), (2.2, 100, 999), (2.2, 200, 999), (3.3, 100, 999), (3.3, 200, 999)], [], [(4.4, 400, 999), (4.4, 400, 888), (4.4, 400, 777), (5.5, 400, 999), (5.5, 400, 888), (5.5, 400, 777)]]
assert a.cross(b, nested=True).cross(c).tolist() == [[[(1.1, 100, 999), (1.1, 200, 999)], [(2.2, 100, 999), (2.2, 200, 999)], [(3.3, 100, 999), (3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888), (4.4, 400, 777)], [(5.5, 400, 999), (5.5, 400, 888), (5.5, 400, 777)]]]
assert a.cross(b).cross(c, nested=True).tolist() == [[[(1.1, 100, 999)], [(1.1, 200, 999)], [(2.2, 100, 999)], [(2.2, 200, 999)], [(3.3, 100, 999)], [(3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888), (4.4, 400, 777)], [(5.5, 400, 999), (5.5, 400, 888), (5.5, 400, 777)]]]
def test_jagged_zip(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
b = awkward0.JaggedArray([1, 5, 5], [4, 5, 7], [999, 10, 20, 30, 999, 40, 50, 999])
c = numpy.array([100, 200, 300])
d = 1000
assert awkward0.JaggedArray.zip(one=a, two=b).tolist() == [[{"one": 1.1, "two": 10}, {"one": 2.2, "two": 20}, {"one": 3.3, "two": 30}], [], [{"one": 4.4, "two": 40}, {"one": 5.5, "two": 50}]]
assert awkward0.JaggedArray.zip(one=b, two=a).tolist() == [[{"one": 10, "two": 1.1}, {"one": 20, "two": 2.2}, {"one": 30, "two": 3.3}], [], [{"one": 40, "two": 4.4}, {"one": 50, "two": 5.5}]]
assert awkward0.JaggedArray.zip(one=b, two=c).tolist() == [[{"one": 10, "two": 100}, {"one": 20, "two": 100}, {"one": 30, "two": 100}], [], [{"one": 40, "two": 300}, {"one": 50, "two": 300}]]
assert awkward0.JaggedArray.zip(one=b, two=d).tolist() == [[{"one": 10, "two": 1000}, {"one": 20, "two": 1000}, {"one": 30, "two": 1000}], [], [{"one": 40, "two": 1000}, {"one": 50, "two": 1000}]]
assert a.zip(b).tolist() == [[(1.1, 10), (2.2, 20), (3.3, 30)], [], [(4.4, 40), (5.5, 50)]]
assert b.zip(a).tolist() == [[(10, 1.1), (20, 2.2), (30, 3.3)], [], [(40, 4.4), (50, 5.5)]]
assert b.zip(c).tolist() == [[(10, 100), (20, 100), (30, 100)], [], [(40, 300), (50, 300)]]
assert b.zip(d).tolist() == [[(10, 1000), (20, 1000), (30, 1000)], [], [(40, 1000), (50, 1000)]]
def test_physics_matching(self):
gen_pt = awkward0.fromiter([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
reco_pt = awkward0.fromiter([[20.2, 10.1, 30.3, 50.5], [50.5], [50.5]])
gen_eta = awkward0.fromiter([[-3.0, -2.0, 2.0], [], [-1.0, 1.0]])
reco_eta = awkward0.fromiter([[-2.2, -3.3, 2.2, 0.0], [0.0], [1.1]])
gen_phi = awkward0.fromiter([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
reco_phi = awkward0.fromiter([[0.1, -1.4, 1.4, 0.78], [0.78], [-0.77]])
gen = uproot_methods.TLorentzVectorArray.from_ptetaphim(
gen_pt, gen_eta, gen_phi, 0.2)
reco = uproot_methods.TLorentzVectorArray.from_ptetaphim(
reco_pt, reco_eta, reco_phi, 0.2)
("gen", gen)
("reco", reco)
("gen.cross(reco)", gen.cross(reco))
pairing = gen.cross(reco, nested=True)
("pairing = gen.cross(reco, nested=True)", gen.cross(reco,
nested=True))
metric = pairing.i0.delta_r(pairing.i1)
("metric = pairing.i0.delta_r(pairing.i1)", metric)
index_of_minimized = metric.argmin()
("index_of_minimized = metric.argmin()", index_of_minimized)
assert index_of_minimized.tolist() == [[[1], [0], [2]], [], [[0], [0]]]
("metric[index_of_minimized]", metric[index_of_minimized])
passes_cut = (metric[index_of_minimized] < 0.5)
("passes_cut = (metric[index_of_minimized] < 0.5)", passes_cut)
assert passes_cut.tolist() == [[[True], [True], [True]], [],
[[False], [True]]]
best_pairings_that_pass_cut = pairing[index_of_minimized][passes_cut]
("best_pairings_that_pass_cut = pairing[index_of_minimized][passes_cut]",
best_pairings_that_pass_cut)
genrecos = best_pairings_that_pass_cut.flatten(axis=1)
("genrecos = best_pairings_that_pass_cut.flatten(axis=1)", genrecos)
("genrecos.counts", genrecos.counts)
("gen.counts", gen.counts)
assert genrecos.counts.tolist() == [3, 0, 1]
assert gen.counts.tolist() == [3, 0, 2]
("genrecos.i0.pt", genrecos.i0.pt)
("genrecos.i1.pt", genrecos.i1.pt)
def make_jets_and_electrons():
jet_m = awkward0.fromiter([[60.0, 70.0, 80.0], [], [90.0, 100.0]])
jet_pt = awkward0.fromiter([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
e_pt = awkward0.fromiter([[20.2, 50.5], [50.5], [50.5]])
jet_eta = awkward0.fromiter([[-3.0, -2.0, 2.0], [], [-1.0, 1.0]])
e_eta = awkward0.fromiter([[-2.2, 0.0], [0.0], [1.1]])
jet_phi = awkward0.fromiter([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
e_phi = awkward0.fromiter([[0.1, 0.78], [0.78], [-0.77]])
jets = uproot_methods.TLorentzVectorArray.from_ptetaphim(
jet_pt, jet_eta, jet_phi, jet_m)
electrons = uproot_methods.TLorentzVectorArray.from_ptetaphim(
e_pt, e_eta, e_phi, 0.000511)
return jets, electrons
def test_likenumpy_slices(self):
print()
np = numpy.array([[1, 10, 100], [2, 20, 200], [3, 30, 300]])
aw = awkward0.fromiter(np)
assert np.tolist() == aw.tolist()
assert np[:2].tolist() == aw[:2].tolist()
assert np[:2, :2].tolist() == aw[:2, :2].tolist()
assert np[:2, 2].tolist() == aw[:2, 2].tolist()
assert np[2, :2].tolist() == aw[2, :2].tolist()
assert np[:2, [0, 1]].tolist() == aw[:2, [0, 1]].tolist()
assert np[[0, 1], :2].tolist() == aw[[0, 1], :2].tolist()
assert np[:2, [0, 1, 2]].tolist() == aw[:2, [0, 1, 2]].tolist()
assert np[[0, 1, 2], :2].tolist() == aw[[0, 1, 2], :2].tolist()
assert np[[0, 1], [0, 1]].tolist() == aw[[0, 1], [0, 1]].tolist()
assert np[[0, 1, 2], [0, 1, 2]].tolist() == aw[[0, 1, 2], [0, 1, 2]].tolist()
assert np[:2, [True, False, True]].tolist() == aw[:2, [True, False, True]].tolist()
assert np[[True, False, True], :2].tolist() == aw[[True, False, True], :2].tolist()
assert np[[True, False, True], [True, False, True]].tolist() == aw[[True, False, True], [True, False, True]].tolist()
np = numpy.array([[[1, 10, 100], [2, 20, 200], [3, 30, 300]], [[4, 40, 400], [5, 50, 500], [6, 60, 600]], [[7, 70, 700], [8, 80, 800], [9, 90, 900]]])
aw = awkward0.fromiter(np)
assert np.tolist() == aw.tolist()
assert np[:2].tolist() == aw[:2].tolist()
assert np[:2, :2].tolist() == aw[:2, :2].tolist()
assert np[:2, 2].tolist() == aw[:2, 2].tolist()
assert np[2, :2].tolist() == aw[2, :2].tolist()
assert np[:2, [0, 1]].tolist() == aw[:2, [0, 1]].tolist()
assert np[[0, 1], :2].tolist() == aw[[0, 1], :2].tolist()
assert np[:2, [0, 1, 2]].tolist() == aw[:2, [0, 1, 2]].tolist()
assert np[[0, 1, 2], :2].tolist() == aw[[0, 1, 2], :2].tolist()
assert np[[0, 1], [0, 1]].tolist() == aw[[0, 1], [0, 1]].tolist()
assert np[[0, 1, 2], [0, 1, 2]].tolist() == aw[[0, 1, 2], [0, 1, 2]].tolist()
assert np[:2, [True, False, True]].tolist() == aw[:2, [True, False, True]].tolist()
assert np[[True, False, True], :2].tolist() == aw[[True, False, True], :2].tolist()
assert np[[True, False, True], [True, False, True]].tolist() == aw[[True, False, True], [True, False, True]].tolist()
assert np[:2, :2, 0].tolist() == aw[:2, :2, 0].tolist()
assert np[:2, 2, 0].tolist() == aw[:2, 2, 0].tolist()
assert np[2, :2, 0].tolist() == aw[2, :2, 0].tolist()
assert np[:2, [0, 1], 0].tolist() == aw[:2, [0, 1], 0].tolist()
assert np[[0, 1], :2, 0].tolist() == aw[[0, 1], :2, 0].tolist()
assert np[:2, [0, 1, 2], 0].tolist() == aw[:2, [0, 1, 2], 0].tolist()
assert np[[0, 1, 2], :2, 0].tolist() == aw[[0, 1, 2], :2, 0].tolist()
assert np[[0, 1], [0, 1], 0].tolist() == aw[[0, 1], [0, 1], 0].tolist()
assert np[[0, 1, 2], [0, 1, 2], 0].tolist() == aw[[0, 1, 2], [0, 1, 2], 0].tolist()
assert np[:2, [True, False, True], 0].tolist() == aw[:2, [True, False, True], 0].tolist()
assert np[[True, False, True], :2, 0].tolist() == aw[[True, False, True], :2, 0].tolist()
assert np[[True, False, True], [True, False, True], 0].tolist() == aw[[True, False, True], [True, False, True], 0].tolist()
assert np[:2, :2, 1].tolist() == aw[:2, :2, 1].tolist()
assert np[:2, 2, 1].tolist() == aw[:2, 2, 1].tolist()
assert np[2, :2, 1].tolist() == aw[2, :2, 1].tolist()
assert np[:2, [0, 1], 1].tolist() == aw[:2, [0, 1], 1].tolist()
assert np[[0, 1], :2, 1].tolist() == aw[[0, 1], :2, 1].tolist()
assert np[:2, [0, 1, 2], 1].tolist() == aw[:2, [0, 1, 2], 1].tolist()
assert np[[0, 1, 2], :2, 1].tolist() == aw[[0, 1, 2], :2, 1].tolist()
assert np[[0, 1], [0, 1], 1].tolist() == aw[[0, 1], [0, 1], 1].tolist()
assert np[[0, 1, 2], [0, 1, 2], 1].tolist() == aw[[0, 1, 2], [0, 1, 2], 1].tolist()
assert np[:2, [True, False, True], 1].tolist() == aw[:2, [True, False, True], 1].tolist()
assert np[[True, False, True], :2, 1].tolist() == aw[[True, False, True], :2, 1].tolist()
assert np[[True, False, True], [True, False, True], 1].tolist() == aw[[True, False, True], [True, False, True], 1].tolist()
def test_jagged_sort(self):
a = awkward0.fromiter([[2.,3.,1.], [4., -numpy.inf, 5.], [numpy.inf, 4., numpy.nan, -numpy.inf], [numpy.nan], [3., None, 4., -1.]])
assert a.argsort().tolist() == [[1, 0, 2], [2, 0, 1], [0, 1, 3, 2], [0], [2, 0, 3]]
assert a.argsort(True).tolist() == [[2, 0, 1], [1, 0, 2], [3, 1, 0, 2], [0], [3, 0, 2]]
def test_jagged_parents(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3, 4.4, 5.5], [], [6.6, 7.7, 8.8], [9.9]])
assert a.parents.tolist() == [0, 0, 0, 0, 0, 2, 2, 2, 3]
b = a[[False, True, False, True]]
assert b.parents.tolist() == [-1, -1, -1, -1, -1, -1, -1, -1, 1]
def test_jagged_localindex(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3, 4.4, 5.5], [], [6.6, 7.7, 8.8], [9.9]])
assert a.localindex.tolist() == [[0, 1, 2, 3, 4], [], [0, 1, 2], [0]]
b = a[[False, True, False, True]]
assert b.localindex.tolist() == [[], [0]]
def test_jagged_fillna(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
assert a.pad(4).fillna(999).tolist() == [[1.1, 2.2, 3.3, 999], [999, 999, 999, 999], [4.4, 5.5, 999, 999]]
assert a.pad(4, numpy.ma.masked).fillna(999).regular().tolist() == [[1.1, 2.2, 3.3, 999], [999, 999, 999, 999], [4.4, 5.5, 999, 999]]
def test_jagged_distincts_argnested(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
assert a.distincts().tolist() == [[(1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5)]]
assert a.argdistincts().tolist() == [[(0, 1), (0, 2), (1, 2)], [], [(0, 1)]]
assert a.distincts(nested=True).tolist() == [[[(1.1, 2.2), (1.1, 3.3)], [(2.2, 3.3)]], [], [[(4.4, 5.5)]]]
assert a.argdistincts(nested=True).tolist() == [[[(0, 1), (0, 2)], [(1, 2)]], [], [[(0, 1)]]]
def test_jagged_pairs_argnested(self):
a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
assert a.pairs().tolist() == [[(1.1, 1.1), (1.1, 2.2), (1.1, 3.3), (2.2, 2.2), (2.2, 3.3), (3.3, 3.3)], [], [(4.4, 4.4), (4.4, 5.5), (5.5, 5.5)]]
assert a.argpairs().tolist() == [[(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)], [], [(0, 0), (0, 1), (1, 1)]]
assert a.pairs(nested=True).tolist() == [[[(1.1, 1.1), (1.1, 2.2), (1.1, 3.3)], [(2.2, 2.2), (2.2, 3.3)], [(3.3, 3.3)]], [], [[(4.4, 4.4), (4.4, 5.5)], [(5.5, 5.5)]]]
assert a.argpairs(nested=True).tolist() == [[[(0, 0), (0, 1), (0, 2)], [(1, 1), (1, 2)], [(2, 2)]], [], [[(0, 0), (0, 1)], [(1, 1)]]]
stardicts.append({
"name": starname,
"ra": row["star", "right asc. (deg)"],
"dec": row["star", "declination (deg)"],
"dist": row["star", "distance (pc)"],
"mass": row["star", "mass (solar)"],
"radius": row["star", "radius (solar)"],
"planets": []
})
stardicts[-1]["planets"].append({
"name": planetname,
"orbit": row["planet", "orbit (AU)"],
"eccen": row["planet", "eccen."],
"period": row["planet", "period (days)"],
"mass": row["planet", "mass (Jupiter)"],
"radius": row["planet", "radius (Jupiter)"]
})
stars = awkward0.fromiter(stardicts)
arrowstars = awkward0.toarrow(stars)
with open("exoplanets.arrow", "wb") as sink:
writer = pyarrow.RecordBatchFileWriter(sink, arrowstars.schema)
for batch in arrowstars.to_batches():
writer.write_batch(batch)
writer.close()
# to read it back again:
awkward0.fromarrow(
pyarrow.ipc.open_file(open("exoplanets.arrow", "rb")).get_batch(0))
slices = [
2,
slice(None),
slice(2, 4),
slice(1, None, 2),
slice(None, None, -1),
numpy.array([2, 0, 0]),
numpy.array([3, 1, 2]),
numpy.array([True, False, True, True]),
numpy.array([True, True, True, False])
]
a = numpy.arange(4**4).reshape(4, 4, 4, 4)
a2 = awkward0.fromiter(a)
for x in slices:
print(x)
assert a[x, ].tolist() == getitem_enter(a2, (x, )).tolist()
for y in slices:
print(x, y)
assert a[x, y].tolist() == getitem_enter(a2, (x, y)).tolist()
for z in slices:
print(x, y, z)
assert a[x, y, z].tolist() == getitem_enter(a2, (x, y, z)).tolist()
tmp = a2.content.content.content.tolist()
tmp.insert(4, 999)
a2.content.content.stops = a2.content.content.stops.copy()
a2.content.content.starts[1:] += 1