Add I/O functions to simplify benchmarking

docs/pipeline/index.rst

@@ -110,4 +110,4 @@ Reference/API
 .. automodapi:: protopipe.pipeline
     :no-inheritance-diagram:
     :include-all-objects:
-    :skip: EventSource
+    :skip: Path, Table, EventSource

protopipe/pipeline/__init__.py

@@ -4,3 +4,4 @@
 from .image_cleaning import *
 from .event_preparer import *
 from .utils import *
+from .io import *

protopipe/pipeline/io.py

@@ -0,0 +1,95 @@
+"""Utility functions mainly used in benchmarking notebooks."""
+
+from pathlib import Path
+
+import tables
+import pandas
+from astropy.table import Table
+
+
+def get_camera_names(input_directory=None,
+                     file_name=None):
+    """Read the names of the cameras.
+
+    Parameters
+    ==========
+    input_directory : str or pathlib.Path
+        Path of the input file.
+    file_name : str
+        Name of the input file.
+
+    Returns
+    =======
+    camera_names : list(str)
+        Table names as a list.
+    """
+    if (input_directory is None) or (file_name is None):
+        print("ERROR: check input")
+
+    input_file = Path(input_directory) / file_name
+
+    with tables.open_file(input_file, 'r') as f:
+        camera_names = [cam.name for cam in f.root]
+
+    return camera_names
+
+
+def read_protopipe_TRAINING_per_tel_type(input_directory=None,
+                                         file_name=None,
+                                         camera_names=None):
+    """Read a TRAINING file and extract the data per telescope type.
+
+    Parameters
+    ==========
+    input_directory : str or pathlib.Path
+        Path of the input directory where the TRAINING file is located.
+    file_name : str
+        Name of the input TRAINING file.
+
+    Returns
+    =======
+    dataFrames : dict(pandas.DataFrame)
+        Dictionary of tables per camera.
+    """
+    if (input_directory is None) or (file_name is None):
+        print("ERROR: check input")
+    if camera_names is None:
+        print("ERROR: no cameras specified")
+    # load DL1 images
+    input_file = Path(input_directory) / file_name
+    dataFrames = {}
+    for camera in camera_names:
+        dataFrames[camera] = pandas.read_hdf(input_file, f"/{camera}")
+    return dataFrames
+
+
+def read_TRAINING_per_tel_type_with_images(input_directory=None,
+                                           input_filename=None,
+                                           camera_names=None):
+    """Read a TRAINING file and extract the data per telescope type.
+
+    Parameters
+    ==========
+    input_directory : str or pathlib.Path
+        Path of the input directory where the TRAINING file is located.
+    file_name : str
+        Name of the input TRAINING file.
+    camera_names: list
+        List of camera names corresponding to the table names in the file.
+
+    Returns
+    =======
+    table : dict(astropy.Table)
+        Dictionary of astropy tables per camera.
+    """
+    input_file = Path(input_directory) / input_filename
+
+    table = {}
+
+    with tables.open_file(input_file, mode='r') as h5file:
+        for camera in camera_names:
+            table[camera] = Table()
+            for key in h5file.get_node(f"/{camera}").colnames:
+                table[camera][key] = h5file.get_node(f"/{camera}").col(key)
+
+    return table

protopipe/pipeline/utils.py

@@ -4,6 +4,7 @@
 import math
 import joblib
 
+import numpy
 import astropy.units as u
 import matplotlib.pyplot as plt
 import os.path as path
@@ -503,14 +504,36 @@ def load_models(path, cam_id_list):
     -------
     model_dict: dict
         Dictionary with `cam_id` as keys and pickled models as values.
-
     """
 
     model_dict = {}
     for key in cam_id_list:
-            try:
-                model_dict[key] = joblib.load(path.format(cam_id=key))
-            except IndexError:
-                model_dict[key] = joblib.load(path.format(key))
+        try:
+            model_dict[key] = joblib.load(path.format(cam_id=key))
+        except IndexError:
+            model_dict[key] = joblib.load(path.format(key))
 
     return model_dict
+
+
+def add_stats(data, ax, x=0.70, y=0.85, fontsize=10):
+    """Add a textbox containing statistical information."""
+    mu = data.mean()
+    median = np.median(data)
+    sigma = data.std()
+    textstr = '\n'.join((
+        r'$\mu=%.2f$' % (mu, ),
+        r'$\mathrm{median}=%.2f$' % (median, ),
+        r'$\sigma=%.2f$' % (sigma, )))
+
+    # these are matplotlib.patch.Patch properties
+    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
+
+    # place a text box in upper left in axes coords
+    ax.text(x, y,
+            textstr,
+            transform=ax.transAxes,
+            fontsize=fontsize,
+            horizontalalignment='left',
+            verticalalignment='center',
+            bbox=props)