diff --git a/xgm.py b/xgm.py index f9d341ed6a362798e540432f5b8a6e9796c5904f..a8866963a2d47e10a5a02097bd8df5b24217a4a7 100644 --- a/xgm.py +++ b/xgm.py @@ -29,12 +29,13 @@ def cleanXGMdata(data, npulses=None, sase3First=True): ''' dropList = [] mergeList = [] - dedicated = False + load_sa1 = True if 'sase3' in data: if np.all(data['npulses_sase1'].where(data['npulses_sase3'] !=0, drop=True) == 0): - dedicated = True + load_sa1 = False print('Dedicated trains, skip loading SASE 1 data.') + npulses_sa1 = None #pulse-resolved signals from XGMs keys = ["XTD10_XGM", "XTD10_SA3", "XTD10_SA1", "XTD10_XGM_sigma", "XTD10_SA3_sigma", "XTD10_SA1_sigma", @@ -43,26 +44,37 @@ def cleanXGMdata(data, npulses=None, sase3First=True): if ("SCS_SA3" not in data and "SCS_XGM" in data): #no SASE-resolved arrays available + if not 'sase3' in data: + npulses_xgm = data['SCS_XGM'].where(data['SCS_XGM'], drop=True).shape[1] + npulses_sa1 = npulses_xgm - npulses + if npulses_sa1<=0: + load_sa1 = False sa3 = selectSASEinXGM(data, xgm='SCS_XGM', sase='sase3', npulses=npulses, - sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'}) + sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'}).rename("SCS_SA3") mergeList.append(sa3) - if not dedicated: + if load_sa1: sa1 = selectSASEinXGM(data, xgm='SCS_XGM', sase='sase1', - npulses=npulses, sase3First=sase3First).rename( - {'XGMbunchId':'sa1_pId'}) + npulses=npulses_sa1, sase3First=sase3First).rename( + {'XGMbunchId':'sa1_pId'}).rename("SCS_SA1") mergeList.append(sa1) dropList.append('SCS_XGM') keys.remove('SCS_XGM') + load_sa1 = True if ("XTD10_SA3" not in data and "XTD10_XGM" in data): #no SASE-resolved arrays available + if not 'sase3' in data: + npulses_xgm = data['XTD10_XGM'].where(data['XTD10_XGM'], drop=True).shape[1] + npulses_sa1 = npulses_xgm - npulses + if npulses_sa1<=0: + load_sa1 = False sa3 = selectSASEinXGM(data, xgm='XTD10_XGM', sase='sase3', npulses=npulses, - sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'}) + sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'}).rename("XTD10_SA3") mergeList.append(sa3) - if not dedicated: + if load_sa1: sa1 = selectSASEinXGM(data, xgm='XTD10_XGM', sase='sase1', - npulses=npulses, sase3First=sase3First).rename( - {'XGMbunchId':'sa1_pId'}) + npulses=npulses_sa1, sase3First=sase3First).rename( + {'XGMbunchId':'sa1_pId'}).rename("XTD10_SA1") mergeList.append(sa1) dropList.append('XTD10_XGM') keys.remove('XTD10_XGM') @@ -74,7 +86,7 @@ def cleanXGMdata(data, npulses=None, sase3First=True): sase = 'sa3_' elif "sa1" in key.lower(): sase = 'sa1_' - if dedicated: + if not load_sa1: dropList.append(key) continue else: