From a9feab8760e3c9591b76b76badb3ba1f6a5d0c6c Mon Sep 17 00:00:00 2001
From: Tamas Gal <tgal@km3net.de>
Date: Wed, 8 Jan 2020 10:25:01 +0100
Subject: [PATCH] Add preliminary attempts for slicing magic
---
km3io/offline.py | 6 ++++++
tests/test_offline.py | 5 +++++
2 files changed, 11 insertions(+)
diff --git a/km3io/offline.py b/km3io/offline.py
index 8bd986d..7d42d99 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -1,5 +1,6 @@
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)
diff --git a/tests/test_offline.py b/tests/test_offline.py
index 5257b4d..69867ed 100644
--- a/tests/test_offline.py
+++ b/tests/test_offline.py
@@ -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):
--
GitLab