Skip to content
Snippets Groups Projects
Commit 3bf1f243 authored by Tamas Gal's avatar Tamas Gal :speech_balloon:
Browse files

Final final refactoring

parent 486d30dd
No related branches found
No related tags found
1 merge request!45Adapt best track root access
Pipeline #14564 passed with warnings
Stage: test
Stage: coverage
Stage: doc
Hide whitespace changes
Inline Side-by-side
km3io/tools.py
+ 72
− 124
View file @ 3bf1f243
......@@ -151,24 +151,12 @@ def fitinf(fitparam, tracks):
return params[:, index]
def fitparams():
"""name of the fit parameters as defined in the official
KM3NeT-Dataformat.
Returns
-------
dict_keys
fit parameters keys.
"""
return kfit.keys()
def count_nested(Array, axis=0):
def count_nested(arr, axis=0):
"""Count elements in a nested awkward Array.
Parameters
----------
Array : awkward1.Array
arr : awkward1.Array
Array of data. Example tracks.fitinf or tracks.rec_stages.
axis : int, optional
axis = 0: to count elements in the outmost level of nesting.
......@@ -181,16 +169,18 @@ def count_nested(Array, axis=0):
counts of elements found in a nested awkward1 Array.
"""
if axis == 0:
return ak1.num(Array, axis=0)
return ak1.num(arr, axis=0)
if axis == 1:
return ak1.num(Array, axis=1)
return ak1.num(arr, axis=1)
if axis == 2:
return ak1.count(Array, axis=2)
return ak1.count(arr, axis=2)
def get_multiplicity(tracks, rec_stages):
"""Tracks selection based on specific reconstruction stages (for multiplicity
calculations).
"""Tracks selection based on specific reconstruction stages.
Counts how many tracks with the specific reconstructions stages
are found per event.
Parameters
----------
......@@ -214,79 +204,55 @@ def get_multiplicity(tracks, rec_stages):
return out
def best_track(tracks,
start_stages=None,
end_stages=None,
min_stages=None,
max_stages=None,
stages=None):
def best_track(tracks, startend=None, minmax=None, stages=None):
"""Best track selection.
Parameters
----------
tracks : km3io.offline.OfflineBranch
tracks of interest. tracks can be from multiple events, or from one event, or a slice of tracks.
start_stages : int, optional
the exact starting step of rec_stages, as in tracks.rec_stages.
end_stages : int, optional
the exact ending step of rec_stages, as in tracks.rec_stages.
min_stages : int, optional
the minimum range of rec_stages.
max_stages : int, optional
the maximum range of rec_stages.
Array of tracks or jagged array of tracks (multiple events).
startend: tuple(int, int), optional
The required first and last stage in tracks.rec_stages.
minmax: tuple(int, int), optional
The range (minimum and maximum) value of rec_stages to take into account.
stages : list or set, optional
either a list or a set of stages:
- list: the order of the rec_stages is conserved.
- if stages in a set, the order is irrelevant.
- list: the order of the rec_stages is respected.
- set: the order is irrelevant.
Returns
-------
km3io.offline.OfflineBranch
the best tracks based on the rec_stages selection. The logest track and the highest likelihood track is returned.
The best tracks based on the selection.
Raises
------
ValueError
valueError raised when:
- too many inputs specified.
- no inputs are specified.
- too many inputs specified.
- no inputs are specified.
"""
inputs = [stages, start_stages, end_stages, min_stages, max_stages]
min_max = [min_stages, max_stages]
start_end = [start_stages, end_stages]
inputs = (stages, startend, minmax)
if all(v is None for v in inputs):
raise ValueError("No reconstruction stages were specified")
raise ValueError("either stages, startend or minmax must be specified.")
if all(v is not None for v in inputs) or all(v is not None
for v in inputs[0:3]):
raise ValueError(
"Please specify either a range or a set of rec stages.")
if stages is not None and (startend is not None or minmax is not None):
raise ValueError("Please specify either a range or a set of rec stages.")
if all(v is None for v in inputs[1:]) and (stages is not None):
if stages is not None and startend is None and minmax is None:
selected_tracks = tracks[mask(tracks, stages=stages)]
if all(v is not None
for v in start_end) and all(v is None
for v in min_max) and (stages is None):
selected_tracks = tracks[mask(tracks,
start_stages=start_stages,
end_stages=end_stages)]
if startend is not None and minmax is None and stages is None:
selected_tracks = tracks[mask(tracks, startend=startend)]
if all(v is not None
for v in min_max) and all(v is None
for v in start_end) and (stages is None):
selected_tracks = tracks[mask(tracks,
min_stages=start_stages,
max_stages=end_stages)]
if minmax is not None and startend is None and stages is None:
selected_tracks = tracks[mask(tracks, minmax=minmax)]
return _max_lik_track(_longest_tracks(selected_tracks))
def _longest_tracks(tracks):
""" select the longest reconstructed track"""
"""Select the longest reconstructed track"""
if tracks.is_single:
stages_nesting_level = 1
tracks_nesting_level = 0
......@@ -303,7 +269,7 @@ def _longest_tracks(tracks):
def _max_lik_track(tracks):
""" select the track with the highest likelihood """
"""Select the track with the highest likelihood """
if tracks.is_single:
tracks_nesting_level = 0
else:
......@@ -312,13 +278,8 @@ def _max_lik_track(tracks):
return tracks[tracks.lik == ak1.max(tracks.lik, axis=tracks_nesting_level)]
def mask(tracks,
stages=None,
start_stages=None,
end_stages=None,
min_stages=None,
max_stages=None):
"""create a mask for tracks.rec_stages .
def mask(tracks, stages=None, startend=None, minmax=None):
"""Create a mask for tracks.rec_stages.
Parameters
----------
......@@ -326,45 +287,34 @@ def mask(tracks,
tracks, or one track, or slice of tracks, or slice of one track.
stages : list or set
reconstruction stages of interest:
- if stages is a list: the order of rec_stages in conserved.
- if stages is a set: the order of rec_stages in irrelevant.
start_stages : int, optional
the exact starting step of rec_stages, as in tracks.rec_stages.
end_stages : int, optional
the exact ending step of rec_stages, as in tracks.rec_stages.
min_stages : int, optional
the minimum range of rec_stages.
max_stages : int, optional
the maximum range of rec_stages.
- list: the order of rec_stages in respected.
- set: the order of rec_stages in irrelevant.
startend: tuple(int, int), optional
The required first and last stage in tracks.rec_stages.
minmax: tuple(int, int), optional
The range (minimum and maximum) value of rec_stages to take into account.
Returns
-------
awkward1.Array
awkward1.Array(bool)
an awkward1 Array mask where True corresponds to the positions
where stages were found. False otherwise.
Raises
------
ValueError
valueError raised when:
- too many inputs specified.
- no inputs are specified.
- too many inputs specified.
- no inputs are specified.
"""
inputs = [stages, start_stages, end_stages, min_stages, max_stages]
min_max = [min_stages, max_stages]
start_end = [start_stages, end_stages]
inputs = (stages, startend, minmax)
if all(v is None for v in inputs):
raise ValueError(
"either stages or (start_stages and end_stages) or (min_stages and max_stages) must be specified."
)
raise ValueError("either stages, startend or minmax must be specified.")
if all(v is not None for v in inputs) or all(v is not None
for v in inputs[0:3]):
raise ValueError("too many inputs are specified.")
if stages is not None and (startend is not None or minmax is not None):
raise ValueError("Please specify either a range or a set of rec stages.")
if (stages is not None) and all(v is None for v in inputs[1:]):
if stages is not None and startend is None and minmax is None:
if isinstance(stages, list):
# order of stages is conserved
return _mask_explicit_rec_stages(tracks, stages)
......@@ -373,22 +323,16 @@ def mask(tracks,
return _mask_rec_stages_in_range_min_max(tracks,
valid_stages=stages)
if all(v is not None
for v in start_end) and all(v is None
for v in min_max) and (stages is None):
return _mask_rec_stages_between_start_end(tracks, start_stages,
end_stages)
if startend is not None and minmax is None and stages is None:
return _mask_rec_stages_between_start_end(tracks, *startend)
if all(v is None
for v in start_end) and all(v is not None
for v in min_max) and (stages is None):
return _mask_rec_stages_in_range_min_max(tracks,
min_stages=min_stages,
max_stages=max_stages)
if minmax is not None and startend is None and stages is None:
return _mask_rec_stages_in_range_min_max(tracks, *minmax)
def _mask_rec_stages_between_start_end(tracks, start, end):
""" mask tracks.rec_stages that start exactly with start and end exactly with end. ie [start, a, b ...,z , end] """
"""Mask tracks.rec_stages that start exactly with start and end exactly
with end. ie [start, a, b ...,z , end] """
builder = ak1.ArrayBuilder()
if tracks.is_single:
_find_between_single(tracks.rec_stages, start, end, builder)
......@@ -400,7 +344,7 @@ def _mask_rec_stages_between_start_end(tracks, start, end):
@nb.jit(nopython=True)
def _find_between(rec_stages, start, end, builder):
""" find tracks.rec_stages where rec_stages[0] == start and rec_stages[-1] == end."""
"""Find tracks.rec_stages where rec_stages[0] == start and rec_stages[-1] == end."""
for s in rec_stages:
builder.begin_list()
......@@ -418,7 +362,8 @@ def _find_between(rec_stages, start, end, builder):
@nb.jit(nopython=True)
def _find_between_single(rec_stages, start, end, builder):
""" find tracks.rec_stages where rec_stages[0] == start and rec_stages[-1] == end in a single track. """
"""Find tracks.rec_stages where rec_stages[0] == start and
rec_stages[-1] == end in a single track. """
builder.begin_list()
for s in rec_stages:
......@@ -434,7 +379,7 @@ def _find_between_single(rec_stages, start, end, builder):
def _mask_explicit_rec_stages(tracks, stages):
"""mask explicit rec_stages .
"""Mask explicit rec_stages .
Parameters
----------
......@@ -494,7 +439,8 @@ def _find(rec_stages, stages, builder):
@nb.jit(nopython=True)
def _find_single(rec_stages, stages, builder):
"""construct an awkward1 array with the same structure as tracks.rec_stages.
"""Construct an awkward1 array with the same structure as tracks.rec_stages.
When stages are found, the Array is filled with value 1, otherwise it is filled
with value 0.
......@@ -525,7 +471,7 @@ def _find_single(rec_stages, stages, builder):
def best_jmuon(tracks):
"""select the best JMUON track.
"""Select the best JMUON track.
Parameters
----------
......@@ -545,7 +491,7 @@ def best_jmuon(tracks):
def best_jshower(tracks):
"""select the best JSHOWER track.
"""Select the best JSHOWER track.
Parameters
----------
......@@ -565,7 +511,7 @@ def best_jshower(tracks):
def best_aashower(tracks):
"""select the best AASHOWER track.
"""Select the best AASHOWER track.
Parameters
----------
......@@ -585,7 +531,7 @@ def best_aashower(tracks):
def best_dusjshower(tracks):
"""select the best DISJSHOWER track.
"""Select the best DISJSHOWER track.
Parameters
----------
......@@ -608,7 +554,7 @@ def _mask_rec_stages_in_range_min_max(tracks,
min_stages=None,
max_stages=None,
valid_stages=None):
"""mask tracks where rec_stages are withing the range(min, max).
"""Mask tracks where rec_stages are withing the range(min, max).
Parameters
----------
......@@ -642,9 +588,10 @@ def _mask_rec_stages_in_range_min_max(tracks,
@nb.jit(nopython=True)
def _find_in_range(rec_stages, valid_stages, builder):
"""construct an awkward1 array with the same structure as tracks.rec_stages.
When stages are within the range(min, max), the Array is filled with value 1, otherwise it is filled
with value 0.
"""Construct an awkward1 array with the same structure as tracks.rec_stages.
When stages are within the range(min, max), the Array is filled with
value 1, otherwise it is filled with value 0.
Parameters
----------
......@@ -676,9 +623,10 @@ def _find_in_range(rec_stages, valid_stages, builder):
@nb.jit(nopython=True)
def _find_in_range_single(rec_stages, valid_stages, builder):
"""construct an awkward1 array with the same structure as tracks.rec_stages.
When stages are within the range(min, max), the Array is filled with value 1, otherwise it is filled
with value 0.
"""Construct an awkward1 array with the same structure as tracks.rec_stages.
When stages are within the range(min, max), the Array is filled with
value 1, otherwise it is filled with value 0.
Parameters
----------
......
tests/test_tools.py
+ 12
− 23
View file @ 3bf1f243
......@@ -9,7 +9,7 @@ from km3net_testdata import data_path
from km3io import OfflineReader
from km3io.tools import (to_num, cached_property, unfold_indices, unique,
uniquecount, fitinf, fitparams, count_nested, _find,
uniquecount, fitinf, count_nested, _find,
mask, best_track, get_w2list_param, get_multiplicity,
best_jmuon, best_jshower, best_aashower,
best_dusjshower)
......@@ -36,11 +36,6 @@ class TestFitinf(unittest.TestCase):
assert best_beta[1] == self.best_fit[1][0]
assert best_beta[2] == self.best_fit[2][0]
def test_fitparams(self):
keys = set(fitparams())
assert "JGANDALF_BETA0_RAD" in keys
class TestBestTrackSelection(unittest.TestCase):
def setUp(self):
......@@ -110,7 +105,7 @@ class TestBestTrackSelection(unittest.TestCase):
assert best3.rec_stages[3] == [1, 3]
def test_best_track_selection_from_multiple_events_with_start_end(self):
best = best_track(self.events.tracks, start_stages=1, end_stages=4)
best = best_track(self.events.tracks, startend=(1, 4))
assert len(best) == 10
......@@ -120,7 +115,7 @@ class TestBestTrackSelection(unittest.TestCase):
assert best.rec_stages[3] == [1, 3, 5, 4]
# test with shorter stages
best2 = best_track(self.events.tracks, start_stages=1, end_stages=3)
best2 = best_track(self.events.tracks, startend=(1, 3))
assert len(best2) == 10
......@@ -130,7 +125,7 @@ class TestBestTrackSelection(unittest.TestCase):
assert best2.rec_stages[3] == [1, 3]
# test the importance of start as a real start of rec_stages
best3 = best_track(self.events.tracks, start_stages=0, end_stages=3)
best3 = best_track(self.events.tracks, startend=(0, 3))
assert len(best3) == 10
......@@ -140,7 +135,7 @@ class TestBestTrackSelection(unittest.TestCase):
assert best3.rec_stages[3] is None
# test the importance of end as a real end of rec_stages
best4 = best_track(self.events.tracks, start_stages=1, end_stages=10)
best4 = best_track(self.events.tracks, startend=(1, 10))
assert len(best4) == 10
......@@ -165,7 +160,7 @@ class TestBestTrackSelection(unittest.TestCase):
assert best2.rec_stages[0] == [1, 3, 5, 4]
# stages with start and end
best3 = best_track(self.one_event.tracks, start_stages=1, end_stages=4)
best3 = best_track(self.one_event.tracks, startend=(1, 4))
assert len(best3) == 1
assert best3.lik == ak.max(self.one_event.tracks.lik)
......@@ -193,7 +188,7 @@ class TestBestTrackSelection(unittest.TestCase):
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, start_stages=1, end_stages=4)
best = best_track(tracks_slice, startend=(1, 4))
assert len(best) == 1
assert best.lik == ak.max(tracks_slice.lik)
......@@ -228,7 +223,7 @@ class TestBestTrackSelection(unittest.TestCase):
assert best.rec_stages[0] == [1, 3, 5, 4]
# using start and end
best = best_track(tracks_slice, start_stages=1, end_stages=4)
best = best_track(tracks_slice, startend=(1, 4))
assert len(best) == 5
......@@ -241,10 +236,7 @@ class TestBestTrackSelection(unittest.TestCase):
def test_best_track_raises_when_too_many_inputs(self):
with self.assertRaises(ValueError):
best_track(self.events.tracks,
start_stages=1,
end_stages=4,
stages=[1, 3, 5, 4])
best_track(self.events.tracks, startend=(1, 4), stages=[1, 3, 5, 4])
class TestBestJmuon(unittest.TestCase):
......@@ -384,7 +376,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, start_stages=1, end_stages=4)
masks = mask(self.tracks, 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))
......@@ -394,7 +386,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, start_stages=1, end_stages=4)
masks = mask(track, startend=(1, 4))
assert track[masks].rec_stages[0][0] == 1
assert track[masks].rec_stages[0][-1] == 4
......@@ -410,10 +402,7 @@ class TestRecStagesMasks(unittest.TestCase):
def test_mask_raises_when_too_many_inputs(self):
with self.assertRaises(ValueError):
mask(self.tracks,
start_stages=1,
end_stages=4,
stages=[1, 3, 5, 4])
mask(self.tracks, startend=(1, 4), stages=[1, 3, 5, 4])
def test_mask_raises_when_no_inputs(self):
with self.assertRaises(ValueError):
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment