Best reco hits events

This addresses issues #22 (closed) and #24 (closed)

This is ready for merge into master. Although I am not sure I understand why the code style fails here on git (I updated yapf on my computer and used it to correct the code style but maybe I am missing a detail here?).

km3io/offline.py

@@ -11,6 +11,7 @@ BASKET_CACHE_SIZE = 110 * 1024**2
 
 class OfflineKeys:
     """wrapper for offline keys"""
+
     def __init__(self, file_path):
         """OfflineKeys is a class that reads all the available keys in an offline
         file and adapts the keys format to Python format.
@@ -60,7 +61,7 @@ class OfflineKeys:
             except those found in fake branches.
         """
         if self._events_keys is None:
-            fake_branches = ['Evt', 'AAObject', 'TObject', 't']
+            fake_branches = set(['Evt', 'AAObject', 'TObject', 't'])
             t_baskets = ['t.fSec', 't.fNanoSec']
             tree = uproot.open(self._file_path)['E']['Evt']
             self._events_keys = [
@@ -80,9 +81,9 @@ class OfflineKeys:
             except those found in fake branches.
         """
         if self._hits_keys is None:
-            fake_branches = [
+            fake_branches = set([
                 'hits.usr', 'hits.usr_names'
-            ]  # to be treated like trks.usr and trks.usr_names
+            ])  # to be treated like trks.usr and trks.usr_names
             tree = uproot.open(self._file_path)['E']['hits']
             self._hits_keys = [
                 key.decode('utf8') for key in tree.keys()
@@ -103,9 +104,9 @@ class OfflineKeys:
         if self._tracks_keys is None:
             # a solution can be tree['trks.usr_data'].array(
             # uproot.asdtype(">i4"))
-            fake_branches = [
+            fake_branches = set([
                 'trks.usr_data', 'trks.usr', 'trks.usr_names'
-            ]  # can be accessed using tree['trks.usr_names'].array()
+            ])  # can be accessed using tree['trks.usr_names'].array()
             tree = uproot.open(self._file_path)['E']['Evt']['trks']
             self._tracks_keys = [
                 key.decode('utf8') for key in tree.keys()
@@ -124,7 +125,7 @@ class OfflineKeys:
             except those found in fake branches.
         """
         if self._mc_hits_keys is None:
-            fake_branches = ['mc_hits.usr', 'mc_hits.usr_names']
+            fake_branches = set(['mc_hits.usr', 'mc_hits.usr_names'])
             tree = uproot.open(self._file_path)['E']['Evt']['mc_hits']
             self._mc_hits_keys = [
                 key.decode('utf8') for key in tree.keys()
@@ -143,9 +144,9 @@ class OfflineKeys:
             except those found in fake branches.
         """
         if self._mc_tracks_keys is None:
-            fake_branches = [
-                'mc_trks.usr_data', 'mc_trks.usr', 'mc_trks.usr_names'
-            ]  # same solution as above can be used
+            fake_branches = set(
+                ['mc_trks.usr_data', 'mc_trks.usr',
+                 'mc_trks.usr_names'])  # same solution as above can be used
             tree = uproot.open(self._file_path)['E']['Evt']['mc_trks']
             self._mc_tracks_keys = [
                 key.decode('utf8') for key in tree.keys()
@@ -273,6 +274,7 @@ class OfflineKeys:
 
 class Reader:
     """Reader for one offline ROOT file"""
+
     def __init__(self, file_path):
         """ Reader class is an offline ROOT file reader. This class is a
         "very" low level I/O.
@@ -313,7 +315,7 @@ class Reader:
             do not contain data. Therefore, the keys corresponding to these
             fake branches are not read.
         """
-        if key not in self.keys.valid_keys and not isinstance(key, int):
+        if key not in set(self.keys.valid_keys) and not isinstance(key, int):
             raise KeyError(
                 "'{}' is not a valid key or is a fake branch.".format(key))
         return self._data[key]
@@ -340,6 +342,7 @@ class Reader:
 
 class OfflineReader:
     """reader for offline ROOT files"""
+
     def __init__(self, file_path, data=None):
         """ OfflineReader class is an offline ROOT file wrapper
 
@@ -378,8 +381,6 @@ class OfflineReader:
             if 'Head' in fobj:
                 self._header = {}
                 for n, x in fobj['Head']._map_3c_string_2c_string_3e_.items():
-                    print("{:15s} {}".format(n.decode("utf-8"),
-                                             x.decode("utf-8")))
                     self._header[n.decode("utf-8")] = x.decode("utf-8")
             else:
                 warnings.warn("Your file header has an unsupported format")
@@ -570,15 +571,134 @@ class OfflineReader:
                     str(stages)))
         else:
             fit_data = np.array([
-                i[k]
-                for i, j, k in zip(fit_info, rec_stages[:, 0], rec_stages[:,
-                                                                          1])
-                if k is not None
+                i[k] for i, j, k in zip(fit_info, rec_stages[:, 0],
+                                        rec_stages[:, 1]) if k is not None
             ])
             rec_array = np.core.records.fromarrays(fit_data.transpose(),
                                                    names=keys)
             return rec_array
 
+    def get_reco_hits(self, stages, keys):
+        """construct a dictionary of hits class data based on the reconstruction
+        stages of interest. For example, if the reconstruction stages of interest
+        are [1, 2, 3, 4, 5], then get_reco_hits method will select the hits data 
+        from the events that were reconstructed following these stages (i.e 
+        [1, 2, 3, 4, 5]).
+
+        Parameters
+        ----------
+        stages : list
+            list of reconstruction stages of interest. for example
+            [1, 2, 3, 4, 5].
+        keys : list of str
+            list of the hits class attributes.
+
+        Returns
+        -------
+        dict
+            dictionary of lazyarrays containing data for each hits attribute requested.
+
+        Raises
+        ------
+        ValueError
+            ValueError raised when the reconstruction stages of interest
+            are not found in the file.
+        """
+        lazy_d = {}
+        rec_stages = np.array(
+            [match for match in self._find_rec_stages(stages)])
+        mask = rec_stages[:, 1] != None
+        if np.all(rec_stages[:, 1] == None):
+            raise ValueError(
+                "The stages {} are not found in your file.".format(
+                    str(stages)))
+        else:
+            for key in keys:
+                lazy_d[key] = getattr(self.hits, key)[mask]
+        return lazy_d
+
+    def get_reco_events(self, stages, keys):
+        """construct a dictionary of events class data based on the reconstruction
+        stages of interest. For example, if the reconstruction stages of interest
+        are [1, 2, 3, 4, 5], then get_reco_events method will select the events data 
+        that were reconstructed following these stages (i.e [1, 2, 3, 4, 5]).
+
+        Parameters
+        ----------
+        stages : list
+            list of reconstruction stages of interest. for example
+            [1, 2, 3, 4, 5].
+        keys : list of str
+            list of the events class attributes.
+
+        Returns
+        -------
+        dict
+            dictionary of lazyarrays containing data for each events attribute requested.
+
+        Raises
+        ------
+        ValueError
+            ValueError raised when the reconstruction stages of interest
+            are not found in the file.
+        """
+        lazy_d = {}
+        rec_stages = np.array(
+            [match for match in self._find_rec_stages(stages)])
+        mask = rec_stages[:, 1] != None
+        if np.all(rec_stages[:, 1] == None):
+            raise ValueError(
+                "The stages {} are not found in your file.".format(
+                    str(stages)))
+        else:
+            for key in keys:
+                lazy_d[key] = getattr(self.events, key)[mask]
+        return lazy_d
+
+    def get_reco_tracks(self, stages, keys):
+        """construct a dictionary of tracks class data based on the reconstruction
+        stages of interest. For example, if the reconstruction stages of interest
+        are [1, 2, 3, 4, 5], then get_reco_tracks method will select tracks data 
+        from the events that were reconstructed following these stages (i.e 
+        [1, 2, 3, 4, 5]).
+
+        Parameters
+        ----------
+        stages : list
+            list of reconstruction stages of interest. for example
+            [1, 2, 3, 4, 5].
+        keys : list of str
+            list of the tracks class attributes.
+
+        Returns
+        -------
+        dict
+            dictionary of lazyarrays containing data for each tracks attribute requested.
+
+        Raises
+        ------
+        ValueError
+            ValueError raised when the reconstruction stages of interest
+            are not found in the file.
+        """
+        lazy_d = {}
+        rec_stages = np.array(
+            [match for match in self._find_rec_stages(stages)])
+        mask = rec_stages[:, 1] != None
+        if np.all(rec_stages[:, 1] == None):
+            raise ValueError(
+                "The stages {} are not found in your file.".format(
+                    str(stages)))
+        else:
+            for key in keys:
+                lazy_d[key] = np.array([
+                    i[k] for i, k in zip(
+                        getattr(self.tracks, key)[mask], rec_stages[:, 1]
+                        [mask])
+                ])
+
+        return lazy_d
+
     def _find_rec_stages(self, stages):
         """find the index of reconstruction stages of interest in a
         list of multiple reconstruction stages.
@@ -639,6 +759,7 @@ class OfflineReader:
 
 class OfflineEvents:
     """wrapper for offline events"""
+
     def __init__(self, keys, values):
         """wrapper for offline events.
 
@@ -673,6 +794,7 @@ class OfflineEvents:
 
 class OfflineEvent:
     """wrapper for an offline event"""
+
     def __init__(self, keys, values):
         """wrapper for one offline event.
 
@@ -700,6 +822,7 @@ class OfflineEvent:
 
 class OfflineHits:
     """wrapper for offline hits"""
+
     def __init__(self, keys, values):
         """wrapper for offline hits.
 
@@ -734,6 +857,7 @@ class OfflineHits:
 
 class OfflineHit:
     """wrapper for an offline hit"""
+
     def __init__(self, keys, values):
         """wrapper for one offline hit.
 
@@ -770,6 +894,7 @@ class OfflineHit:
 
 class OfflineTracks:
     """wrapper for offline tracks"""
+
     def __init__(self, keys, values, fitparameters=None):
         """wrapper for offline tracks
 
@@ -814,6 +939,7 @@ class OfflineTracks:
 
 class OfflineTrack:
     """wrapper for an offline track"""
+
     def __init__(self, keys, values, fitparameters=None):
         """wrapper for one offline track.