Merge branch '35-support-for-aanet-usr-fields-in-offlinereader' into 'master'

Resolve "Support for aanet 'usr' fields in OfflineReader"

Closes #35

See merge request !20

CHANGELOG.rst

@@ -3,6 +3,10 @@ Unreleased changes
 
 Version 0
 ---------
+0.9.0 / 2020-03-03
+~~~~~~~~~~~~~~~~~~
+* Added support for the ``usr`` field of events
+
 0.8.3 / 2020-02-25
 ~~~~~~~~~~~~~~~~~~
 * The times of snapshot and triggered hits were parsed as big endian (standard)

README.rst

@@ -66,6 +66,8 @@ Tutorial
 
   * `reading events data <#reading-events-data>`__
 
+  * `reading usr data of events <#reading-usr-data-of-events>`__
+
   * `reading hits data <#reading-hits-data>`__
 
   * `reading tracks data <#reading-tracks-data>`__
@@ -75,6 +77,7 @@ Tutorial
   * `reading mc tracks data <#reading-mc-tracks-data>`__
 
 
+
 Introduction
 ------------
 
@@ -628,6 +631,44 @@ or the number of hits:
     176
 
 
+reading usr data of events
+""""""""""""""""""""""""""
+
+To access the ``usr`` data of events, use the ``.usr`` property which behaves
+like a dictionary and returns ``lazyarray``, compatible to the ``numpy.array``
+interface. The available keys can be accessed either as attributes or via a
+dictionary lookup:
+
+.. code-block:: python3
+
+    >>> import km3io
+    >>> f = km3io.OfflineReader("tests/samples/usr-sample.root")
+    >>> f.usr
+    <km3io.offline.Usr at 0x7efd53a41eb0>
+    >>> print(f.usr)
+    RecoQuality: [85.45957235835593 68.74744265572737 50.18704013646688]
+    RecoNDF: [37.0 37.0 29.0]
+    CoC: [118.6302815337638 44.33580521344907 99.93916717621543]
+    ToT: [825.0 781.0 318.0]
+    ChargeAbove: [176.0 278.0 53.0]
+    ChargeBelow: [649.0 503.0 265.0]
+    ChargeRatio: [0.21333333333333335 0.3559539052496799 0.16666666666666666]
+    DeltaPosZ: [37.51967774166617 -10.280346193553832 13.67595659707355]
+    FirstPartPosZ: [135.29499707179326 41.46665612378939 107.39596803432326]
+    LastPartPosZ: [97.77531933012709 51.747002317343224 93.72001143724971]
+    NSnapHits: [51.0 107.0 98.0]
+    NTrigHits: [30.0 32.0 14.0]
+    NTrigDOMs: [7.0 11.0 7.0]
+    NTrigLines: [6.0 5.0 4.0]
+    NSpeedVetoHits: [0.0 0.0 0.0]
+    NGeometryVetoHits: [0.0 0.0 0.0]
+    ClassficationScore: [0.16863382173469108 0.17944356593281038 0.08155750660727408]
+    >>> f.usr.DeltaPosZ
+    <ChunkedArray [37.51967774166617 -10.280346193553832 13.67595659707355] at 0x7efd54013eb0>
+    >>> f.usr['RecoQuality']
+    <ChunkedArray [85.45957235835593 68.74744265572737 50.18704013646688] at 0x7efd54034b50>
+
+
 reading hits data
 """""""""""""""""
 

km3io/offline.py

@@ -365,6 +365,7 @@ class OfflineReader:
         self._keys = None
         self._best_reco = None
         self._header = None
+        self._usr = None
 
     def __getitem__(self, item):
         return OfflineReader(file_path=self._file_path, data=self._data[item])
@@ -474,6 +475,12 @@ class OfflineReader:
                 fitparameters=self.keys.fitparameters)
         return self._mc_tracks
 
+    @property
+    def usr(self):
+        if self._usr is None:
+            self._usr = Usr(self._file_path)
+        return self._usr
+
     @property
     def best_reco(self):
         """returns the best reconstructed track fit data. The best fit is defined
@@ -816,6 +823,44 @@ class OfflineReader:
             yield i, j
 
 
+class Usr:
+    """Helper class to access AAObject usr stuff"""
+    def __init__(self, filepath):
+        self._f = uproot.open(filepath)
+        # Here, we assume that every event has the same names in the same order
+        # to massively increase the performance. This needs triple check if it's
+        # always the case; the usr-format is simply a very bad design.
+        try:
+            self._usr_names = [
+                n.decode("utf-8")
+                for n in self._f['E']['Evt']['usr_names'].array()[0]
+            ]
+        except (KeyError, IndexError):  # e.g. old aanet files
+            self._usr_names = []
+        else:
+            self._usr_idx_lookup = {
+                name: index
+                for index, name in enumerate(self._usr_names)
+            }
+            self._usr_data = self._f['E']['Evt']['usr'].lazyarray(
+                basketcache=uproot.cache.ThreadSafeArrayCache(
+                    BASKET_CACHE_SIZE))
+            for name in self._usr_names:
+                setattr(self, name, self[name])
+
+    def __getitem__(self, item):
+        return self._usr_data[:, self._usr_idx_lookup[item]]
+
+    def keys(self):
+        return self._usr_names
+
+    def __str__(self):
+        entries = []
+        for name in self.keys():
+            entries.append("{}: {}".format(name, self[name]))
+        return '\n'.join(entries)
+
+
 class OfflineEvents:
     """wrapper for offline events"""
     def __init__(self, keys, values):

tests/test_offline.py

@@ -7,6 +7,7 @@ from km3io import OfflineReader
 
 SAMPLES_DIR = Path(__file__).parent / 'samples'
 OFFLINE_FILE = SAMPLES_DIR / 'aanet_v2.0.0.root'
+OFFLINE_USR = SAMPLES_DIR / 'usr-sample.root'
 OFFLINE_NUMUCC = SAMPLES_DIR / "numucc.root"  # with mc data
 
 
@@ -478,3 +479,40 @@ class TestOfflineTrack(unittest.TestCase):
     def test_str(self):
         self.assertEqual(repr(self.track).split('\n\t')[0], 'offline track:')
         self.assertTrue("JGANDALF_LAMBDA" in repr(self.track))
+
+
+class TestUsr(unittest.TestCase):
+    def setUp(self):
+        self.f = OfflineReader(OFFLINE_USR)
+
+    def test_str(self):
+        print(self.f.usr)
+
+    def test_nonexistent_usr(self):
+        f = OfflineReader(SAMPLES_DIR / "daq_v1.0.0.root")
+        self.assertListEqual([], f.usr.keys())
+
+    def test_keys(self):
+        self.assertListEqual([
+            'RecoQuality', 'RecoNDF', 'CoC', 'ToT', 'ChargeAbove',
+            'ChargeBelow', 'ChargeRatio', 'DeltaPosZ', 'FirstPartPosZ',
+            'LastPartPosZ', 'NSnapHits', 'NTrigHits', 'NTrigDOMs',
+            'NTrigLines', 'NSpeedVetoHits', 'NGeometryVetoHits',
+            'ClassficationScore'
+        ], self.f.usr.keys())
+
+    def test_getitem(self):
+        assert np.allclose(
+            [118.6302815337638, 44.33580521344907, 99.93916717621543],
+            self.f.usr['CoC'])
+        assert np.allclose(
+            [37.51967774166617, -10.280346193553832, 13.67595659707355],
+            self.f.usr['DeltaPosZ'])
+
+    def test_attributes(self):
+        assert np.allclose(
+            [118.6302815337638, 44.33580521344907, 99.93916717621543],
+            self.f.usr.CoC)
+        assert np.allclose(
+            [37.51967774166617, -10.280346193553832, 13.67595659707355],
+            self.f.usr.DeltaPosZ)