From eec7f32caecb2f9f5a87191af39e78fb8bf3b487 Mon Sep 17 00:00:00 2001
From: Laurent Mercadier <laurent.mercadier@xfel.eu>
Date: Thu, 1 Dec 2022 06:46:18 +0100
Subject: [PATCH] Add undulator config function
---
src/toolbox_scs/misc/__init__.py | 2 +
src/toolbox_scs/misc/undulators.py | 82 ++++++++++++++++++++++++++++++
2 files changed, 84 insertions(+)
create mode 100644 src/toolbox_scs/misc/undulators.py
diff --git a/src/toolbox_scs/misc/__init__.py b/src/toolbox_scs/misc/__init__.py
index 3ff302d..0f544c8 100644
--- a/src/toolbox_scs/misc/__init__.py
+++ b/src/toolbox_scs/misc/__init__.py
@@ -1,9 +1,11 @@
from .bunch_pattern import *
from .bunch_pattern_external import *
from .laser_utils import *
+from .undulators import *
__all__ = (
bunch_pattern.__all__
+ bunch_pattern_external.__all__
+ laser_utils.__all__
+ + undulators.__all__
)
diff --git a/src/toolbox_scs/misc/undulators.py b/src/toolbox_scs/misc/undulators.py
new file mode 100644
index 0000000..9169044
--- /dev/null
+++ b/src/toolbox_scs/misc/undulators.py
@@ -0,0 +1,82 @@
+__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')
+
+
+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
--
GitLab