Adds subset option on load to selectively load some trains with by_index

Load.py

@@ -6,7 +6,7 @@
     Copyright (2019) SCS Team.
 """
 import numpy as np
-from karabo_data import RunDirectory
+from karabo_data import RunDirectory by_index
 import xarray as xr
 
 mnemonics = {
@@ -171,7 +171,8 @@ mnemonics = {
 }
 
 def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False,
-    validate=False, runpath='/gpfs/exfel/exp/{}/{}/{}/raw/r{:04d}/'):
+    validate=False, runpath='/gpfs/exfel/exp/{}/{}/{}/raw/r{:04d}/',
+    subset=by_index[:]):
     """ Load a run and extract the data. Output is an xarray with aligned trainIds
 
         Inputs:
@@ -186,13 +187,15 @@ def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False,
             validate: boolean, whether to run karabo-data-validate or not
             runpath: a string to fromat the run folder path with topic,
                 semesterNB, proposalNB and runNB
+            subset: a subset of train that can be load with by_index[:5] for the
+                first 5 trains.
 
         Outputs:
             res: an xarray DataSet with aligned trainIds
     """
 
     runFolder = runpath.format(topic, semesterNB, proposalNB, runNB)
-    run = RunDirectory(runFolder)
+    run = RunDirectory(runFolder).select_trains(subset)
 
     if validate:
         get_ipython().system('karabo-data-validate ' + runFolder)