diff --git a/Load.py b/Load.py index 62ec295c9db1d8ad071b14a16634d352809c0fec..a26a487b4a4b7ff7f62f6cf9ae568ecfd0b91890 100644 --- a/Load.py +++ b/Load.py @@ -6,8 +6,10 @@ Copyright (2019) SCS Team. """ import numpy as np -from karabo_data import RunDirectory, by_index +from karabo_data import by_index, RunDirectory +from karabo_data.read_machinery import find_proposal import xarray as xr +import os mnemonics = { # Machine @@ -362,23 +364,20 @@ mnemonics = { 'dim': ['gott_pId','pixelId']} } -def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, - validate=False, runpath='/gpfs/exfel/exp/{}/{}/{}/raw/r{:04d}/', - subset=by_index[:], rois={}): +def load(fields, runNB, proposalNB, subFolder='raw', display=False, validate=False, + subset=by_index[:], rois={}): """ Load a run and extract the data. Output is an xarray with aligned trainIds Inputs: fields: list of mnemonic strings to load specific data such as "fastccd", "SCS_XGM", or dictionnaries defining a custom mnemonic such as {"extra": {'SCS_CDIFFT_MAG/SUPPLY/CURRENT', 'actual_current.value', None}} - runNB: run number as integer - proposalNB: string or int of the proposal number e.g. 'p002252' or 2252 - semesterNB: string or int of the semester number where the proposal data are saved - topic: string of the topic - display: boolean, whether to show the run.info or not - 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 + runNB: (str, int) run number as integer + proposalNB: (str, int) of the proposal number e.g. 'p002252' or 2252 + subFolder: (str) sub-folder from which to load the data. Use 'raw' for raw + data or 'proc' for processed data. + display: (bool) whether to show the run.info or not + validate: (bool) whether to run karabo-data-validate or not subset: a subset of train that can be load with by_index[:5] for the first 5 trains rois: a dictionnary of mnemonics with a list of rois definition and the desired @@ -389,15 +388,18 @@ def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, Outputs: res: an xarray DataSet with aligned trainIds """ - - if type(proposalNB)==int: + + if isinstance(runNB, int): + runNB = 'r{:04d}'.format(runNB) + if isinstance(proposalNB,int): proposalNB = 'p{:06d}'.format(proposalNB) - runFolder = runpath.format(topic, semesterNB, proposalNB, runNB) + runFolder = os.path.join(find_proposal(proposalNB), subFolder, runNB) run = RunDirectory(runFolder).select_trains(subset) if validate: get_ipython().system('karabo-data-validate ' + runFolder) if display: + print('Loading data from {}'.format(runFolder)) run.info() keys = []