Skip to content
Snippets Groups Projects

Resolve "uproot4 integration"

Merged Resolve "uproot4 integration"
58-uproot4-integration-2 into master
Merged Tamas Gal requested to merge 58-uproot4-integration-2 into master
Compare
and
Show latest version
1 file
+ 16
− 16
Compare changes
  • Side-by-side
  • Inline
tests/test_tools.py
+ 125
− 87
@@ -18,7 +18,6 @@ from km3io.tools import (
uniquecount,
fitinf,
count_nested,
_find,
mask,
best_track,
get_w2list_param,
@@ -28,9 +27,16 @@ from km3io.tools import (
best_aashower,
best_dusjshower,
is_cc,
usr,
)
OFFLINE_FILE = OfflineReader(data_path("offline/km3net_offline.root"))
OFFLINE_USR = OfflineReader(data_path("offline/usr-sample.root"))
OFFLINE_MC_TRACK_USR = OfflineReader(
data_path(
"offline/mcv5.11r2.gsg_muonCChigherE-CC_50-5000GeV.km3_AAv1.jterbr00004695.jchain.aanet.498.root"
)
)
GENHEN_OFFLINE_FILE = OfflineReader(
data_path("offline/mcv5.1.genhen_anumuNC.sirene.jte.jchain.aashower.sample.root")
)
@@ -58,18 +64,12 @@ class TestFitinf(unittest.TestCase):
assert beta[1] == self.best_fit[1][0]
assert beta[2] == self.best_fit[2][0]
def test_fitinf_from_one_event_and_one_track(self):
beta = fitinf(kfit.JGANDALF_BETA0_RAD, self.tracks[0][0])
assert beta == self.tracks[0][0].fitinf[0]
class TestBestTrackSelection(unittest.TestCase):
def setUp(self):
self.events = OFFLINE_FILE.events
self.one_event = OFFLINE_FILE.events[0]
@unittest.skip
def test_best_track_selection_from_multiple_events_with_explicit_stages_in_list(
self,
):
@@ -102,58 +102,46 @@ class TestBestTrackSelection(unittest.TestCase):
assert best3.rec_stages[2] is None
assert best3.rec_stages[3] is None
@unittest.skip
def test_best_track_selection_from_multiple_events_with_explicit_stages_in_set(
def test_best_track_selection_from_multiple_events_with_a_set_of_stages(
self,
):
best = best_track(self.events.tracks, stages={1, 3, 4, 5})
assert len(best) == 10
assert best.rec_stages[0].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[1].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[2].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[3].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert best.rec_stages[1].tolist() == [1, 3, 5, 4]
assert best.rec_stages[2].tolist() == [1, 3, 5, 4]
assert best.rec_stages[3].tolist() == [1, 3, 5, 4]
# test with a shorter set of rec_stages
best2 = best_track(self.events.tracks, stages={1, 3})
assert len(best2) == 10
assert best2.rec_stages[0].tolist() == [[1, 3]]
assert best2.rec_stages[1].tolist() == [[1, 3]]
assert best2.rec_stages[2].tolist() == [[1, 3]]
assert best2.rec_stages[3].tolist() == [[1, 3]]
# test the irrelevance of order in rec_stages in sets
best3 = best_track(self.events.tracks, stages={3, 1})
assert len(best3) == 10
assert best3.rec_stages[0].tolist() == [[1, 3]]
assert best3.rec_stages[1].tolist() == [[1, 3]]
assert best3.rec_stages[2].tolist() == [[1, 3]]
assert best3.rec_stages[3].tolist() == [[1, 3]]
for rec_stages in best2.rec_stages:
for stage in {1, 3}:
assert stage in rec_stages
def test_best_track_selection_from_multiple_events_with_start_end(self):
best = best_track(self.events.tracks, startend=(1, 4))
assert len(best) == 10
assert best.rec_stages[0].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[1].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[2].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[3].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert best.rec_stages[1].tolist() == [1, 3, 5, 4]
assert best.rec_stages[2].tolist() == [1, 3, 5, 4]
assert best.rec_stages[3].tolist() == [1, 3, 5, 4]
# test with shorter stages
best2 = best_track(self.events.tracks, startend=(1, 3))
assert len(best2) == 10
assert best2.rec_stages[0].tolist() == [[1, 3]]
assert best2.rec_stages[1].tolist() == [[1, 3]]
assert best2.rec_stages[2].tolist() == [[1, 3]]
assert best2.rec_stages[3].tolist() == [[1, 3]]
assert best2.rec_stages[0].tolist() == [1, 3]
assert best2.rec_stages[1].tolist() == [1, 3]
assert best2.rec_stages[2].tolist() == [1, 3]
assert best2.rec_stages[3].tolist() == [1, 3]
# test the importance of start as a real start of rec_stages
best3 = best_track(self.events.tracks, startend=(0, 3))
@@ -179,23 +167,20 @@ class TestBestTrackSelection(unittest.TestCase):
# stages as a list
best = best_track(self.one_event.tracks, stages=[1, 3, 5, 4])
assert len(best) == 1
assert best.lik == ak.max(self.one_event.tracks.lik)
assert np.allclose(best.rec_stages[0].tolist(), [1, 3, 5, 4])
assert np.allclose(best.rec_stages.tolist(), [1, 3, 5, 4])
# stages as a set
best2 = best_track(self.one_event.tracks, stages={1, 3, 4, 5})
assert len(best2) == 1
assert best2.lik == ak.max(self.one_event.tracks.lik)
assert np.allclose(best2.rec_stages[0].tolist(), [1, 3, 5, 4])
assert np.allclose(best2.rec_stages.tolist(), [1, 3, 5, 4])
# stages with start and end
best3 = best_track(self.one_event.tracks, startend=(1, 4))
assert len(best3) == 1
assert best3.lik == ak.max(self.one_event.tracks.lik)
assert np.allclose(best3.rec_stages[0].tolist(), [1, 3, 5, 4])
assert np.allclose(best3.rec_stages.tolist(), [1, 3, 5, 4])
def test_best_track_on_slices_one_event(self):
tracks_slice = self.one_event.tracks[self.one_event.tracks.rec_type == 4000]
@@ -203,63 +188,90 @@ class TestBestTrackSelection(unittest.TestCase):
# test stages with list
best = best_track(tracks_slice, stages=[1, 3, 5, 4])
assert len(best) == 1
assert best.lik == ak.max(tracks_slice.lik)
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert best.rec_stages.tolist() == [1, 3, 5, 4]
# test stages with set
best2 = best_track(tracks_slice, stages={1, 3, 4, 5})
assert len(best2) == 1
assert best2.lik == ak.max(tracks_slice.lik)
assert best2.rec_stages[0].tolist() == [1, 3, 5, 4]
assert best2.rec_stages.tolist() == [1, 3, 5, 4]
def test_best_track_on_slices_with_start_end_one_event(self):
tracks_slice = self.one_event.tracks[0:5]
best = best_track(tracks_slice, startend=(1, 4))
assert len(best) == 1
assert best.lik == ak.max(tracks_slice.lik)
assert best.rec_stages[0][0] == 1
assert best.rec_stages[0][-1] == 4
assert best.rec_stages[0] == 1
assert best.rec_stages[-1] == 4
def test_best_track_on_slices_with_explicit_rec_stages_one_event(self):
tracks_slice = self.one_event.tracks[0:5]
best = best_track(tracks_slice, stages=[1, 3, 5, 4])
assert best.lik == ak.max(tracks_slice.lik)
assert best.rec_stages[0][0] == 1
assert best.rec_stages[0][-1] == 4
assert best.rec_stages[0] == 1
assert best.rec_stages[-1] == 4
@unittest.skip
def test_best_track_on_slices_multiple_events(self):
tracks_slice = self.events.tracks[0:5]
tracks_slice = self.events[0:5].tracks
# stages in list
best = best_track(tracks_slice, stages=[1, 3, 5, 4])
assert len(best) == 5
assert best.lik == ak.max(tracks_slice.lik)
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert np.allclose(
best.lik.tolist(),
[
294.6407542676734,
96.75133289411137,
560.2775306614813,
278.2872951665753,
99.59098153341449,
],
)
for i in range(len(best)):
assert best.rec_stages[i].tolist() == [1, 3, 5, 4]
# stages in set
best = best_track(tracks_slice, stages={1, 3, 4, 5})
best = best_track(tracks_slice, stages={3, 4, 5})
assert len(best) == 5
assert best.lik == ak.max(tracks_slice.lik)
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert np.allclose(
best.lik.tolist(),
[
294.6407542676734,
96.75133289411137,
560.2775306614813,
278.2872951665753,
99.59098153341449,
],
)
for i in range(len(best)):
assert best.rec_stages[i].tolist() == [1, 3, 5, 4]
# using start and end
best = best_track(tracks_slice, startend=(1, 4))
start, end = (1, 4)
best = best_track(tracks_slice, startend=(start, end))
assert len(best) == 5
assert best.lik == ak.max(tracks_slice.lik)
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert np.allclose(
best.lik.tolist(),
[
294.6407542676734,
96.75133289411137,
560.2775306614813,
278.2872951665753,
99.59098153341449,
],
)
for i in range(len(best)):
rs = best.rec_stages[i].tolist()
assert rs[0] == start
assert rs[-1] == end
def test_best_track_raises_when_unknown_stages(self):
with self.assertRaises(ValueError):
@@ -276,10 +288,10 @@ class TestBestJmuon(unittest.TestCase):
assert len(best) == 10
assert best.rec_stages[0].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[1].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[2].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[3].tolist() == [[1, 3, 5, 4]]
assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
assert best.rec_stages[1].tolist() == [1, 3, 5, 4]
assert best.rec_stages[2].tolist() == [1, 3, 5, 4]
assert best.rec_stages[3].tolist() == [1, 3, 5, 4]
assert best.lik[0] == ak.max(OFFLINE_FILE.events.tracks.lik[0])
assert best.lik[1] == ak.max(OFFLINE_FILE.events.tracks.lik[1])
@@ -378,28 +390,18 @@ class TestRecStagesMasks(unittest.TestCase):
self.tracks = OFFLINE_FILE.events.tracks
def test_find(self):
builder = ak.ArrayBuilder()
_find(self.nested, ak.Array([1, 2, 3]), builder)
labels = builder.snapshot()
assert labels[0][0] == 1
assert labels[0][1] == 1
assert labels[0][2] == 0
assert labels[1][0] == 0
def test_mask_with_explicit_rec_stages_in_list_with_multiple_events(self):
rec_stages = self.tracks.rec_stages
stages = [1, 3, 5, 4]
masks = mask(self.tracks, stages=stages)
masks = mask(self.tracks.rec_stages, sequence=stages)
assert masks[0][0] == all(rec_stages[0][0] == ak.Array(stages))
assert masks[1][0] == all(rec_stages[1][0] == ak.Array(stages))
assert masks[0][1] == False
def test_mask_with_explicit_rec_stages_in_set_with_multiple_events(self):
stages = {1, 3, 4, 5}
masks = mask(self.tracks, stages=stages)
def test_mask_with_atleast_on_multiple_events(self):
stages = [1, 3, 4, 5]
masks = mask(self.tracks.rec_stages, atleast=stages)
tracks = self.tracks[masks]
assert 1 in tracks.rec_stages[0][0]
@@ -410,7 +412,7 @@ class TestRecStagesMasks(unittest.TestCase):
def test_mask_with_start_and_end_of_rec_stages_with_multiple_events(self):
rec_stages = self.tracks.rec_stages
stages = [1, 3, 5, 4]
masks = mask(self.tracks, startend=(1, 4))
masks = mask(self.tracks.rec_stages, startend=(1, 4))
assert masks[0][0] == all(rec_stages[0][0] == ak.Array(stages))
assert masks[1][0] == all(rec_stages[1][0] == ak.Array(stages))
@@ -420,7 +422,7 @@ class TestRecStagesMasks(unittest.TestCase):
rec_stages = self.tracks.rec_stages[0][0]
stages = [1, 3, 5, 4]
track = self.tracks[0]
masks = mask(track, startend=(1, 4))
masks = mask(track.rec_stages, startend=(1, 4))
assert track[masks].rec_stages[0][0] == 1
assert track[masks].rec_stages[0][-1] == 4
@@ -429,20 +431,37 @@ class TestRecStagesMasks(unittest.TestCase):
rec_stages = self.tracks.rec_stages[0][0]
stages = [1, 3]
track = self.tracks[0]
masks = mask(track, stages=stages)
masks = mask(track.rec_stages, sequence=stages)
assert track[masks].rec_stages[0][0] == stages[0]
assert track[masks].rec_stages[0][1] == stages[1]
def test_mask_raises_when_too_many_inputs(self):
with self.assertRaises(ValueError):
mask(self.tracks, startend=(1, 4), stages=[1, 3, 5, 4])
def test_mask_raises_when_no_inputs(self):
with self.assertRaises(ValueError):
mask(self.tracks)
class TestMask(unittest.TestCase):
def test_minmax_2dim_mask(self):
arr = ak.Array([[1, 2, 3, 4], [3, 4, 5], [1, 2, 5]])
m = mask(arr, minmax=(1, 4))
self.assertListEqual(m.tolist(), [True, False, False])
def test_minmax_3dim_mask(self):
arr = ak.Array([[[1, 2, 3, 4], [3, 4, 5], [1, 2, 5]], [[1, 2, 3]]])
m = mask(arr, minmax=(1, 4))
self.assertListEqual(m.tolist(), [[True, False, False], [True]])
def test_minmax_4dim_mask(self):
arr = ak.Array(
[[[[1, 2, 3, 4], [3, 4, 5], [1, 2, 5]], [[1, 2, 3]]], [[[1, 9], [3, 3]]]]
)
m = mask(arr, minmax=(1, 4))
self.assertListEqual(
m.tolist(), [[[True, False, False], [True]], [[False, True]]]
)
class TestUnique(unittest.TestCase):
def run_random_test_with_dtype(self, dtype):
max_range = 100
@@ -543,3 +562,22 @@ class TestIsCC(unittest.TestCase):
all(NC_file) == True
) # this test fails because the CC flags are not reliable in old files
self.assertTrue(all(CC_file) == True)
class TestUsr(unittest.TestCase):
def test_event_usr(self):
assert np.allclose(
[118.6302815337638, 44.33580521344907, 99.93916717621543],
usr(OFFLINE_USR.events, "CoC").tolist(),
)
assert np.allclose(
[37.51967774166617, -10.280346193553832, 13.67595659707355],
usr(OFFLINE_USR.events, "DeltaPosZ").tolist(),
)
def test_mc_tracks_usr(self):
assert np.allclose(
[0.0487],
usr(OFFLINE_MC_TRACK_USR.mc_tracks[0], "bx").tolist(),
atol=0.0001,
)