Test setting plot title

Author: Tom-Willemsen

src/ibex_bluesky_core/callbacks/__init__.py

@@ -229,6 +229,10 @@ def start(self, doc: RunStart) -> None:
                     # where a fit result can be returned before
                     # the QtAwareCallback has had a chance to process it.
                     self._subs.append(self._live_fit)
+
+                # Sample 5000 points as this strikes a reasonable balance between displaying
+                # 'enough' points for almost any scan (even after zooming in on a peak), while
+                # not taking 'excessive' compute time to generate these samples.
                 self._subs.append(LiveFitPlot(livefit=self._live_fit, ax=ax, num_points=5000))
             else:
                 self._subs.append(self._live_fit)

src/ibex_bluesky_core/callbacks/_plotting.py

@@ -206,7 +206,7 @@ def __init__(
         y: str,
         ax: Axes,
         postfix: str,
-        output_dir: str | os.PathLike[str] | None,
+        output_dir: str | os.PathLike[str] | None = None,
     ) -> None:
         """Initialise the PlotPNGSaver callback.
 

src/ibex_bluesky_core/plans/reflectometry/_det_map_align.py

@@ -1,5 +1,6 @@
 """Reflectometry detector-mapping alignment plans."""
 
+import functools
 import logging
 from collections.abc import Generator
 from typing import Any, TypedDict
@@ -40,6 +41,19 @@
 logger = logging.getLogger(__name__)
 
 
+def _set_title_to_fit_result(
+    name: str, doc: dict[str, Any], *, fit_callback: LiveFit, ax: Axes
+) -> None:
+    fit_result = fit_callback.result
+    if fit_result is not None:
+        ax.set_title(
+            f"Best x0: {fit_result.params['x0'].value:.4f} +/- {fit_result.params['x0'].stderr:.4f}"
+        )
+    else:
+        ax.set_title("Fit failed")
+    plt.draw()
+
+
 def _height_scan_callback_and_fit(
     reducer: PeriodSpecIntegralsReducer,
     height: NamedMovable[float],
@@ -68,18 +82,12 @@ def _height_scan_callback_and_fit(
     for cb in height_scan_callbacks.subs:
         height_scan_ld.subscribe(cb)
 
-    def set_title_to_height_fit_result(name: str, doc: dict[str, Any]) -> None:
-        fit_result = height_scan_callbacks.live_fit.result
-        if fit_result is not None:
-            ax.set_title(
-                f"Best x0: {fit_result.params['x0'].value:.4f} "
-                f"+/- {fit_result.params['x0'].stderr:.4f}"
-            )
-        else:
-            ax.set_title("Fit failed")  # pragma: no cover
-        plt.draw()
-
-    height_scan_ld.subscribe(set_title_to_height_fit_result, "stop")
+    height_scan_ld.subscribe(
+        functools.partial(
+            _set_title_to_fit_result, fit_callback=height_scan_callbacks.live_fit, ax=ax
+        ),
+        "stop",
+    )
 
     return height_scan_ld, height_scan_callbacks.live_fit
 
@@ -135,24 +143,16 @@ def gaussian_max_y_guess(
     for cb in angle_scan_callbacks.subs:
         angle_scan_ld.subscribe(cb)
 
-    def set_title_to_angle_fit_result(name: str, doc: dict[str, Any]) -> None:
-        fit_result = angle_scan_callbacks.live_fit.result
-        if fit_result is not None:
-            ax.set_title(
-                f"Best x0: {fit_result.params['x0'].value:.4f} "
-                f"+/- {fit_result.params['x0'].stderr:.4f}"
-            )
-        else:
-            ax.set_title("Fit failed")  # pragma: no cover
-        plt.draw()
-
-    angle_scan_ld.subscribe(set_title_to_angle_fit_result, "stop")
-
-    # Make sure the Plot PNG saving happens *after* setting plot title to fit result...
     angle_scan_ld.subscribe(
-        PlotPNGSaver(x=angle_name, y=counts_name, ax=ax, postfix="", output_dir=None)
+        functools.partial(
+            _set_title_to_fit_result, fit_callback=angle_scan_callbacks.live_fit, ax=ax
+        ),
+        "stop",
     )
 
+    # Make sure the Plot PNG saving happens *after* setting plot title to fit result...
+    angle_scan_ld.subscribe(PlotPNGSaver(x=angle_name, y=counts_name, ax=ax, postfix=""))
+
     return angle_scan_ld, angle_scan_callbacks.live_fit
 
 

tests/plans/test_det_map_align.py

@@ -16,6 +16,7 @@
     angle_scan_plan,
     height_and_angle_scan_plan,
 )
+from ibex_bluesky_core.plans.reflectometry._det_map_align import _set_title_to_fit_result
 
 
 @pytest.fixture
@@ -122,3 +123,24 @@ def test_det_map_align_bad_angle_map_shape(RE, dae, height):
                     angle_map=np.array([21, 22, 23, 24]),
                 )
             )
+
+
+def test_set_title_to_fit_result_failed_fit():
+    ax = MagicMock()
+    live_fit = MagicMock()
+    live_fit.result = None
+
+    _set_title_to_fit_result(fit_callback=live_fit, ax=ax, name="stop", doc={})
+
+    ax.set_title.assert_called_once_with("Fit failed")
+
+
+def test_set_title_to_fit_result_good_fit():
+    ax = MagicMock()
+    live_fit = MagicMock()
+    live_fit.result.params["x0"].value = 1.23456789
+    live_fit.result.params["x0"].stderr = 9.87654321
+
+    _set_title_to_fit_result(fit_callback=live_fit, ax=ax, name="stop", doc={})
+
+    ax.set_title.assert_called_once_with("Best x0: 1.2346 +/- 9.8765")