Load.py

# -*- coding: utf-8 -*-
""" Toolbox for SCS.

    Various utilities function to quickly process data measured at the SCS instruments.

    Copyright (2019) SCS Team.
"""
import numpy as np
from extra_data import by_index, RunDirectory
from extra_data.read_machinery import find_proposal
import xarray as xr
import os
from ToolBox.bunch_pattern import extractBunchPattern

mnemonics = {
    # Machine
    "sase3": {'source':'SCS_RR_UTC/MDL/BUNCH_DECODER',
              'key':'sase3.pulseIds.value',
              'dim':['bunchId']},
    "sase2": {'source':'SCS_RR_UTC/MDL/BUNCH_DECODER',
              'key':'sase2.pulseIds.value',
              'dim':['bunchId']},
    "sase1": {'source':'SCS_RR_UTC/MDL/BUNCH_DECODER',
              'key':'sase1.pulseIds.value',
              'dim':['bunchId']},
    "maindump": {'source':'SCS_RR_UTC/MDL/BUNCH_DECODER',
                 'key':'maindump.pulseIds.value',
                 'dim':['bunchId']},
    "bunchpattern": {'source':'SCS_RR_UTC/TSYS/TIMESERVER',
                     'key':'readBunchPatternTable.value',
                     'dim':None},
    "bunchPatternTable": {'source':'SCS_RR_UTC/TSYS/TIMESERVER',
                     'key':'bunchPatternTable.value',
                     'dim':['pulse_slot']},
    "npulses_sase3": {'source':'SCS_RR_UTC/MDL/BUNCH_DECODER',
                      'key':'sase3.nPulses.value',
                      'dim':None},
    "npulses_sase1": {'source':'SCS_RR_UTC/MDL/BUNCH_DECODER',
                      'key':'sase1.nPulses.value',
                      'dim':None},

    #Bunch Arrival Monitors
    "BAM5": {'source':'SCS_ILH_LAS/DOOCS/BAM_414_B2:output',
                      'key':'data.lowChargeArrivalTime',
                      'dim':['BAMbunchId']},
    "BAM6": {'source':'SCS_ILH_LAS/DOOCS/BAM_1932M_TL:output',
                      'key':'data.lowChargeArrivalTime',
                      'dim':['BAMbunchId']},
    "BAM7": {'source':'SCS_ILH_LAS/DOOCS/BAM_1932S_TL:output',
                      'key':'data.lowChargeArrivalTime',
                      'dim':['BAMbunchId']},
    
    # SA3
    "nrj": {'source':'SA3_XTD10_MONO/MDL/PHOTON_ENERGY',
            'key':'actualEnergy.value',
            'dim':None},

    "M2BEND": {'source': 'SA3_XTD10_MIRR-2/MOTOR/BENDER',
               'key': 'actualPosition.value',
               'dim':None},
    "VSLIT": {'source':'SA3_XTD10_VSLIT/MDL/BLADE',
              'key':'actualGap.value',
              'dim':None},
    "ESLIT": {'source':'SCS_XTD10_ESLIT/MDL/MAIN',
              'key':'actualGap.value',
              'dim':None},
    "HSLIT": {'source':'SCS_XTD10_HSLIT/MDL/BLADE',
              'key':'actualGap.value',
              'dim':None},
    "transmission": {'source':'SA3_XTD10_GATT/MDL/GATT_TRANSMISSION_MONITOR',
                     'key':'Estimated_Tr.value',
                     'dim':None},
    "GATT_pressure": {'source':'P_GATT',
                      'key':'value.value',
                      'dim':None},
    "navitar": {'source':'SCS_XTD10_IMGES/CAM/BEAMVIEW_NAVITAR:daqOutput',
               'key':'data.image.pixels',
               'dim':['x','y']},
    "UND": {'source':'SA3_XTD10_UND/DOOCS/PHOTON_ENERGY',
                'key':'actualPosition.value',
                'dim':None},
    #DPS imagers
    "DPS2CAM2": {'source':'SCS_BLU_DPS-2/CAM/IMAGER2CAMERA:daqOutput',
                'key':'data.image.pixels',
                'dim':['dps2cam2_y', 'dps2cam2_x']},
    # XTD10 XGM
    ## keithley
    "XTD10_photonFlux": {'source':'SA3_XTD10_XGM/XGM/DOOCS',
                     'key':'pulseEnergy.photonFlux.value',
                     'dim':None},
    "XTD10_photonFlux_sigma": {'source':'SA3_XTD10_XGM/XGM/DOOCS',
                     'key':'pulseEnergy.photonFluxSigma.value',
                     'dim':None},
    ## ADC
    "XTD10_XGM": {'source':'SA3_XTD10_XGM/XGM/DOOCS:output',
                'key':'data.intensityTD',
                'dim':['XGMbunchId']},
    "XTD10_XGM_sigma": {'source':'SA3_XTD10_XGM/XGM/DOOCS:output',
                'key':'data.intensitySigmaTD',
                'dim':['XGMbunchId']},
    "XTD10_SA3": {'source':'SA3_XTD10_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa3TD',
                'dim':['XGMbunchId']},
    "XTD10_SA3_sigma": {'source':'SA3_XTD10_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa3SigmaTD',
                'dim':['XGMbunchId']},
    "XTD10_SA1": {'source':'SA3_XTD10_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa1TD',
                'dim':['XGMbunchId']},
    "XTD10_SA1_sigma": {'source':'SA3_XTD10_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa1SigmaTD',
                'dim':['XGMbunchId']},
    ## low pass averaged ADC
    "XTD10_slowTrain": {'source':'SA3_XTD10_XGM/XGM/DOOCS',
                    'key':'controlData.slowTrain.value',
                    'dim':None},
    "XTD10_slowTrain_SA1": {'source':'SA3_XTD10_XGM/XGM/DOOCS',
                    'key':'controlData.slowTrainSa1.value',
                    'dim':None},
    "XTD10_slowTrain_SA3": {'source':'SA3_XTD10_XGM/XGM/DOOCS',
                    'key':'controlData.slowTrainSa3.value',
                    'dim':None},

    # SCS XGM
    ## keithley
    "SCS_photonFlux": {'source':'SCS_BLU_XGM/XGM/DOOCS',
                     'key':'pulseEnergy.photonFlux.value',
                     'dim':None},
    "SCS_photonFlux_sigma": {'source':'SCS_BLU_XGM/XGM/DOOCS',
                     'key':'pulseEnergy.photonFluxSigma.value',
                     'dim':None},
    ## ADC
    "SCS_XGM": {'source':'SCS_BLU_XGM/XGM/DOOCS:output',
                'key':'data.intensityTD',
                'dim':['XGMbunchId']},
    "SCS_XGM_sigma": {'source':'SCS_BLU_XGM/XGM/DOOCS:output',
                'key':'data.intensitySigmaTD',
                'dim':['XGMbunchId']},
    "SCS_SA1": {'source':'SCS_BLU_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa1TD',
                'dim':['XGMbunchId']},
    "SCS_SA1_sigma": {'source':'SCS_BLU_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa1SigmaTD',
                'dim':['XGMbunchId']},
    "SCS_SA3": {'source':'SCS_BLU_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa3TD',
                'dim':['XGMbunchId']},
    "SCS_SA3_sigma": {'source':'SCS_BLU_XGM/XGM/DOOCS:output',
                'key':'data.intensitySa3SigmaTD',
                'dim':['XGMbunchId']},
    ## low pass averaged ADC
    "SCS_slowTrain": {'source':'SCS_BLU_XGM/XGM/DOOCS',
                    'key':'controlData.slowTrain.value',
                    'dim':None},
    "SCS_slowTrain_SA1": {'source':'SCS_BLU_XGM/XGM/DOOCS',
                    'key':'controlData.slowTrainSa1.value',
                    'dim':None},
    "SCS_slowTrain_SA3": {'source':'SCS_BLU_XGM/XGM/DOOCS',
                    'key':'controlData.slowTrainSa3.value',
                    'dim':None},

    # KBS
    "HFM_CAPB": {'source':'SCS_KBS_HFM/ASENS/CAPB',
                 'key':'value.value',
                 'dim':None},
    "HFM_CAPF": {'source':'SCS_KBS_HFM/ASENS/CAPF',
                 'key':'value.value',
                 'dim':None},
    "HFM_CAPM": {'source':'SCS_KBS_HFM/ASENS/CAPM',
                 'key':'value.value',
                 'dim':None},
    "HFM_BENDERB": {'source':'SCS_KBS_HFM/MOTOR/BENDERB',
                    'key':'encoderPosition.value',
                    'dim':None},
    "HFM_BENDERF": {'source':'SCS_KBS_HFM/MOTOR/BENDERF',
                    'key':'encoderPosition.value',
                    'dim':None},
    "VFM_CAPB": {'source':'SCS_KBS_VFM/ASENS/CAPB',
                 'key':'value.value',
                 'dim':None},
    "VFM_CAPF": {'source':'SCS_KBS_VFM/ASENS/CAPF',
                 'key':'value.value',
                 'dim':None},
    "VFM_CAPM": {'source':'SCS_KBS_VFM/ASENS/CAPM',
                 'key':'value.value',
                 'dim':None},
    "VFM_BENDERB": {'source':'SCS_KBS_VFM/MOTOR/BENDERB',
                    'key':'encoderPosition.value',
                    'dim':None},
    "VFM_BENDERF": {'source':'SCS_KBS_VFM/MOTOR/BENDERF',
                    'key':'encoderPosition.value',
                    'dim':None},
    
    # AFS LASER
    "AFS_PhaseShifter": {'source':'SCS_ILH_LAS/PHASESHIFTER/DOOCS',
                 'key':'actualPosition.value',
                 'dim':None},
    "AFS_DelayLine": {'source':'SCS_ILH_LAS/MOTOR/LT3',
                 'key':'AActualPosition.value',
                 'dim':None},
    "AFS_HalfWP": {'source':'SCS_ILH_LAS/MOTOR/ROT_OPA_BWP1',
                 'key':'actualPosition.value',
                 'dim':None},
    "AFS_FocusLens": {'source':'SCS_ILH_LAS/MOTOR/LT_SPARE1',
                 'key':'actualPosition.value',
                 'dim':None},
    # 2nd lens of telescope
    "AFS_TeleLens": {'source':'SCS_ILH_LAS/MOTOR/LT2',
                 'key':'actualPosition.value',
                 'dim':None},

    # PP LASER 800 nm path
    "PP800_PhaseShifter": {'source':'SCS_ILH_LAS/DOOCS/PP800_PHASESHIFTER',
                 'key':'actualPosition.value',
                 'dim':None},
    "PP800_SynchDelayLine": {'source':'SCS_ILH_LAS/DOOCS/PPL_OPT_DELAY',
                 'key':'actualPosition.value',
                 'dim':None},
    "PP800_DelayLine": {'source':'SCS_ILH_LAS/MOTOR/LT3',
                 'key':'AActualPosition.value',
                 'dim':None},
    "PP800_HalfWP": {'source':'SCS_ILH_LAS/MOTOR/ROT8WP1',
                 'key':'actualPosition.value',
                 'dim':None},
    "PP800_FocusLens": {'source':'SCS_ILH_LAS/MOTOR/LT_SPARE1',
                 'key':'actualPosition.value',
                 'dim':None},
    # 1st lens of telescope (setup of August 2019)
    "PP800_TeleLens": {'source':'SCS_ILH_LAS/MOTOR/LT7',
                 'key':'actualPosition.value',
                 'dim':None},
    "ILH_8CAM1": {'source':'SCS_ILH_LAS/CAM/8CAM1:daqOutput',
                'key':'data.image.pixels',
                'dim':['8cam1_y', '8cam1_x']},

    # GPC 
    "GPC_EOS_DelayLine": {'source':'SCS_CDIDET_GRD/MOTOR/IMAGER',
                 'key':'actualPosition.value',
                 'dim':None},
    "GPC_X": {'source':'SCS_GPC_MOV/MOTOR/X',
                 'key':'actualPosition.value',
                 'dim':None},
    "GPC_Y": {'source':'SCS_GPC_MOV/MOTOR/Y',
                 'key':'actualPosition.value',
                 'dim':None},
    "GPC_THETA": {'source':'SCS_GPC_MOV/MOTOR/THETA',
                 'key':'actualPosition.value',
                 'dim':None},
    "GPC_THETAMAG": {'source':'SCS_GPC_MOV/MOTOR/THETAMAG',
                 'key':'actualPosition.value',
                 'dim':None},
    
    # FFT
    "scannerX": {'source':'SCS_CDIFFT_SAM/LMOTOR/SCANNERX',
                 'key':'actualPosition.value',
                 'dim':None},
    "scannerY": {'source':'SCS_CDIFFT_SAM/MOTOR/SCANNERY',
                 'key':'actualPosition.value',
                 'dim':None},
    "scannerY_enc": {'source':'SCS_CDIFFT_SAM/ENC/SCANNERY',
                     'key':'value.value',
                     'dim':None},
    "SAM-Z": {'source':'SCS_CDIFFT_MOV/ENC/SAM_Z',
              'key':'value.value',
              'dim':None},
    "magnet": {'source':'SCS_CDIFFT_MAG/ASENS/CURRENT',
               'key':'value.value',
               'dim':None},
    "magnet_old": {'source':'SCS_CDIFFT_MAG/SUPPLY/CURRENT',
               'key':'actualCurrent.value',
               'dim':None},
    
    "Vertical_FDM": {'source':'SCS_CDIFFT_LDM/CAM/CAMERA1A:daqOutput',
                'key':'data.image.pixels',
                'dim':['vfdm_y', 'vfdm_x']},

    # FastCCD, if in raw folder, raw images
    #          if in proc folder, dark substracted and relative gain corrected
    "fastccd": {'source':'SCS_CDIDET_FCCD2M/DAQ/FCCD:daqOutput',
                'key':'data.image.pixels',
                'dim':['x', 'y']},
    # FastCCD with common mode correction
    "fastccd_cm": {'source':'SCS_CDIDET_FCCD2M/DAQ/FCCD:daqOutput',
                'key':'data.image.pixels_cm',
                'dim':['x', 'y']},
    # FastCCD charge split correction in very low photon count regime
    "fastccd_classified": {'source':'SCS_CDIDET_FCCD2M/DAQ/FCCD:daqOutput',
                'key':'data.image.pixels_classified',
                'dim':['x', 'y']},
    # FastCCD event multiplicity from the charge split correction:
    # 0: no events
    # 100, 101: single events
    # 200-203: charge split into two pixels in four different orientations
    # 300-303: charge split into three pixels in four different orientations
    # 400-403: charge split into four pixels in four different orientations
    # 1000: charge in more than four neighboring pixels. Cannot be produced by a single photon alone.
    "fastccd_patterns": {'source':'SCS_CDIDET_FCCD2M/DAQ/FCCD:daqOutput',
                'key':'data.image.patterns',
                'dim':['x', 'y']},
    # FastCCD gain map, 0 high gain, 1 medium gain, 2 low gain
    "fastccd_gain": {'source':'SCS_CDIDET_FCCD2M/DAQ/FCCD:daqOutput',
                'key':'data.image.gain',
                'dim':['x', 'y']},
    # FastCCD mask, bad pixel map to be ignored if > 0
    "fastccd_mask": {'source':'SCS_CDIDET_FCCD2M/DAQ/FCCD:daqOutput',
                'key':'data.image.mask',
                'dim':['x', 'y']},

    # TIM
    "MCP1apd": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_D.apd.pulseIntegral',
                'dim':['apdId']},
    "MCP1raw": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_D.raw.samples',
                'dim':['samplesId']},
    "MCP2apd": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_C.apd.pulseIntegral',
                'dim':['apdId']},
    "MCP2raw": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_C.raw.samples',
                'dim':['samplesId']},
    "MCP3apd": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_B.apd.pulseIntegral',
                'dim':['apdId']},
    "MCP3raw": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_B.raw.samples',
                'dim':['samplesId']},
    "MCP4apd": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_A.apd.pulseIntegral',
                'dim':['apdId']},
    "MCP4raw": {'source':'SCS_UTC1_ADQ/ADC/1:network',
                'key':'digitizers.channel_1_A.raw.samples',
                'dim': ['samplesId']},

    # FastADC
    "FastADC0peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_0.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC0raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_0.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC1peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_1.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC1raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_1.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC2peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_2.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC2raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_2.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC3peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_3.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC3raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_3.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC4peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_4.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC4raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_4.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC5peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_5.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC5raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_5.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC6peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_6.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC6raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_6.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC7peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_7.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC7raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_7.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC8peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_8.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC8raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_8.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},
    "FastADC9peaks": {'source':'SCS_UTC1_MCP/ADC/1:channel_9.output',
                'key':'data.peaks',
                'dim':['peakId']},
    "FastADC9raw": {'source':'SCS_UTC1_MCP/ADC/1:channel_9.output',
                'key':'data.rawData',
                'dim':['fadc_samplesId']},

    # KARABACON
    "KARABACON": {'source':'SCS_DAQ_SCAN/MDL/KARABACON',
                    'key': 'actualStep.value',
                    'dim': None},
    
    #GOTTHARD
    "Gotthard1": {'source':'SCS_PAM_XOX/DET/GOTTHARD_RECEIVER1:daqOutput',
                    'key': 'data.adc',
                    'dim': ['gott_pId','pixelId']},
    "Gotthard2": {'source':'SCS_PAM_XOX/DET/GOTTHARD_RECEIVER2:daqOutput',
                    'key': 'data.adc',
                    'dim': ['gott_pId','pixelId']}
}

def load(fields, runNB, proposalNB, subFolder='raw', display=False, validate=False,
         subset=by_index[:], rois={}, useBPTable=True):
    """ 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: (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 extra-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
                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']}}}
            useBPTable: If True, uses the raw bunch pattern table to extract sase pulse
                number and indices in the trains. If false, load the data from BUNCH_DECODER
                middle layer device.

        Outputs:
            res: an xarray DataSet with aligned trainIds
    """
    
    if isinstance(runNB, int):
        runNB = 'r{:04d}'.format(runNB)
    if isinstance(proposalNB,int):
        proposalNB = 'p{:06d}'.format(proposalNB)
    runFolder = os.path.join(find_proposal(proposalNB), subFolder, runNB)
    run = RunDirectory(runFolder).select_trains(subset)

    if validate:
        get_ipython().system('extra-data-validate ' + runFolder)
    if display:
        print('Loading data from {}'.format(runFolder))
        run.info()

    keys = []
    vals = []

    # load pulse pattern info
    if useBPTable:
        bp_mnemo = mnemonics['bunchPatternTable']
        if bp_mnemo['source'] not in run.all_sources:
            print('Source {} not found in run. Skipping!'.format(
                                mnemonics['bunchPatternTable']['source']))
        else:
            bp_table = run.get_array(bp_mnemo['source'],bp_mnemo['key'], 
                                        extra_dims=bp_mnemo['dim'])
            sase1, npulses_sase1, dummy = extractBunchPattern(bp_table, 'sase1')
            sase3, npulses_sase3, dummy = extractBunchPattern(bp_table, 'sase3')
            keys += ["sase1", "npulses_sase1", "sase3", "npulses_sase3"]
            vals += [sase1, npulses_sase1, sase3, npulses_sase3]
    else:
        fields += ["sase1", "sase3", "npulses_sase3", "npulses_sase1"]

    for f in fields:

        if type(f) == dict:
            # extracting mnemomic defined on the spot
            if len(f.keys()) > 1:
                print('Loading only one "on-the-spot" mnemonic at a time, skipping all others !')
            k = list(f.keys())[0]
            v = f[k]
        else:
            # extracting mnemomic from the table
            if f in mnemonics:
                v = mnemonics[f]
                k = f
            else:
                print('Unknow mnemonic "{}". Skipping!'.format(f))
                continue
                
        if k in keys:
            continue # already loaded, skip

        if display:
            print('Loading {}'.format(k))

        if v['source'] not in run.all_sources:
            print('Source {} not found in run. Skipping!'.format(v['source']))
            continue

        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)
    result.attrs['run'] = run
    result.attrs['runFolder'] = runFolder
    return result

def concatenateRuns(runs):
    """ Sorts and concatenate a list of runs with identical data variables along the
        trainId dimension.
        
        Input:
            runs: (list) the xarray Datasets to concatenate
        Output:
            a concatenated xarray Dataset
    """
    firstTid = {i: int(run.trainId[0].values) for i,run in enumerate(runs)}
    orderedDict = dict(sorted(firstTid.items(), key=lambda t: t[1]))
    orderedRuns = [runs[i] for i in orderedDict]
    keys = orderedRuns[0].keys()
    for run in orderedRuns[1:]:
        if run.keys() != keys:
            print('data fields between different runs are not identical. Cannot combine runs.')
            return
    
    result = xr.concat(orderedRuns, dim='trainId')
    for k in orderedRuns[0].attrs.keys():
        result.attrs[k] = [run.attrs[k] for run in orderedRuns]
    return result