__all__ = [
'get_undulator_config',
]
import numpy as np
import xarray as xr
from toolbox_scs.load import load_run_values
import matplotlib.pyplot as plt
def get_undulator_cells(run, park_pos=62.0):
rvalues = run.get_run_values('SA3_XTD4_UND/DOOCS/UNDULATOR_CELLS')
rvalues = {k: rvalues[k] for k in rvalues.keys()
if '.gapApplied.value' in k or '.kApplied.value' in k}
cells = list(range(2, 13)) + list(range(14, 24))
keys = np.unique([k.replace('.gapApplied.value',
'').replace('.kApplied.value',
'') for k in rvalues])
assert len(keys) == len(cells)
result = [[], [], []]
names = ['gap', 'K', 'cell_name']
for i, k in enumerate(keys):
result[0].append(rvalues[k + '.gapApplied.value'])
result[1].append(rvalues[k + '.kApplied.value'])
result[2].append(k)
result = xr.merge([xr.DataArray(result[i], dims='cell',
coords={'cell': cells},
name=names[i]) for i in range(3)])
result['closed'] = result['gap'] < park_pos
return result
def plot_undulator_config(ds, park_pos):
fig, ax = plt.subplots(figsize=(6, 3))
ax.bar(ds.cell, ds.gap - park_pos-1, bottom=park_pos+1, alpha=0.5)
for c in ds.cell:
ax.text(c-.25, park_pos/2, f"K={ds.sel(cell=c).K.values:.4f}",
rotation='vertical')
ax.set_ylim(0, park_pos+1)
ax.invert_yaxis()
ax.set_xlabel('CELL #')
ax.set_ylabel('gap size')
[docs]def get_undulator_config(run, park_pos=62.0, plot=True):
'''
Extract the undulator cells configuration from a given run.
The gap size and K factor as well as the magnetic chicane delay and photon
energy of colors 1, 2 and 3 are compiled into an xarray Dataset.
Note:
This function looks at run control values, it does not reflect any change
of values during the run. Do not use to extract configuration when scanning
the undulator.
Parameters
----------
run: EXtra-Data DataCollection
The run containing the undulator information
park_pos: float, optional
The parked position of a cell (i.e. when fully opened)
plot: bool, optional
If True, plot the undulator cells configuration
Returns
-------
cells: xarray Dataset
The resulting dataset of the undulator configuration
'''
ds = get_undulator_cells(run, park_pos)
rvalues = load_run_values(run)
attrs = {}
if 'UND' in rvalues:
attrs = {f'color_{i+1}_keV': rvalues[k] for
i, k in enumerate(['UND', 'UND2', 'UND3'])}
if 'MAG_CHICANE_DELAY' in rvalues:
attrs['MAG_CHICANE_DELAY'] = rvalues['MAG_CHICANE_DELAY']
for k in attrs:
ds.attrs[k] = attrs[k]
if plot:
plot_undulator_config(ds, park_pos)
return ds