From 8e8e5de83e669707bd0f6075b792b88865ca4eff Mon Sep 17 00:00:00 2001
From: augustin_s <sven.augustin@psi.ch>
Date: Tue, 12 Apr 2022 12:09:47 +0000
Subject: [PATCH] added q-space device

---
 qspace.py | 168 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 168 insertions(+)
 create mode 100644 qspace.py

diff --git a/qspace.py b/qspace.py
new file mode 100644
index 0000000..08ef38c
--- /dev/null
+++ b/qspace.py
@@ -0,0 +1,168 @@
+from types import SimpleNamespace
+
+from diffcalc.hkl.calc import HklCalculation
+from diffcalc.hkl.constraints import Constraints
+from diffcalc.hkl.geometry import Position
+from diffcalc.ub.calc import UBCalculation
+
+from slic.core.adjustable import Adjustable, PVAdjustable
+from slic.devices.general.motor import Motor
+from slic.devices.device import Device
+from slic.devices.simpledevice import SimpleDevice
+
+
+INDICES = {
+    "h": 0,
+    "k": 1,
+    "l": 2
+}
+
+
+class QSpace3D(Device):
+
+    def __init__(self, ID, eta, chi, phi, wavelength, **kwargs): # if the diffcalc objects need some initial parameters, add them here and fill them in below...
+        super().__init__(ID, **kwargs)
+
+        # collect the motors in a device for nicer printing
+        self.motors = motors = SimpleDevice(
+            ID,
+            eta = eta,
+            chi = chi,
+            phi = phi
+        )
+
+        self.wavelength = wavelength
+
+        # some diffcalc objects ...
+        ub = UBCalculation()
+        cons = Constraints()
+        hkl = HklCalculation(ub, cons)
+
+        # ... collected in a namespace
+        self.dc = dc = SimpleNamespace(
+            ub = ub,
+            cons = cons,
+            hkl = hkl
+        )
+
+        # and the individual coordinates:
+        self.h = QSpace1D(ID + "-H", "h", dc, motors, wavelength)
+        self.k = QSpace1D(ID + "-K", "k", dc, motors, wavelength)
+        self.l = QSpace1D(ID + "-L", "l", dc, motors, wavelength)
+
+
+    # it might be nice (but optional) to add some shortcuts like this:
+    def set_lattice(self, *args, **kwargs):
+        self.dc.ub.set_lattice(*args, **kwargs)
+        self.dc.ub.calc_ub() # not sure whether this needs to be called explicitly?
+
+    # if I understand the examples/code correctly, then some more method calls are mandatory?
+    # those should probably all get shortcuts...
+
+
+
+class QSpace1D(Adjustable):
+
+    def __init__(self, ID, index, dc, motors, wavelength):
+        super().__init__(ID)
+        self.index = index
+        self.dc = dc
+        self.motors = motors
+        self.wavelength = wavelength
+
+
+    # the following three methods are mandatory:
+
+    def get_current_value(self):
+        hkl = self._get_hkl()
+        i = self._get_index()
+        return hkl[i]
+
+    def set_target_value(self, value):
+        # get all current indices
+        hkl = self._get_hkl()
+        i = self._get_index()
+        # insert the target value into the right spot
+        hkl[i] = value
+        self._set_hkl(self, *hkl)
+
+    def is_moving(self):
+        ms = self.motors
+        return ms.eta.moving or ms.chi.moving or ms.phi.moving
+
+
+    # some helpful things:
+
+    def get_wavelength(self):
+        return self.wavelength.get_current_value()
+
+    def _get_index(self):
+        i = self.index
+        if isinstance(i, str):
+            i = i.casefold()
+            i = INDICES.get(i, i)
+        if i not in INDICES.values():
+            allowed = INDICES.keys() | INDICES.values()
+            allowed = sorted(allowed, key=str)
+            raise ValueError(f"index must be from {allowed} but is {repr(i)}")
+        return i
+
+    def _get_hkl(self):
+        angles = self._get_angles()
+        return self._calc_hkl(*angles)
+
+    def _set_hkl(self, h, k, l):
+        angles = self._calc_angles(self, h, k, l)
+        self._set_angles(self, *angles)
+
+    def _get_angles(self):
+        ms = self.motors
+        eta = ms.eta.get_current_value()
+        chi = ms.chi.get_current_value()
+        phi = ms.phi.get_current_value()
+        return eta, chi, phi
+
+    def _set_angles(self, eta, chi, phi):
+        ms = self.motors
+        t_eta = ms.eta.set_target_value(eta)
+        t_chi = ms.chi.set_target_value(chi)
+        t_phi = ms.phi.set_target_value(phi)
+        # wait for all motors to arrive
+        tasks = (t_eta, t_chi, t_phi)
+        for t in tasks:
+            t.wait()
+
+    def _calc_hkl(self, eta, chi, phi):
+        wl = self.get_wavelength()
+        pos = Position(eta=eta, chi=chi, phi=phi)
+        hkl = self.dc.hkl.get_hkl(pos, wl)
+        return hkl
+
+    def _calc_angles(self, h, k, l):
+        wl = self.get_wavelength()
+        pos, _virtual_angles = next(iter(
+            self.dc.hkl.get_position(h, k, l, wl)
+        ))
+        return pos.eta, pos.chi, pos.phi
+
+
+
+
+
+# these point to the different motors
+eta = Motor("SOMETHING:ETA")
+chi = Motor("SOMETHING:CHI")
+phi = Motor("SOMETHING:PHI")
+
+# and this to the machine wavelength (maybe this needs to be calculated from the energy? then we should add a Wavelength wrapper...)
+wl = PVAdjustable("MACHINE:WAVELENGTH")
+
+# put it all together:
+q = QSpace3D("SOMETHING:Q", eta, chi, phi, wl)
+
+## in ipython
+#q.set_lattice("SiO2", 4.913, 5.405)
+#q.set_target_value(1.5)
+
+
+