# -*- coding: utf-8 -*-
""" Toolbox for SCS.
Various utilities function to quickly process data
measured at the SCS instruments.
Copyright (2019) SCS Team.
"""
import os
import logging
import numpy as np
import xarray as xr
from extra_data.read_machinery import find_proposal
from extra_data import RunDirectory
# import and hide variable, such that it does not alter namespace.
from ..constants import mnemonics as _mnemonics_bp
from ..mnemonics_machinery import mnemonics_for_run
from .bunch_pattern_external import is_pulse_at
__all__ = [
'extractBunchPattern',
'get_sase_pId',
'npulses_has_changed',
'pulsePatternInfo',
'repRate',
]
log = logging.getLogger(__name__)
[docs]def npulses_has_changed(run, loc='sase3', run_mnemonics=None):
"""
Checks if the number of pulses has changed during the run for
a specific location `loc` (='sase1', 'sase3', 'scs_ppl' or 'laser')
If the source is not found in the run, returns True.
Parameters
----------
run: extra_data.DataCollection
DataCollection containing the data.
loc: str
The location where to check: {'sase1', 'sase3', 'scs_ppl'}
run_mnemonics: dict
the mnemonics for the run (see `menonics_for_run`)
Returns
-------
ret: bool
True if the number of pulses has changed or the source was not
found, False if the number of pulses did not change.
"""
sase_list = ['sase1', 'sase3', 'laser', 'scs_ppl']
if loc not in sase_list:
raise ValueError(f"Unknow sase location '{loc}'. Expected one in"
f"{sase_list}")
if run_mnemonics is None:
run_mnemonics = mnemonics_for_run(run)
if loc == 'scs_ppl':
loc = 'laser'
if loc not in run_mnemonics:
return True
if run_mnemonics[loc]['key'] not in run[run_mnemonics[loc]['source']].keys():
log.info(f'Mnemonic {loc} not found in run.')
return True
npulses = run.get_array(*run_mnemonics['npulses_'+loc].values())
if len(np.unique(npulses)) == 1:
return False
return True
def get_unique_sase_pId(run, loc='sase3', run_mnemonics=None):
"""
Assuming that the number of pulses did not change during the run,
returns the pulse Ids as the run value of the sase mnemonic.
Parameters
----------
run: extra_data.DataCollection
DataCollection containing the data.
loc: str
The location where to check: {'sase1', 'sase3', 'scs_ppl'}
run_mnemonics: dict
the mnemonics for the run (see `menonics_for_run`)
Returns
-------
pulseIds: np.array
the pulse ids at the specified location. Returns None if the
mnemonic is not in the run.
"""
if run_mnemonics is None:
run_mnemonics = mnemonics_for_run(run)
if loc == 'scs_ppl':
loc = 'laser'
if loc not in run_mnemonics:
# bunch pattern not recorded
return None
npulses = run.get_run_value(run_mnemonics['npulses_'+loc]['source'],
run_mnemonics['npulses_'+loc]['key'])
pulseIds = run.get_run_value(run_mnemonics[loc]['source'],
run_mnemonics[loc]['key'])[:npulses]
return pulseIds
[docs]def get_sase_pId(run, loc='sase3', run_mnemonics=None,
bpt=None, merge_with=None):
"""
Returns the pulse Ids of the specified `loc` during a run.
If the number of pulses has changed during the run, it loads the
bunch pattern table and extract all pulse Ids used.
Parameters
----------
run: extra_data.DataCollection
DataCollection containing the data.
loc: str
The location where to check: {'sase1', 'sase3', 'scs_ppl'}
run_mnemonics: dict
the mnemonics for the run (see `menonics_for_run`)
bpt: 2D-array
The bunch pattern table. Used only if the number of pulses
has changed. If None, it is loaded on the fly.
merge_with: xarray.Dataset
dataset that may contain the bunch pattern table to use in
case the number of pulses has changed. If merge_with does
not contain the bunch pattern table, it is loaded and added
as a variable 'bunchPatternTable' to merge_with.
Returns
-------
pulseIds: np.array
the pulse ids at the specified location. Returns None if the
mnemonic is not in the run.
"""
if npulses_has_changed(run, loc, run_mnemonics) is False:
return get_unique_sase_pId(run, loc, run_mnemonics)
if bpt is None:
bpt = load_bpt(run, merge_with, run_mnemonics)
if bpt is not None:
mask = is_pulse_at(bpt, loc)
return np.unique(np.nonzero(mask.values)[1])
return None
def load_bpt(run, merge_with=None, run_mnemonics=None):
"""
Load the bunch pattern table. It returns the one contained in
merge_with if possible. Or, it adds it to merge_with once it is
loaded.
Parameters
----------
run: extra_data.DataCollection
DataCollection containing the data.
merge_with: xarray.Dataset
dataset that may contain the bunch pattern table or to which
add the bunch pattern table once loaded.
run_mnemonics: dict
the mnemonics for the run (see `menonics_for_run`)
Returns
-------
bpt: xarray.Dataset
the bunch pattern table as specified by the mnemonics
'bunchPatternTable'
"""
if run_mnemonics is None:
run_mnemonics = mnemonics_for_run(run)
for key in ['bunchPatternTable', 'bunchPatternTable_SA3']:
if merge_with is not None and key in merge_with:
log.debug(f'Using {key} from merge_with dataset.')
return merge_with[key]
if key in run_mnemonics:
bpt = run.get_array(*run_mnemonics[key].values(),
name='bunchPatternTable')
log.debug(f'Loaded {key} from DataCollection.')
if merge_with is not None:
merge_with.update(merge_with.merge(bpt, join='inner'))
return bpt
log.debug('Could not find bunch pattern table.')
return None
[docs]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))
plt.plot(data['npulses_sase3'].trainId, data['npulses_sase3'], 'o-',
ms=3, label='SASE 3')
plt.xlabel('trainId')
plt.ylabel('pulses per train')
plt.plot(data['npulses_sase1'].trainId, data['npulses_sase1'], '^-',
ms=3, color='C2', label='SASE 1')
plt.legend()
[docs]def repRate(data=None, runNB=None, proposalNB=None, key='sase3'):
''' Calculates the pulse repetition rate (in kHz) in sase
according to the bunch pattern and assuming a grid of
4.5 MHz.
Inputs:
-------
data: xarray Dataset containing pulse pattern, needed if runNB is none
runNB: int or str, run number. Needed if data is None
proposal: int or str, proposal where to find the run. Needed if data is None
key: str in [sase1, sase2, sase3, scs_ppl], source for which the
repetition rate is calculated
Output:
-------
f: repetition rate in kHz
'''
if runNB is None and data is None:
raise ValueError('Please provide either the runNB + proposal or the data argument.')
if runNB is not None and proposalNB is None:
raise ValueError('Proposal is missing.')
if runNB is not None:
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), 'raw', runNB)
runDir = RunDirectory(runFolder)
bp_mnemo = _mnemonics_bp['bunchPatternTable']
if bp_mnemo['source'] not in runDir.all_sources:
raise ValueError('Source {} not found in run'.format(
bp_mnemo['source']))
else:
bp_table = runDir.get_array(bp_mnemo['source'],bp_mnemo['key'],
extra_dims=bp_mnemo['dim'])
a, b, mask = extractBunchPattern(bp_table, key=key)
else:
if key not in ['sase1', 'sase3']:
a, b, mask = extractBunchPattern(key=key, runDir=data.attrs['run'])
else:
a = data[key]
b = data[f'npulses_{key}']
a = a.where(b > 1, drop = True).values
if len(a)==0:
print('Not enough pulses to extract repetition rate')
return 0
f = 1/((a[0,1] - a[0,0])*12e-3/54.1666667)
return f