Skip to content
Snippets Groups Projects
Commit 4c1cf2df authored by Laurent Mercadier's avatar Laurent Mercadier
Browse files

Fixes some bugs

parent 971893a8
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 89 additions and 36 deletions
xgm.py
+ 89
− 36
View file @ 4c1cf2df
......@@ -133,33 +133,42 @@ def cleanXGMdata(data, npulses=None, sase3First=True):
'''
dropList = []
mergeList = []
if ("XTD10_SA3" not in data and "XTD10_XGM" in data) or (
"SCS_SA3" not in data and "SCS_XGM" in data):
dedicated = False
if 'sase3' in data:
if np.all(data['npulses_sase1'].where(data['npulses_sase3'] !=0,
drop=True) == 0):
dedicated = True
print('Dedicated trains, skip loading SASE 1 data.')
keys = ["XTD10_XGM", "XTD10_SA3", "XTD10_SA1",
"XTD10_XGM_sigma", "XTD10_SA3_sigma", "XTD10_SA1_sigma",
"SCS_XGM", "SCS_SA3", "SCS_SA1",
"SCS_XGM_sigma", "SCS_SA3_sigma", "SCS_SA1_sigma"]
if ("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')
sa3 = selectSASEinXGM(data, xgm='SCS_XGM', sase='sase3', npulses=npulses,
sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'})
mergeList.append(sa3)
if not dedicated:
sa1 = selectSASEinXGM(data, xgm='SCS_XGM', sase='sase1',
npulses=npulses, sase3First=sase3First).rename(
{'XGMbunchId':'sa1_pId'})
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')
dropList.append('SCS_XGM')
keys.remove('SCS_XGM')
if ("XTD10_SA3" not in data and "XTD10_XGM" in data):
#no SASE-resolved arrays available
sa3 = selectSASEinXGM(data, xgm='XTD10_XGM', sase='sase3', npulses=npulses,
sase3First=sase3First).rename({'XGMbunchId':'sa3_pId'})
mergeList.append(sa3)
if not dedicated:
sa1 = selectSASEinXGM(data, xgm='XTD10_XGM', sase='sase1',
npulses=npulses, sase3First=sase3First).rename(
{'XGMbunchId':'sa1_pId'})
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"]
dropList.append('XTD10_XGM')
keys.remove('XTD10_XGM')
for key in keys:
if key not in data:
......@@ -168,6 +177,9 @@ def cleanXGMdata(data, npulses=None, sase3First=True):
sase = 'sa3_'
elif "sa1" in key.lower():
sase = 'sa1_'
if dedicated:
dropList.append(key)
continue
else:
dropList.append(key)
continue
......@@ -176,7 +188,7 @@ def cleanXGMdata(data, npulses=None, sase3First=True):
dropList.append(key)
mergeList.append(res)
mergeList.append(data.drop(dropList))
subset = xr.merge(mergeList, join='outer')
subset = xr.merge(mergeList, join='inner')
for k in data.attrs.keys():
subset.attrs[k] = data.attrs[k]
return subset
......@@ -237,13 +249,51 @@ def selectSASEinXGM(data, sase='sase3', xgm='SCS_XGM', sase3First=True, npulses=
dims=['trainId', 'bunchId'],
coords={'trainId':data.trainId},
name='sase_all')
mask_sa1 = sa_all.sel(trainId=sa1.trainId).isin(sa1_val).rename({'bunchId':'XGMbunchId'})
mask_sa3 = sa_all.sel(trainId=sa3.trainId).isin(sa3_val).rename({'bunchId':'XGMbunchId'})
idxListAll, invAll = np.unique(sa_all.fillna(-1), axis=0, return_inverse=True)
idxList3 = np.unique(sa3)
idxList1 = np.unique(sa1)
if sase=='sase3':
big_sa3 = []
for i,idxXGM in enumerate(idxListAll):
idxXGM = np.isin(idxXGM, idxList3)
idxTid = invAll==i
mask = xr.DataArray(np.zeros((data.dims['trainId'], sa_all['bunchId'].shape[0]), dtype=bool),
dims=['trainId', 'XGMbunchId'],
coords={'trainId':data.trainId,
'XGMbunchId':sa_all['bunchId'].values},
name='mask')
mask[idxTid, idxXGM] = True
sa3 = xgm_arr.where(mask, drop=True)
if sa3.trainId.size > 0:
sa3 = sa3.assign_coords(XGMbunchId=np.arange(sa3.XGMbunchId.size))
big_sa3.append(sa3)
if len(big_sa3) > 0:
da_sa3 = xr.concat(big_sa3, dim='trainId').rename('{}_SA3'.format(xgm.split('_')[0]))
else:
da_sa3 = xr.DataArray([], dims=['trainId'], name='{}_SA3'.format(xgm.split('_')[0]))
return da_sa3
if sase=='sase1':
return xgm_arr.where(mask_sa1, drop=True).rename('{}_SA1'.format(xgm.split('_')[0]))
else:
return xgm_arr.where(mask_sa3, drop=True).rename('{}_SA3'.format(xgm.split('_')[0]))
big_sa1 = []
for i,idxXGM in enumerate(idxListAll):
idxXGM = np.isin(idxXGM, idxList1)
idxTid = invAll==i
mask = xr.DataArray(np.zeros((data.dims['trainId'], sa_all['bunchId'].shape[0]), dtype=bool),
dims=['trainId', 'XGMbunchId'],
coords={'trainId':data.trainId,
'XGMbunchId':sa_all['bunchId'].values},
name='mask')
mask[idxTid, idxXGM] = True
sa1 = xgm_arr.where(mask, drop=True)
if sa1.trainId.size > 0:
sa1 = sa1.assign_coords(XGMbunchId=np.arange(sa1.XGMbunchId.size))
big_sa1.append(sa1)
if len(big_sa1) > 0:
da_sa1 = xr.concat(big_sa1, dim='trainId').rename('{}_SA1'.format(xgm.split('_')[0]))
else:
da_sa1 = xr.DataArray([], dims=['trainId'], name='{}_SA1'.format(xgm.split('_')[0]))
return da_sa1
def saseContribution(data, sase='sase1', xgm='XTD10_XGM'):
''' Calculate the relative contribution of SASE 1 or SASE 3 pulses
......@@ -539,14 +589,17 @@ def getTIMapd(data, mcp=1, use_apd=True, intstart=None, intstop=None,
#2.1 case where apd is used:
if use_apd:
return data[f'MCP{mcp}apd'].where(mask, drop=True)
res = data[f'MCP{mcp}apd'].where(mask, drop=True)
res = res.assign_coords(apdId=np.arange(res['apdId'].shape[0]))
return res
#2.2 case where integration is performed on raw trace:
else:
peaks = mcpPeaks(data, intstart, intstop, bkgstart, bkgstop, mcp=mcp,
t_offset=period, npulses=data.dims['apdId'])
peaks = mcpPeaks(data, intstart, intstop, bkgstart, bkgstop, mcp=mcp, t_offset=period,
npulses=data.dims['apdId'])
mask = mask.rename({'apdId':f'MCP{mcp}fromRaw'})
return peaks.where(mask, drop=True)
res = peaks.where(mask, drop=True)
res = res.assign_coords({f'MCP{mcp}fromRaw':np.arange(res[f'MCP{mcp}fromRaw'].shape[0])})
return res
def calibrateTIM(data, rollingWindow=200, mcp=1, plot=False, use_apd=True, intstart=None,
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

Legal Notice | Contacts