diff --git a/src/toolbox_scs/load.py b/src/toolbox_scs/load.py index 68cca3960239556e017107d67d62ab73e6ddfa72..0e6d5c029dbb12b85b9e72c67ca884d392fe764e 100644 --- a/src/toolbox_scs/load.py +++ b/src/toolbox_scs/load.py @@ -39,7 +39,7 @@ def load(proposalNB=None, runNB=None, subFolder='raw', display=False, validate=False, - subset=ed.by_index[:], + subset=None, rois={}, extract_adq412=True, extract_fadc=True, @@ -55,7 +55,8 @@ def load(proposalNB=None, runNB=None, laser_bp=None, ): """ - Load a run and extract the data. Output is an xarray with aligned + Load a run and extract the data. Output is an xarray with align + trainIds Parameters @@ -79,9 +80,9 @@ def load(proposalNB=None, runNB=None, whether to show the run.info or not validate: bool whether to run extra-data-validate or not - subset: - a subset of train that can be load with by_index[:5] for the first 5 - trains + subset: slice or extra_data.by_index or numpy.s_ + a subset of train that can be loaded with extra_data.by_index[:5] for + the first 5 trains. If None, all trains are retrieved. rois: dict a dictionnary of mnemonics with a list of rois definition and the desired names, for example: @@ -136,7 +137,9 @@ def load(proposalNB=None, runNB=None, """ runFolder = find_run_path(proposalNB, runNB, subFolder) - run = ed.RunDirectory(runFolder).select_trains(subset) + run = ed.RunDirectory(runFolder) + if subset is not None: + run = run.select_trains(subset) if fields is None: return run, xr.Dataset() if isinstance(fields, str): @@ -301,7 +304,7 @@ def find_run_path(proposalNB, runNB, data='raw'): return os.path.join(find_proposal(proposalNB), data, runNB) -def open_run(proposalNB, runNB, subset=ed.by_index[:], **kwargs): +def open_run(proposalNB, runNB, subset=None, **kwargs): """ Get extra_data.DataCollection in a given proposal. Wraps the extra_data open_run routine and adds subset selection, out of @@ -314,9 +317,9 @@ def open_run(proposalNB, runNB, subset=ed.by_index[:], **kwargs): proposal number e.g. 'p002252' or 2252 runNB: (str, int) run number e.g. 17 or 'r0017' - subset: - a subset of train that can be load with by_index[:5] for the first 5 - trains + subset: slice or extra_data.by_index or numpy.s_ + a subset of train that can be loaded with extra_data.by_index[:5] for + the first 5 trains. If None, all trains are retrieved. **kwargs -------- @@ -331,11 +334,14 @@ def open_run(proposalNB, runNB, subset=ed.by_index[:], **kwargs): DataCollection object containing information about the specified run. Data can be loaded using built-in class methods. """ - return ed.open_run(proposalNB, runNB, **kwargs).select_trains(subset) + run = ed.open_run(proposalNB, runNB, **kwargs) + if subset is not None: + run = run.select_trains(subset) + return run def get_array(run=None, mnemonic=None, stepsize=None, - subset=ed.by_index[:], subFolder='raw', + subset=None, subFolder='raw', proposalNB=None, runNB=None): """ Loads one data array for the specified mnemonic and rounds its values to @@ -354,9 +360,9 @@ def get_array(run=None, mnemonic=None, stepsize=None, stepsize : float nominal stepsize of the array data - values will be rounded to integer multiples of this value. - subset: - a subset of train that can be load with by_index[:5] for the first 5 - trains + subset: slice or extra_data.by_index or numpy.s_ + a subset of train that can be loaded with extra_data.by_index[:5] for + the first 5 trains. If None, all trains are retrieved. subFolder: (str) sub-folder from which to load the data. Use 'raw' for raw data or 'proc' for processed data.