wip - this doesnt work because fig is None?

Author: rerpha

src/ibex_bluesky_core/callbacks/__init__.py

@@ -24,7 +24,7 @@
 from ibex_bluesky_core.callbacks.fitting.livefit_logger import (
     LiveFitLogger,
 )
-from ibex_bluesky_core.callbacks.plotting import LivePlot
+from ibex_bluesky_core.callbacks.plotting import LivePlot, PlotPNGSaver
 
 logger = logging.getLogger(__name__)
 
@@ -55,6 +55,9 @@ def __init__(
         live_fit_logger_output_dir: str | PathLike[str] | None = None,
         live_fit_logger_postfix: str = "",
         human_readable_file_postfix: str = "",
+        save_plot_to_png: bool = True,
+        plot_png_output_dir : str | PathLike[str] | None = None,
+        plot_png_postfix: str = ""
     ) -> None:
         """A collection of ISIS standard callbacks for use within plans.
 
@@ -114,6 +117,7 @@ def _inner():
             live_fit_logger_postfix: the postfix to add to live fit logger.
             human_readable_file_postfix: optional postfix to add to human-readable file logger.
         """  # noqa
+        fig = None
         self._subs = []
         self._peak_stats = None
         self._live_fit = None
@@ -202,6 +206,9 @@ def start(self, doc: RunStart) -> None:
                     yerr=yerr,
                 )
             )
+            if save_plot_to_png:
+                self._subs.append(PlotPNGSaver(plot=fig, output_dir=plot_png_output_dir, postfix=plot_png_postfix
+                ))
 
     @property
     def live_fit(self) -> LiveFit:

src/ibex_bluesky_core/callbacks/_utils.py

@@ -1,6 +1,11 @@
 import os
+from datetime import datetime
 from pathlib import Path
 from platform import node
+from typing import Union
+from zoneinfo import ZoneInfo
+
+from event_model import Event, RunStart, RunStop
 
 OUTPUT_DIR_ENV_VAR = "IBEX_BLUESKY_CORE_OUTPUT"
 
@@ -31,3 +36,9 @@ def get_default_output_path() -> Path:
         if output_dir_env is None
         else Path(output_dir_env)
     )
+
+
+def format_time(doc: Union[Event, RunStart, RunStop]):
+    datetime_obj = datetime.fromtimestamp(doc[TIME])
+    title_format_datetime = datetime_obj.astimezone(ZoneInfo("UTC")).strftime("%Y-%m-%d_%H-%M-%S")
+    return title_format_datetime

src/ibex_bluesky_core/callbacks/file_logger.py

@@ -30,6 +30,7 @@
     UNITS,
     UNKNOWN_RB,
     get_instrument,
+    format_time,
 )
 
 logger = logging.getLogger(__name__)
@@ -67,19 +68,18 @@ def start(self, doc: RunStart) -> None:
         self.current_start_document = doc[UID]
 
         datetime_obj = datetime.fromtimestamp(doc[TIME])
-        title_format_datetime = datetime_obj.astimezone(ZoneInfo("UTC")).strftime(
-            "%Y-%m-%d_%H-%M-%S"
-        )
         rb_num = doc.get(RB, UNKNOWN_RB)
 
         # motors is a tuple, we need to convert to a list to join the two below
         motors = list(doc.get(MOTORS, []))
 
+        formatted_time = format_time(doc)
+
         self.filename = (
             self.output_dir
             / f"{rb_num}"
             / f"{get_instrument()}{'_' + '_'.join(motors) if motors else ''}_"
-            f"{title_format_datetime}Z{self.postfix}.txt"
+            f"{formatted_time}Z{self.postfix}.txt"
         )
         if rb_num == UNKNOWN_RB:
             logger.warning('No RB number found, saving to "%s"', UNKNOWN_RB)
@@ -91,7 +91,6 @@ def start(self, doc: RunStart) -> None:
         ]
         header_data = {k: v for k, v in doc.items() if k not in exclude_list}
 
-        formatted_time = datetime_obj.astimezone(ZoneInfo("UTC")).strftime("%Y-%m-%d %H:%M:%S")
         header_data[START_TIME] = formatted_time
 
         # make sure the parent directory exists, create it if not

src/ibex_bluesky_core/callbacks/fitting/livefit_logger.py

@@ -3,10 +3,8 @@
 import csv
 import logging
 import os
-from datetime import datetime
 from pathlib import Path
 from typing import Optional
-from zoneinfo import ZoneInfo
 
 import numpy as np
 from bluesky.callbacks import CallbackBase
@@ -17,11 +15,11 @@
 from ibex_bluesky_core.callbacks._utils import (
     DATA,
     RB,
-    TIME,
     UID,
     UNKNOWN_RB,
     get_default_output_path,
     get_instrument,
+    format_time,
 )
 from ibex_bluesky_core.callbacks.fitting import LiveFit
 
@@ -74,10 +72,7 @@ def start(self, doc: RunStart) -> None:
             doc (RunStart): The start bluesky document.
 
         """
-        datetime_obj = datetime.fromtimestamp(doc[TIME])
-        title_format_datetime = datetime_obj.astimezone(ZoneInfo("UTC")).strftime(
-            "%Y-%m-%d_%H-%M-%S"
-        )
+        title_format_datetime = format_time(doc)
         self.output_dir.mkdir(parents=True, exist_ok=True)
         self.current_start_document = doc[UID]
         file = f"{get_instrument()}_{self.x}_{self.y}_{title_format_datetime}Z{self.postfix}.txt"

src/ibex_bluesky_core/callbacks/plotting.py

@@ -1,14 +1,19 @@
 """IBEX plotting callbacks."""
 
 import logging
+import os
+from pathlib import Path
 from typing import Any
 
 import matplotlib
 import matplotlib.pyplot as plt
 from bluesky.callbacks import LivePlot as _DefaultLivePlot
-from bluesky.callbacks.core import get_obj_fields, make_class_safe
+from bluesky.callbacks.core import get_obj_fields, make_class_safe, CallbackBase
+from event_model import RunStop
 from event_model.documents import Event, RunStart
 
+from ibex_bluesky_core.callbacks._utils import get_instrument, get_default_output_path, format_time
+
 logger = logging.getLogger(__name__)
 
 
@@ -74,3 +79,24 @@ def start(self, doc: RunStart) -> None:
         """Process an start document (delegate to superclass, then show the plot)."""
         super().start(doc)
         show_plot()
+
+
+class PlotPNGSaver(CallbackBase):
+    """Save plots to PNG files on a run end."""
+
+    def __init__(
+        self,
+        plot: plt.Figure | None,
+        postfix: str,
+        output_dir: str | os.PathLike[str] | None,
+    ):
+        if plot is None:
+            raise ValueError("Plot does not exist - cannot save a PNG of it.")
+        self.plot = plot
+        self.postfix = postfix
+        self.output_dir = Path(output_dir or get_default_output_path())
+        super().__init__()
+
+    def stop(self, doc: RunStop):
+        filename = f"{get_instrument()}_{self.x}_{self.y}_{format_time(doc)}Z{self.postfix}.png"
+        self.plot.savefig(fname=self.output_dir / filename)

tests/callbacks/test_write_log_callback.py

@@ -39,7 +39,7 @@ def test_header_data_all_available_on_start(cb):
     mock_file.assert_called_with(result, "a", newline="\n", encoding="utf-8")
     writelines_call_args = mock_file().writelines.call_args[0][0]
     # time should have been renamed to start_time and converted to human readable
-    assert "start_time: 2024-10-04 13:43:43\n" in writelines_call_args
+    assert "start_time: 2024-10-04_13-43-43\n" in writelines_call_args
     assert f"uid: {uid}\n" in writelines_call_args
 
 
@@ -60,7 +60,7 @@ def test_no_rb_number_folder(cb):
     mock_file.assert_called_with(result, "a", newline="\n", encoding="utf-8")
     # time should have been renamed to start_time and converted to human readable
     writelines_call_args = mock_file().writelines.call_args[0][0]
-    assert "start_time: 2024-10-04 13:43:43\n" in writelines_call_args
+    assert "start_time: 2024-10-04_13-43-43\n" in writelines_call_args
     assert f"uid: {uid}\n" in writelines_call_args
 
 
@@ -81,7 +81,7 @@ def test_no_motors_doesnt_append_to_filename(cb):
     mock_file.assert_called_with(result, "a", newline="\n", encoding="utf-8")
     # time should have been renamed to start_time and converted to human readable
     writelines_call_args = mock_file().writelines.call_args[0][0]
-    assert "start_time: 2024-10-04 13:43:43\n" in writelines_call_args
+    assert "start_time: 2024-10-04_13-43-43\n" in writelines_call_args
     assert f"uid: {uid}\n" in writelines_call_args