diff --git a/Load.py b/Load.py index c63a54a58f900bcd0d92870f6139478c6804cc65..a02cfb9a0b411dd8d112714fee449d6f8d96ec0e 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, open_run, RunDirectory +from karabo_data.read_machinery import find_proposal import xarray as xr +import os mnemonics = { # Machine @@ -362,7 +364,7 @@ mnemonics = { 'dim': ['gott_pId','pixelId']} } -def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, +def load(fields, runNB, proposalNB, semesterNB=None, topic='SCS', useRaw=True, display=False, validate=False, runpath='/gpfs/exfel/exp/{}/{}/{}/raw/r{:04d}/', subset=by_index[:], rois={}): """ Load a run and extract the data. Output is an xarray with aligned trainIds @@ -371,10 +373,13 @@ def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, 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 + runNB: (str, int) run number as integer + proposalNB: (str, int) of the proposal number e.g. 'p002252' or 2252 + semesterNB: (str, int) optional, depricated. The semester number where the proposal + data are saved. Depricated since karabo_data version 0.5 + topic: string of the topic. + useRaw: (bool) if True, access data from 'raw' folder of the run, otherwise + from 'proc' folder (for processed data) 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, @@ -389,10 +394,13 @@ 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) + dataSubFolder = 'raw' if useRaw else 'proc' + runFolder = os.path.join(find_proposal(proposalNB), dataSubFolder, runNB) run = RunDirectory(runFolder).select_trains(subset) if validate: