From b8f7dabd47bce21b6485f887e3cd3ce935a2cdf2 Mon Sep 17 00:00:00 2001
From: x10sa <x10sa@psi.ch>
Date: Tue, 25 Aug 2026 15:35:09 +0200
Subject: [PATCH] bug fixes

---
 pxii_bec/macros/mxcomm_plot.py | 34 +++++++++++++++++++++++++++++++++-
 pxii_bec/macros/mxcomm_scan.py | 30 ++++++++++++++++++++++++------
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/pxii_bec/macros/mxcomm_plot.py b/pxii_bec/macros/mxcomm_plot.py
index 9868911..ab9b7b5 100644
--- a/pxii_bec/macros/mxcomm_plot.py
+++ b/pxii_bec/macros/mxcomm_plot.py
@@ -53,7 +53,7 @@ def plot_live_data(
     wf.plot(device_x=motor_name, device_y=signal_name)
 
 
-def plot_fitted_data(data, fit_result, window_name="Fitting", plot_widget=None):
+def plot_fitted_data(data, fit_result, window_name="Fitting", plot_widget=None, title=None):
     """
     Plot completed scan data and its fitted curve in BEC.
 
@@ -93,3 +93,35 @@ def plot_fitted_data(data, fit_result, window_name="Fitting", plot_widget=None):
 
     # Display the fitted curve as a line without symbols.
     wf.get_curve("Fit").set(symbol_size=0)
+
+def plot_scan_data(
+    data,
+    window_name="Fitting",
+    plot_widget=None,
+    title=None,
+):
+    """Plot completed scan data without a fitted curve."""
+
+    if isinstance(plot_widget, str):
+        dock_area = bec.gui.windows[window_name]
+        wf = getattr(dock_area, plot_widget)
+
+    elif plot_widget is not None:
+        wf = plot_widget
+
+    else:
+        wf, text_widget = select_bec_window(window_name)
+        text_widget.set_plain_text("Fit rejected")
+
+    # Clear the live plot/subscription
+    wf.clear_all()
+
+    wf.title = title or f"Scan: {data['scan_number']}"
+    wf.x_label = data["motor_name"]
+    wf.y_label = data["signal_name"]
+
+    wf.plot(
+        x=data["x_data"],
+        y=data["y_data"],
+        label="Data",
+    )
diff --git a/pxii_bec/macros/mxcomm_scan.py b/pxii_bec/macros/mxcomm_scan.py
index 1de4f30..92c49b3 100755
--- a/pxii_bec/macros/mxcomm_scan.py
+++ b/pxii_bec/macros/mxcomm_scan.py
@@ -57,7 +57,7 @@ def calculate_step_size(start: float, stop: float, steps: int) -> float:
     return round((stop - start) / (steps - 1), 5)
 
 
-def move_to_position(motor_device, motor_name: str, position: float, data: dict):
+def move_to_scan_position(motor_device, motor_name: str, position: float, data: dict):
     """
     Function to move a specified motor device to a given position.
 
@@ -100,6 +100,9 @@ def move_to_position(motor_device, motor_name: str, position: float, data: dict)
     scans.umv(motor_device, motor_position, relative=False)
     print(f"\n Moving {motor_name} to position {motor_position: .3f}")
 
+def restore_position(motor_device, position):
+    """Return a motor to its position before the scan."""
+    scans.umv(motor_device, position, relative=False)    
 
 def go_to_peak(
     motor_device,
@@ -218,10 +221,25 @@ def go_to_peak(
         print("Peak fit rejected:")
         for problem in check["problems"]:
             print(f"  - {problem}")
-        print(f"Returning {motor_name} to starting position " f"{starting_position:.4f}")
-        scans.umv(motor_device, starting_position, relative=False)
+
+        if plot:
+            plot_scan_data(
+                data,
+                window_name=window_name,
+                plot_widget=plot_widget,
+                title=f"{motor_name} scan - FIT REJECTED",
+            )
+        print(f"Returning motor to start position, {starting_position:.5g}")
+        restore_position(
+            motor_device,
+            starting_position,
+        )
+
         return None
 
+
+
+
     # Plot the fitted data if plot = True
     if plot:
         plot_fitted_data(data, fit_result, window_name=window_name, plot_widget=plot_widget)
@@ -230,18 +248,18 @@ def go_to_peak(
     if gomax:
         value = fit_result["x_max"]
         print(f"Max position is at {value}")
-        move_to_position(data["motor_device"], data["motor_name"], fit_result["x_max"], data)
+        move_scan_to_position(data["motor_device"], data["motor_name"], fit_result["x_max"], data)
     else:
         if gap:
             if fit_result["centre"] < 4.5:
                 umv(dev.id_gap, 4.5)
             else:
-                move_to_position(
+                move_to_scan_position(
                     data["motor_device"], data["motor_name"], fit_result["centre"], data
                 )
         else:
             # Safely move the motor to the peak position
-            move_to_position(data["motor_device"], data["motor_name"], fit_result["centre"], data)
+            move_to_scan_position(data["motor_device"], data["motor_name"], fit_result["centre"], data)
 
     return {
         "centre": fit_result["centre"],