diff --git a/Load.py b/Load.py index d43dc0f882c5744350c05ed362fb452d927fa734..2ced1ef0b0c56c45c93f342e59392d3fda3b252b 100644 --- a/Load.py +++ b/Load.py @@ -172,7 +172,7 @@ mnemonics = { def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, validate=False, runpath='/gpfs/exfel/exp/{}/{}/{}/raw/r{:04d}/', - subset=by_index[:]): + subset=by_index[:], rois={}): """ Load a run and extract the data. Output is an xarray with aligned trainIds Inputs: @@ -188,7 +188,11 @@ def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, 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. + first 5 trains + rois: a dictionnary of mnemonics with a list of rois definition and the desired + names, for example {'fastccd':{'ref':{'roi':by_index[730:890, 535:720], + 'dim': ['ref_x', 'ref_y']}, 'sam':{'roi':by_index[1050:1210, 535:720], + 'dim': ['sam_x', 'sam_y']}}} Outputs: res: an xarray DataSet with aligned trainIds @@ -234,10 +238,16 @@ def load(fields, runNB, proposalNB, semesterNB, topic='SCS', display=False, print('Source {} not found in run. Skipping!'.format(v['source'])) continue - vals.append(run.get_array(v['source'], v['key'], extra_dims=v['dim'])) - - keys.append(k) - + if k not in rois: + # no ROIs selection, we read everything + vals.append(run.get_array(v['source'], v['key'], extra_dims=v['dim'])) + keys.append(k) + else: + # ROIs selection, for each ROI we select a region of the data and save it with new name and dimensions + for nk,nv in rois[k].items(): + vals.append(run.get_array(v['source'], v['key'], extra_dims=nv['dim'], roi=nv['roi'])) + keys.append(nk) + aligned_vals = xr.align(*vals, join='inner') result = dict(zip(keys, aligned_vals)) result = xr.Dataset(result)