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

Add preliminary attempts for slicing magic

parent 17bcedbc
No related branches found
No related tags found
1 merge request!6WIP: Resolve "Slicing of tracks"
Pipeline #8415 failed
import uproot
import numpy as np
import warnings
import km3io.definitions.trigger
import km3io.definitions.fitparameters
......@@ -794,6 +795,11 @@ class OfflineTracks:
if isinstance(item, int):
return OfflineTrack(self._keys, [v[item] for v in self._values],
fitparameters=self._fitparameters)
elif isinstance(item, list) and all(isinstance(i, str) for i in item):
cols = item
data = [getattr(self, c) for c in cols]
dtype = dict(names=cols, formats=[d.dtype for d in data])
return np.rec.fromarrays(data, dtype=dtype)
else:
return OfflineTracks(self._keys, [v[item] for v in self._values],
fitparameters=self._fitparameters)
......
......@@ -419,6 +419,11 @@ class TestOfflineTracks(unittest.TestCase):
self.assertListEqual(list(tracks.E[:, 0][_slice]),
list(tracks[_slice].E[:, 0]))
def test_slicing_via_columns(self):
tracks = self.tracks
data = tracks[['E', 'lik']]
assert 1 == data.E
class TestOfflineTrack(unittest.TestCase):
def setUp(self):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment