wip widget

debye_bec/bec_widgets/widgets/scan_control_advanced/scan_control_advanced.py

@@ -3,14 +3,17 @@ import sys
 import numpy as np
 import pyqtgraph as pg
 from bec_lib import bec_logger
+from bec_lib.endpoints import MessageEndpoints
 from bec_widgets.utils.bec_dispatcher import BECDispatcher
 from bec_widgets.utils.colors import apply_theme
 from bec_widgets.widgets.control.scan_control.scan_control import ScanControl
-from qtpy.QtCore import Qt, Signal
+from qtpy.QtCore import Qt, QTimer, Signal
 
 # pylint: disable=E0611
 from qtpy.QtWidgets import QApplication, QDoubleSpinBox, QPushButton, QVBoxLayout, QWidget
 
+from ....devices.mo1_bragg.mo1_bragg_utils import compute_spline
+
 logger = bec_logger.logger
 
 SHOW_MOTION_PROFILE = [
@@ -28,6 +31,25 @@ class ScanControlAdvanced(ScanControl):
         super().__init__(*args, **kwargs)
         self.show_motion_profile(self._selected_scan)
 
+        # self.bec_dispatcher.connect_slot(
+        #     self.mo1_update, MessageEndpoints.device_readback("mo1_bragg_crystal_current_d_spacing")
+        # )
+
+        # Timer: update d_spacing every 1 second
+        self.d_spacing = 0
+        self._timer = QTimer(self)
+        self._timer.setInterval(1000)
+        self._timer.timeout.connect(self.update_d_spacing)
+        self._timer.start()
+
+    def update_d_spacing(self):
+        d_spacing = self.dev.mo1_bragg.crystal.current_d_spacing.read(cached=True)[
+            "mo1_bragg_crystal_current_d_spacing"
+        ]["value"]
+        if d_spacing != self.d_spacing:
+            self.d_spacing = d_spacing
+            self.motion_profile_widget.update_plot(d_spacing=d_spacing)
+
     def _add_metadata_form(self):
         self.layout.addWidget(self.motion_profile_widget)
         super()._add_metadata_form()
@@ -67,10 +89,15 @@ class ScanControlAdvanced(ScanControl):
     def scan_parameter_changed(self, *_):
         params = self.get_scan_parameters()[1]
         logger.info(f"Scan parameters: {params}")
-        self.motion_profile_widget.update_plot(params)
+        self.motion_profile_widget.update_plot(
+            scan_name=self._selected_scan, scan_parameters=params
+        )
 
 
 PLOT_RESOLUTION = 100
+H = 6.62606957e-34
+E = 1.602176634e-19
+C = 299792458
 
 
 class MotionProfileWidget(QWidget):
@@ -81,17 +108,88 @@ class MotionProfileWidget(QWidget):
         super().__init__()
         self.parent = parent
         layout = QVBoxLayout(self)
-        self.plot = pg.PlotWidget()
-        self.plot.enableAutoRange()
-        layout.addWidget(self.plot)
+        self.plot_pos = pg.PlotWidget()
+        self.plot_vel = pg.PlotWidget()
+        layout.addWidget(self.plot_pos)
+        layout.addWidget(self.plot_vel)
+        self.scan_name = None
+        self.scan_parameters = None
+        self.d_spacing = None
 
-    def update_plot(self, scan_parameters):
-        x = np.linspace(scan_parameters["start"], scan_parameters["stop"], PLOT_RESOLUTION)
-        x_sin = np.linspace(0, np.pi, PLOT_RESOLUTION)
-        y = 1 * np.sin(x_sin)
+    def update_plot(self, scan_name=None, scan_parameters=None, d_spacing=None):
+        if scan_name is not None:
+            self.scan_name = scan_name
+        if scan_parameters is not None:
+            self.scan_parameters = scan_parameters
+        if d_spacing is not None:
+            self.d_spacing = d_spacing
+        if self.scan_name is None:
+            return
+        if self.scan_parameters is None:
+            return
+        if self.d_spacing is None:
+            return
+        if d_spacing == 0:
+            return
 
-        self.plot.clear()
-        self.plot.plot(x, y)
+        x_time = np.linspace(0, self.scan_parameters["scan_time"], PLOT_RESOLUTION)
+
+        self.plot_pos.clear()
+        self.plot_vel.clear()
+
+        start_angle = self.energy_to_angle(self.scan_parameters["stop"])
+        stop_angle = self.energy_to_angle(self.scan_parameters["start"])
+        logger.info(f"start: {start_angle}, stop: {stop_angle}")
+        if "xas_simple_scan" in self.scan_name:
+            amp = (stop_angle - start_angle) / 2
+            two_pi_f = 2 * np.pi * 1 / (2 * self.scan_parameters["scan_time"])
+            pos = (start_angle + stop_angle) / 2 + amp * np.cos(x_time * two_pi_f)
+            logger.info(f"pos: {pos}")
+            vel = two_pi_f * amp * np.sin(x_time * two_pi_f)
+            energy = self.angle_to_energy(pos)
+            self.plot_pos.plot(x_time, energy)
+            self.plot_vel.plot(energy, vel)
+        if "xas_advanced_scan" in self.scan_name:
+            if (
+                self.scan_parameters["e_kink"] < self.scan_parameters["start"]
+                or self.scan_parameters["e_kink"] > self.scan_parameters["stop"]
+            ):
+                return
+            if self.scan_parameters["scan_time"] == 0:
+                return
+            e_kink_deg = self.energy_to_angle(self.scan_parameters["e_kink"])
+            pos, vel, t = compute_spline(
+                start_angle,
+                stop_angle,
+                self.scan_parameters["p_kink"],
+                e_kink_deg,
+                self.scan_parameters["scan_time"],
+            )
+            x_time = np.cumsum(t) / 1000
+            energy = self.angle_to_energy(pos)
+            self.plot_pos.plot(-np.flip(x_time) + x_time[-1], np.flip(energy))
+            self.plot_vel.plot(energy, vel)
+
+    def energy_to_angle(self, energy: int | float):
+        if self.d_spacing is None:
+            return 0
+        if energy <= 0:
+            return 0
+        wl = C * H / (E * energy)
+        val = wl / (2 * self.d_spacing * 1e-10)
+        return np.asin(val) / np.pi * 180
+
+    def angle_to_energy(self, angle: np.ndarray | int | float):
+        if self.d_spacing is None:
+            return 0
+        if isinstance(angle, np.ndarray):
+            if angle.any() <= 0:
+                return 0
+        else:
+            if angle <= 0:
+                return 0
+        wl = 2 * self.d_spacing * 1e-10 * np.sin(angle / 180 * np.pi)
+        return C * H / (E * wl)
 
 
 if __name__ == "__main__":