From ead72320f64d500ccfb10a369043a16a3c55709f Mon Sep 17 00:00:00 2001
From: Sven Augustin <sven.augustin@psi.ch>
Date: Sat, 30 Jan 2021 15:54:21 +0100
Subject: [PATCH] added Undulator(s)

---
 devices/undulator.py | 205 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 205 insertions(+)
 create mode 100644 devices/undulator.py

diff --git a/devices/undulator.py b/devices/undulator.py
new file mode 100644
index 0000000..cc412bd
--- /dev/null
+++ b/devices/undulator.py
@@ -0,0 +1,205 @@
+from time import sleep
+import numpy as np
+from scipy.interpolate import Akima1DInterpolator as Akima
+from epics import PV
+
+from slic.core.adjustable import Adjustable
+from slic.core.adjustable import PVAdjustable
+from slic.core.scanner.scanbackend import wait_for_all #, stop_all
+
+
+n_unds = [
+    8, 9, 10, 11, 12, 13,
+    15, 16, 17, 18, 19, 20, 21, 22
+]
+
+und_names = [f"SATUN{n:02}-UIND030" for n in n_unds]
+
+und_name_cal = "SATUN12-UIND030"
+
+
+energies = [
+    0.5457977557372193,
+    0.5429495415906611,
+    0.5401552739843691,
+    0.537355500731781,
+    0.534605067464375,
+    0.5318574959825555,
+    0.5291128804749312,
+    0.5264167029579244,
+    0.5237231992520139,
+    0.5210645197823921,
+]
+
+k_values = [
+    2.673,
+    2.681888888888889,
+    2.690777777777778,
+    2.699666666666667,
+    2.708555555555556,
+    2.7174444444444443,
+    2.7263333333333333,
+    2.735222222222222,
+    2.744111111111111,
+    2.753,
+]
+
+
+#energies = [
+#    0.549610621973092,
+#    0.5479937271701031,
+#    0.546395126105196,
+#    0.544814685256097,
+#    0.543218728674525,
+#    0.541660979262817,
+#    0.540074069530098,
+#    0.538501771379368,
+#    0.536939177980873,
+#    0.535392458877575,
+#    0.533842152732307,
+#    0.5323060597843,
+#    0.530767931907046,
+#    0.5292440325790481,
+#    0.5277161991879861,
+#    0.526199323076255,
+#    0.524699672767738,
+#]
+
+#k_values = [
+#    2.673,
+#    2.678,
+#    2.683,
+#    2.688,
+#    2.693,
+#    2.698,
+#    2.703,
+#    2.708,
+#    2.713,
+#    2.718,
+#    2.723,
+#    2.728,
+#    2.733,
+#    2.738,
+#    2.743,
+#    2.748,
+#    2.753,
+#]
+
+
+#energies = [
+#    1.1743862662276334,
+#    1.1561455825400897,
+#    1.1381577488951293,
+#    1.1204686653636284,
+#    1.1031371551015796,
+#    1.0860714511221887,
+#    1.069262314526579,
+#    1.05279814063446,
+#    1.0366146962391598,
+#    1.0206897775380315,
+#]
+
+#k_values = [
+#    1.5,
+#    1.52222222,
+#    1.54444444,
+#    1.56666667,
+#    1.58888889,
+#    1.61111111,
+#    1.63333333,
+#    1.65555556,
+#    1.67777778,
+#    1.7,
+#]
+
+
+
+class Undulators(Adjustable):
+
+    def __init__(self, energies=energies, k_values=k_values):
+        super().__init__(name="Athos Undulators")
+        self.adjs = {name: Undulator(name) for name in und_names}
+        self.eV_to_K = make_calib(energies, k_values)
+        self.K_to_eV = make_calib(k_values, energies)
+
+
+    def set_target_value(self, value, hold=False):
+        value /= 1000 # eV -> keV
+        k = self.eV_to_K(value)
+        if np.isnan(k):
+            print("K is nan for", value)
+            return
+        print(f"{k} <- {value}")
+
+        def change():
+            # replace by set_all_target_values_and_wait when print not needed anymore
+            tasks = []
+            for name, a in self.adjs.items():
+                print(f"{name} <- {k}")
+                t = a.set_target_value(k, hold=False)
+                tasks.append(t)
+            wait_for_all(tasks)
+
+        return self._as_task(change, hold=hold)
+
+
+    def get_current_value(self):
+        a = self.adjs[und_name_cal]
+        k = a.get_current_value()
+        energy = self.K_to_eV(k)
+        return float(energy) * 1000 # keV -> eV
+
+
+    def is_moving(self):
+        return any(a.is_moving() for a in self.adjs)
+
+
+
+class Undulator(PVAdjustable):
+
+    def __init__(self, name, accuracy=0.001):
+        pvname_setvalue = name + ":K_SET"
+        pvname_readback = name + ":K_READ"
+        super().__init__(pvname_setvalue, pvname_readback=pvname_readback, accuracy=accuracy, name=name)
+        self.energy = PVAdjustable(name + ":FELPHOTENE")
+
+
+
+def make_calib(x, y):
+    x, y = list(zip(*sorted(zip(x, y))))
+    return Akima(x, y)
+
+
+def get_map(k_min=min(k_values), k_max=max(k_values)):
+    und = Undulator(und_name_cal)
+    start_k_value = und.get_current_value()
+
+    k_values = np.linspace(k_min, k_max, 10)
+    energies = []
+    print()
+    print("mapping = {")
+    for k in k_values:
+        und.set_target_value(k).wait()
+        energy = und.energy.get_current_value()
+        print(f"    {energy}: {k},")
+        energies.append(energy)
+    print("}")
+    energies = np.array(energies)
+
+    print()
+    print("energies = [")
+    for energy in energies:
+        print(f"    {energy},")
+    print("]")
+
+    print()
+    print("k_values = [")
+    for k in k_values:
+        print(f"    {k},")
+    print("]")
+
+    und.set_target_value(start_k_value).wait()
+    return energies, k_values
+
+
+