Skip to content
Snippets Groups Projects
xgm.py 47.8 KiB
Newer Older
# -*- coding: utf-8 -*-
""" Toolbox for SCS.

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

    Copyright (2019) SCS Team.
"""
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr


def pulsePatternInfo(data, plot=False):
    ''' display general information on the pulse patterns operated by SASE1 and SASE3.
        This is useful to track changes of number of pulses or mode of operation of
        SASE1 and SASE3. It also determines which SASE comes first in the train and
        the minimum separation between the two SASE sub-trains.
        
        Inputs:
            data: xarray Dataset containing pulse pattern info from the bunch decoder MDL: 
            {'sase1, sase3', 'npulses_sase1', 'npulses_sase3'}
            plot: bool enabling/disabling the plotting of the pulse patterns
            
        Outputs:
            print of pulse pattern info. If plot==True, plot of the pulse pattern.
    '''
    #Which SASE comes first?
    npulses_sa3 = data['npulses_sase3']       
    npulses_sa1 = data['npulses_sase1']  
    dedicated = False
    if np.all(npulses_sa1.where(npulses_sa3 !=0, drop=True) == 0):
        dedicated = True
        print('No SASE 1 pulses during SASE 3 operation')
    if np.all(npulses_sa3.where(npulses_sa1 !=0, drop=True) == 0):
        dedicated = True
        print('No SASE 3 pulses during SASE 1 operation')
    if dedicated==False:
        pulseIdmin_sa1 = data['sase1'].where(npulses_sa1 != 0).where(data['sase1']>1).min().values
        pulseIdmax_sa1 = data['sase1'].where(npulses_sa1 != 0).where(data['sase1']>1).max().values
        pulseIdmin_sa3 = data['sase3'].where(npulses_sa3 != 0).where(data['sase3']>1).min().values
        pulseIdmax_sa3 = data['sase3'].where(npulses_sa3 != 0).where(data['sase3']>1).max().values
        #print(pulseIdmin_sa1, pulseIdmax_sa1, pulseIdmin_sa3, pulseIdmax_sa3)
        if pulseIdmin_sa1 > pulseIdmax_sa3:
            t = 0.220*(pulseIdmin_sa1 - pulseIdmax_sa3 + 1)
            print('SASE 3 pulses come before SASE 1 pulses (minimum separation %.1f µs)'%t)
        elif pulseIdmin_sa3 > pulseIdmax_sa1:
            t = 0.220*(pulseIdmin_sa3 - pulseIdmax_sa1 + 1)
            print('SASE 1 pulses come before SASE 3 pulses (minimum separation %.1f µs)'%t)
        else:
            print('Interleaved mode')
    
    #What is the pulse pattern of each SASE?
    for key in['sase3', 'sase1']:
        print('\n*** %s pulse pattern: ***'%key.upper())
        npulses = data['npulses_%s'%key]
        sase = data[key]
        if not np.all(npulses == npulses[0]):
            print('Warning: number of pulses per train changed during the run!')
        #take the derivative along the trainId to track changes in pulse number:
        diff = npulses.diff(dim='trainId')
        #only keep trainIds where a change occured:
        diff = diff.where(diff !=0, drop=True)
        #get a list of indices where a change occured:
        idx_change = np.argwhere(np.isin(npulses.trainId.values,
                                         diff.trainId.values, assume_unique=True))[:,0]
        #add index 0 to get the initial pulse number per train:
        idx_change = np.insert(idx_change, 0, 0)
        print('npulses\tindex From\tindex To\ttrainId From\ttrainId To\trep. rate [kHz]')
        for i,idx in enumerate(idx_change):
            n = npulses[idx]
            idxFrom = idx
            trainIdFrom = npulses.trainId[idx]
            if i < len(idx_change)-1:
                idxTo = idx_change[i+1]-1
            else:
                idxTo = npulses.shape[0]-1
            trainIdTo = npulses.trainId[idxTo]
            if n <= 1:
                print('%i\t%i\t\t%i\t\t%i\t%i'%(n, idxFrom, idxTo, trainIdFrom, trainIdTo))
            else:
                f = 1/((sase[idxFrom,1] - sase[idxFrom,0])*222e-6)
                print('%i\t%i\t\t%i\t\t%i\t%i\t%.0f'%(n, idxFrom, idxTo, trainIdFrom, trainIdTo, f))
    print('\n')
    if plot:
        plt.figure(figsize=(6,3))
Mercadier's avatar
Mercadier committed
        plt.plot(data['npulses_sase3'].trainId, data['npulses_sase3'], 'o-', 
                 ms=3, label='SASE 3')
        plt.xlabel('trainId')
        plt.ylabel('pulses per train')
Mercadier's avatar
Mercadier committed
        plt.plot(data['npulses_sase1'].trainId, data['npulses_sase1'], '^-',
                 ms=3, color='C2', label='SASE 1')
Mercadier's avatar
Mercadier committed

def repRate(data, sase='sase3'):
    ''' Calculates the pulse repetition rate in sase according
        to the bunch pattern and assuming a minimum pulse 
Mercadier's avatar
Mercadier committed
        Inputs:
            data: xarray Dataset containing pulse pattern
            sase: sase in which the repetition rate is
                  calculated (1,2 or 3)
        Output:
            f: repetition rate in kHz
    '''
    assert sase in data, 'key "{}" not found in data!'.format(sase)
    sase = data[sase].where(data['npulses_{}'.format(sase)]>1,
                            drop=True).values
    if len(sase)==0:
        print('Not enough pulses to extract repetition rate')
        return 0
    f = 1/((sase[0,1] - sase[0,0])*221.54e-6)
def cleanXGMdata(data, npulses=None, sase3First=True):
    ''' Cleans the XGM data arrays obtained from load() function.
        The XGM "TD" data arrays have arbitrary size of 1000 and default value 1.0
        when there is no pulse. This function sorts the SASE 1 and SASE 3 pulses.
        For DAQ runs after April 2019, sase-resolved arrays can be used. For older runs,
        the function selectSASEinXGM can be used to extract sase-resolved pulses.
        Inputs:
            data: xarray Dataset containing XGM TD arrays.
            npulses: number of pulses, needed if pulse pattern not available.
            sase3First: bool, needed if pulse pattern not available.
        
        Output:
            xarray Dataset containing sase- and pulse-resolved XGM data, with
                dimension names 'sa1_pId' and 'sa3_pId'                
    '''
    dropList = []
    mergeList = []
    if ("XTD10_SA3" not in data and "XTD10_XGM" in data) or (
        "SCS_SA3" not in data and "SCS_XGM" in data):
        #no SASE-resolved arrays available
        if 'SCS_XGM' in data:
            sa3 = selectSASEinXGM(data, xgm='SCS_XGM', sase='sase3', npulses=npulses,
                   sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'}).rename('SCS_SA3')
            mergeList.append(sa3)
            sa1 = selectSASEinXGM(data, xgm='SCS_XGM', sase='sase1', npulses=npulses,
                   sase3First=sase3First).rename({'XGMbunchId':'sa1_pId'}).rename('SCS_SA1')
            mergeList.append(sa1)
            dropList.append('SCS_XGM')

        if 'XTD10_XGM' in data:
            sa3 = selectSASEinXGM(data, xgm='XTD10_XGM', sase='sase3', npulses=npulses,
                       sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'}).rename('XTD10_SA3')
            mergeList.append(sa3)
            sa1 = selectSASEinXGM(data, xgm='XTD10_XGM', sase='sase1', npulses=npulses,
                       sase3First=sase3First).rename({'XGMbunchId':'sa1_pId'}).rename('XTD10_SA1')
            mergeList.append(sa1)
            dropList.append('XTD10_XGM')
        keys = []
        
    else:
        keys = ["XTD10_XGM", "XTD10_SA3", "XTD10_SA1",
                "XTD10_XGM_sigma", "XTD10_SA3_sigma", "XTD10_SA1_sigma"]
        keys += ["SCS_XGM", "SCS_SA3", "SCS_SA1",
                 "SCS_XGM_sigma", "SCS_SA3_sigma", "SCS_SA1_sigma"]
        
    for key in keys:
        if key not in data:
            continue
        if "sa3" in key.lower():
            sase = 'sa3_'
        elif "sa1" in key.lower():
            sase = 'sa1_'
        else:
            dropList.append(key)
            continue
        res = data[key].where(data[key] != 1.0, drop=True).rename(
                {'XGMbunchId':'{}pId'.format(sase)}).rename(key)
        dropList.append(key)
        mergeList.append(res)
    mergeList.append(data.drop(dropList))
    subset = xr.merge(mergeList, join='outer')
    for k in data.attrs.keys():
        subset.attrs[k] = data.attrs[k]
def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=None):
    ''' Extract SASE1- or SASE3-only XGM data.
        There are various cases depending on i) the mode of operation (10 Hz
        with fresh bunch, dedicated trains to one SASE, pulse on demand), 
        ii) the potential change of number of pulses per train in each SASE
        and iii) the order (SASE1 first, SASE3 first, interleaved mode).
        
        Inputs:
            data: xarray Dataset containing xgm data
            sase: key of sase to select: {'sase1', 'sase3'}
            xgm: key of xgm to select: {'SA3_XGM', 'SCS_XGM'}
            sase3First: bool, optional. Used in case no bunch pattern was recorded
            npulses: int, optional. Required in case no bunch pattern was recorded.
            
        Output:
            DataArray that has all trainIds that contain a lasing
Loading
Loading full blame...