Source code for toolbox_scs.misc.undulator

__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