fix(timepix): repeat an acquisition whose TimePixFly frame closes before the exposure ends

superxas_bec/devices/timepix/timepix.py

@@ -12,6 +12,9 @@ import os
 import threading
 import time
 import traceback
+from collections import deque
+from dataclasses import dataclass
+from functools import partial
 from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
@@ -21,6 +24,8 @@ from bec_server.scan_server.scans.scan_base import ScanInfo as ScanServerScanInf
 from ophyd import ADBase
 from ophyd import Component as Cpt
 from ophyd import Staged
+from ophyd.status import StatusTimeoutError, WaitTimeoutError
+from ophyd.utils import InvalidState
 from ophyd_devices import (
     AsyncSignal,
     CompareStatus,
@@ -218,6 +223,31 @@ TIMEPIX_FLY_BACKEND_DEFAULT_REDIS_TOPIC = "public/timepix_fly_backend/raw_data"
 # scan_name of superxas_bec.scans.timepix_scan.TimepixScan
 TIMEPIX_SCAN_NAME = "timepix_scan"
 
+# How long a repeated acquisition waits for the images of the lost one before it starts.
+LOST_ACQUISITION_IMAGE_TIMEOUT = 2.0
+
+
+@dataclass
+class XesAcquisition:
+    """
+    One camera acquisition for which TimePixFly was armed. TimePixFly sends exactly one
+    StartFrame ... EndFrame sequence per armed acquisition, in order, so msg_buffer_callback
+    matches each EndFrame to the oldest pending acquisition.
+
+    Attributes:
+        exposure (float): seconds the detector needs for the images of this trigger.
+        repeat_allowed (bool): the acquisition may be repeated if its raw data are lost.
+        acquire_time (float | None): time.monotonic() when Acquire was set, None before.
+        lost (bool): TimePixFly closed its frame before the exposure could end.
+        done (bool): the trigger status of this acquisition has completed.
+    """
+
+    exposure: float
+    repeat_allowed: bool
+    acquire_time: float | None = None
+    lost: bool = False
+    done: bool = False
+
 
 class Timepix(PSIDeviceBase, TimePixControl):
     """
@@ -329,6 +359,13 @@ class Timepix(PSIDeviceBase, TimePixControl):
         self._settings_before_timepix_scan: dict[str, Any] | None = None
         # Log the "on/off window does not fit" warning only once per scan.
         self._on_off_window_warned = False
+        # Acquisitions armed in TimePixFly whose EndFrame has not arrived yet, oldest first.
+        self._xes_acquisitions: deque[XesAcquisition] = deque()
+        self._xes_acquisitions_lock = threading.Lock()
+        # Images of a lost acquisition still to arrive; they are not published.
+        self._images_to_skip = 0
+        self._exposure_per_trigger = 0.0
+        self._images_per_trigger = 1
 
     def stage(self) -> list[object] | StatusBase:  # type: ignore
         """Stage the device.
@@ -414,6 +451,8 @@ class Timepix(PSIDeviceBase, TimePixControl):
             end_frame (dict): The EndFrame. Dictionary representation of detailed structure
                                 described in model .timepix_fly_client.timepix_fly_interface.TimepixEndFrame
         """
+        if self._frame_of_lost_acquisition(end_frame):
+            return  # on_trigger repeats the acquisition; publish nothing for this one
         n_energy_points = start_frame.get("NumEnergyPoints", None)
         troin = start_frame["TRoiN"]
         if troin != self._troin:
@@ -494,6 +533,59 @@ class Timepix(PSIDeviceBase, TimePixControl):
         )
         logger.debug(f"Device data set for Timepix with {tds_period}, {tds_total_events}")
 
+    def _frame_of_lost_acquisition(self, end_frame: dict) -> bool:
+        """
+        Match the EndFrame to its acquisition and check whether that acquisition was lost.
+
+        A frame without any TDC period holds no XES data. If it arrives before the exposure of
+        its acquisition can have ended, TimePixFly's raw-data connection closed before the
+        detector measured, and the raw data of that acquisition never reached TimePixFly
+        (scan S00437, point 1: the frame closed 1.2 s before the detector delivered the image).
+        Such an acquisition is marked lost and on_trigger repeats it once, if allowed. A frame
+        without TDC periods after the exposure means the detector recorded no TDC pulse; it is
+        published as it is.
+
+        Args:
+            end_frame (dict): EndFrame received from TimePixFly.
+
+        Returns:
+            bool: True if the frame belongs to a lost acquisition that will be repeated.
+        """
+        with self._xes_acquisitions_lock:
+            if not self._xes_acquisitions:
+                return False  # no acquisition armed by this device, e.g. after a stop
+            acquisition = self._xes_acquisitions.popleft()
+            if end_frame.get("periods", None) != 0:
+                return False
+            if acquisition.acquire_time is None:
+                timing = "before Acquire was set"
+            else:
+                elapsed = time.monotonic() - acquisition.acquire_time
+                timing = f"{elapsed:.3f} s after Acquire was set"
+                if elapsed >= acquisition.exposure:
+                    logger.error(
+                        f"{self.name}: TimePixFly frame without any TDC period, {timing} "
+                        f"({acquisition.exposure:.3f} s exposure). The detector recorded no TDC "
+                        "pulse: check the TDC input signal. This point has no XES data."
+                    )
+                    return False
+            if not acquisition.repeat_allowed or acquisition.done:
+                logger.error(
+                    f"{self.name}: TimePixFly closed the frame {timing}, before the "
+                    f"{acquisition.exposure:.3f} s exposure could end, without any TDC period. "
+                    "The raw data of this point did not reach TimePixFly and the acquisition "
+                    "cannot be repeated. This point has no XES data."
+                )
+                return False
+            acquisition.lost = True
+            self._images_to_skip += self._images_per_trigger
+        logger.warning(
+            f"{self.name}: TimePixFly closed the frame {timing}, before the "
+            f"{acquisition.exposure:.3f} s exposure could end, without any TDC period. The raw "
+            "data of this acquisition did not reach TimePixFly; repeating the acquisition."
+        )
+        return True
+
     def _put_on_off_levels(self, data_1: np.ndarray, data_2: np.ndarray) -> None:
         """
         Publish the laser on/off levels of a timepix_scan for both energy groups.
@@ -717,6 +809,11 @@ class Timepix(PSIDeviceBase, TimePixControl):
         """Callback for image unique ID updates to trigger preview update."""
         if value == old_value:
             return  # No new image, or counter reset
+        with self._xes_acquisitions_lock:
+            if self._images_to_skip > 0:
+                # Image of an acquisition whose XES data were lost; its repetition replaces it.
+                self._images_to_skip -= 1
+                return
         try:
             # Get new image data
             width = self.image.array_size.width.get()
@@ -754,6 +851,9 @@ class Timepix(PSIDeviceBase, TimePixControl):
             )
         burst_images = scan_msg.scan_parameters.get("frames_per_trigger", 1)
         self._n_images = scan_msg.num_points * burst_images
+        self._images_per_trigger = burst_images
+        self._exposure_per_trigger = burst_images * exp_time - self._readout_time
+        self._clear_xes_acquisitions()
         # Camera has to be set to burst_images, each step will get an individual trigger
         self.cam.acquire_time.set(exp_time - self._readout_time).wait(timeout=self._pv_timeout)
         self.cam.acquire_period.set(exp_time).wait(timeout=self._pv_timeout)
@@ -868,51 +968,155 @@ class Timepix(PSIDeviceBase, TimePixControl):
         return status
 
     def on_trigger(self) -> DeviceStatus | StatusBase | None:
-        """Called when the device is triggered."""
+        """
+        Called when the device is triggered.
 
-        status_camera = TransitionStatus(
+        With XES enabled, TimePixFly is armed before the camera starts. If TimePixFly closes the
+        frame of the acquisition before its exposure could end (see _frame_of_lost_acquisition),
+        the acquisition is repeated once and the returned status resolves after the repetition.
+        This needs the HDF5 plugin to be off (as in timepix_scan): with image writing on, a
+        repetition would add an image to the file, so the point is kept without XES data.
+        """
+        if self.enable_xes is not True:
+            status_camera = self._camera_acquisition_status()
+            self.cancel_on_stop(status_camera)
+            self.cam.acquire.put(1)
+            return status_camera
+
+        point_status = DeviceStatus(self)
+        self.cancel_on_stop(point_status)
+        repeat_allowed = self.hdf.enable.get() != "Enable"
+        self._start_xes_acquisition(point_status, repeat_allowed=repeat_allowed)
+        return point_status
+
+    def _camera_acquisition_status(self) -> TransitionStatus:
+        """Status that resolves once the camera has acquired (AcquireBusy DONE -> ACQUIRING -> DONE)."""
+        return TransitionStatus(
             self.cam.acquire_busy, [ACQUIRESTATUS.DONE, ACQUIRESTATUS.ACQUIRING, ACQUIRESTATUS.DONE]
         )
-        status_backend = None
+
+    def _start_xes_acquisition(self, point_status: DeviceStatus, repeat_allowed: bool) -> None:
+        """
+        Arm TimePixFly and start one camera acquisition for the trigger status point_status.
+        _on_xes_acquisition_done resolves point_status, or repeats a lost acquisition.
+
+        Args:
+            point_status (DeviceStatus): trigger status returned by on_trigger.
+            repeat_allowed (bool): the acquisition may be repeated if its raw data are lost.
+
+        Raises:
+            TimeoutError: if TimePixFly does not reach the 'config' or 'await_connection' state.
+        """
+        status_camera = self._camera_acquisition_status()
         # First we make sure that the backend reach 'config' state. This needs to happend before each trigger.
-        if self.enable_xes is True:
-            status_backend_config = DeviceStatus(self)
-            self.cancel_on_stop(status_backend_config)
-            self.backend.timepix_fly_client.add_status_callback(
-                status=status_backend_config,
-                success=[TimePixFlyStatus.CONFIG],
-                error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN],
-            )
-            try:
-                status_backend_config.wait(timeout=5)
-            except TimeoutError:
-                # pylint:disable=raise-missing-from
-                raise TimeoutError(
-                    f"TimePixFly backend of device {self.name} failed to reach 'config' state in trigger"
-                )
-
-            # Prepare backend to be ready to receive trigger
-            status = self.backend.on_trigger()
-            try:
-                status.wait(timeout=5)  # Wait until backend trigger is done
+        status_backend_config = DeviceStatus(self)
+        self.cancel_on_stop(status_backend_config)
+        self.backend.timepix_fly_client.add_status_callback(
+            status=status_backend_config,
+            success=[TimePixFlyStatus.CONFIG],
+            error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN],
+        )
+        try:
+            status_backend_config.wait(timeout=5)
+        except TimeoutError:
             # pylint:disable=raise-missing-from
-            except Exception:
-                logger.error(
-                    f"TimePixFly backend of device {self.name} failed to prepare for trigger"
-                )
-                raise TimeoutError(
-                    f"TimePixFly backend of device {self.name} failed to reach 'trigger' state in trigger"
-                )
-            # Status that resolves once the trigger is done
-            status_backend = self.backend.on_trigger_finished()
+            raise TimeoutError(
+                f"TimePixFly backend of device {self.name} failed to reach 'config' state in trigger"
+            )
 
-        if status_backend is not None:
-            return_status = status_camera & status_backend
-        else:
-            return_status = status_camera
-        self.cancel_on_stop(return_status)
+        # Registered before arming: the EndFrame of a lost acquisition can arrive at once.
+        acquisition = XesAcquisition(
+            exposure=self._exposure_per_trigger, repeat_allowed=repeat_allowed
+        )
+        with self._xes_acquisitions_lock:
+            self._xes_acquisitions.append(acquisition)
+        # Prepare backend to be ready to receive trigger
+        status = self.backend.on_trigger()
+        try:
+            status.wait(timeout=5)  # Wait until backend trigger is done
+        # pylint:disable=raise-missing-from
+        except Exception:
+            with self._xes_acquisitions_lock:
+                if acquisition in self._xes_acquisitions:
+                    self._xes_acquisitions.remove(acquisition)
+            logger.error(f"TimePixFly backend of device {self.name} failed to prepare for trigger")
+            raise TimeoutError(
+                f"TimePixFly backend of device {self.name} failed to reach 'trigger' state in trigger"
+            )
+        # Status that resolves once the trigger is done
+        status_backend = self.backend.on_trigger_finished()
+        acquisition_status = status_camera & status_backend
+        self.cancel_on_stop(acquisition_status)
+        acquisition_status.add_callback(
+            partial(self._on_xes_acquisition_done, point_status, acquisition)
+        )
+        acquisition.acquire_time = time.monotonic()
         self.cam.acquire.put(1)
-        return return_status
+
+    def _on_xes_acquisition_done(
+        self, point_status: DeviceStatus, acquisition: XesAcquisition, status: StatusBase
+    ) -> None:
+        """Resolve the trigger status, or repeat the acquisition if TimePixFly lost its data."""
+        with self._xes_acquisitions_lock:
+            acquisition.done = True
+            lost = acquisition.lost
+        if point_status.done:
+            return  # cancelled, e.g. by a stop
+        if not status.success:
+            self._resolve_point_status(point_status, status.exception())
+            return
+        if lost:
+            threading.Thread(
+                target=self._repeat_xes_acquisition,
+                args=(point_status,),
+                daemon=True,
+                name=f"{self.name}_repeat_acquisition",
+            ).start()
+            return
+        self._resolve_point_status(point_status)
+
+    def _repeat_xes_acquisition(self, point_status: DeviceStatus) -> None:
+        """Repeat a lost acquisition once, after the images of the lost one have arrived."""
+        try:
+            deadline = time.monotonic() + LOST_ACQUISITION_IMAGE_TIMEOUT
+            while True:
+                with self._xes_acquisitions_lock:
+                    if self._images_to_skip == 0:
+                        break
+                    if time.monotonic() > deadline:
+                        logger.warning(
+                            f"{self.name}: {self._images_to_skip} image(s) of the lost "
+                            "acquisition did not arrive; repeating without waiting for them."
+                        )
+                        self._images_to_skip = 0
+                        break
+                time.sleep(0.05)
+            if point_status.done or self.stopped:
+                return
+            self._start_xes_acquisition(point_status, repeat_allowed=False)
+        except Exception as exc:  # pylint: disable=broad-except
+            content = traceback.format_exc()
+            logger.error(f"{self.name}: repeating the lost acquisition failed: {content}")
+            self._resolve_point_status(point_status, exc)
+
+    @staticmethod
+    def _resolve_point_status(point_status: DeviceStatus, exc: Exception | None = None) -> None:
+        """Finish point_status (failed if exc is given), unless it is already done."""
+        if isinstance(exc, (StatusTimeoutError, WaitTimeoutError)):
+            exc = TimeoutError(str(exc))  # ophyd reserves these for its own timeouts
+        try:
+            if exc is None:
+                point_status.set_finished()
+            else:
+                point_status.set_exception(exc)
+        except InvalidState:
+            pass  # resolved meanwhile, e.g. cancelled by a stop
+
+    def _clear_xes_acquisitions(self) -> None:
+        """Forget pending acquisitions and images to skip (new scan, or stop)."""
+        with self._xes_acquisitions_lock:
+            self._xes_acquisitions.clear()
+            self._images_to_skip = 0
 
     def on_complete(self) -> DeviceStatus | StatusBase | None:
         """Called to inquire if a device has completed a scans."""
@@ -976,6 +1180,7 @@ class Timepix(PSIDeviceBase, TimePixControl):
         # Backend
         if self.enable_xes is True:
             self.backend.on_stop()
+        self._clear_xes_acquisitions()
         # BEC does not unstage the devices of an aborted scan, so undo the timepix_scan
         # overrides here. Runs after the backend stop, which depends on enable_xes.
         try:

tests/tests_devices/test_timepix.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import threading
+import time
 from contextlib import ExitStack
 from unittest import mock
 
@@ -14,6 +15,7 @@ from bec_lib.messages import ScanStatusMessage
 from bec_server.device_server.tests.utils import DMMock
 from bec_server.scan_server.scans.scan_base import ScanInfo as ScanServerScanInfo
 from ophyd import DeviceStatus, Staged, StatusBase
+from ophyd_devices.interfaces.base_classes.psi_device_base import DeviceStoppedError
 from ophyd_devices.tests.utils import MockPV, patch_dual_pvs
 
 from superxas_bec.devices.timepix.timepix import (
@@ -762,3 +764,155 @@ def test_timepix_on_prescan_returns_correct_status(timepix):
     status.wait(timeout=1)
     assert status.done is True
     assert status.success is True
+
+
+# --- Lost acquisitions: TimePixFly closes the frame before the exposure ends (S00437 point 1) ---
+
+
+def _wait_for(condition, timeout: float = 2.0):
+    """Poll until condition() is true (the repetition runs in its own thread)."""
+    deadline = time.monotonic() + timeout
+    while not condition():
+        assert time.monotonic() < deadline, "condition not met in time"
+        time.sleep(0.01)
+
+
+def _camera_acquires(dev):
+    """The camera runs one acquisition: AcquireBusy DONE -> ACQUIRING -> DONE."""
+    _force_signal_value(dev.cam.acquire_busy, ACQUIRESTATUS.ACQUIRING)
+    _force_signal_value(dev.cam.acquire_busy, ACQUIRESTATUS.DONE)
+
+
+def _empty_frames():
+    """Frames of an acquisition in which TimePixFly saw no TDC pulse (as S00437, point 1)."""
+    start_frame, data_frames, end_frame = _xes_frames(np.zeros((4, 8), dtype=np.float32))
+    data_frames[0]["period"] = 0
+    return start_frame, data_frames, {**end_frame, "periods": 0}
+
+
+@pytest.fixture
+def xes_timepix(timepix):
+    """Timepix armed for XES as in a timepix_scan: image writing off, 1 s exposure, 1 image."""
+    _force_signal_value(timepix.hdf.enable, "Disable")
+    timepix.troin = 4
+    timepix._exposure_per_trigger = 1.0
+    timepix._images_per_trigger = 1
+    backend_done = [StatusBase(), StatusBase()]
+    timepix.backend.on_trigger_finished = mock.Mock(side_effect=backend_done)
+    timepix.backend_done = backend_done
+    return timepix
+
+
+def test_timepix_lost_acquisition_is_repeated_once(xes_timepix):
+    """A frame closed before the exposure ends is dropped and the acquisition repeated."""
+    timepix = xes_timepix
+    with (
+        mock.patch.object(timepix.xes_data, "put") as xes_put,
+        mock.patch.object(timepix.static_spectra, "put") as spectra_put,
+        # a 2 x 2 detector image
+        mock.patch.object(timepix.image.array_size.height, "get", return_value=2),
+        mock.patch.object(timepix.image.array_size.width, "get", return_value=2),
+        mock.patch.object(timepix.image.array_data, "get", return_value=np.arange(4)),
+    ):
+        status = timepix.on_trigger()
+        assert timepix.cam.acquire.get() == 1
+
+        # TimePixFly closes the frame at once, without TDC periods: nothing is published
+        timepix.msg_buffer_callback(*_empty_frames())
+        xes_put.assert_not_called()
+        # the image of the lost acquisition is not published either
+        timepix._on_new_image_received(value=1, old_value=0)
+        spectra_put.assert_not_called()
+
+        _camera_acquires(timepix)
+        timepix.backend_done[0].set_finished()
+        assert status.done is False
+
+        # the repetition arms TimePixFly again and restarts the camera
+        _wait_for(lambda: timepix.backend.on_trigger_finished.call_count == 2)
+        timepix.msg_buffer_callback(*_xes_frames(_level_test_spectra()))
+        timepix._on_new_image_received(value=2, old_value=1)
+        _camera_acquires(timepix)
+        timepix.backend_done[1].set_finished()
+        status.wait(timeout=2)
+
+    assert status.success is True
+    xes_put.assert_called_once()
+    np.testing.assert_array_equal(xes_put.call_args.args[0], _level_test_spectra())
+    spectra_put.assert_called_once()
+    assert not timepix._xes_acquisitions
+
+
+def test_timepix_lost_repetition_is_published(xes_timepix):
+    """If the repetition is lost too, its empty frame is published: the point keeps one row."""
+    timepix = xes_timepix
+    with mock.patch.object(timepix.xes_data, "put") as xes_put:
+        status = timepix.on_trigger()
+        timepix.msg_buffer_callback(*_empty_frames())
+        timepix._on_new_image_received(value=1, old_value=0)
+        _camera_acquires(timepix)
+        timepix.backend_done[0].set_finished()
+        _wait_for(lambda: timepix.backend.on_trigger_finished.call_count == 2)
+
+        with mock.patch("superxas_bec.devices.timepix.timepix.logger") as logger:
+            timepix.msg_buffer_callback(*_empty_frames())
+        _camera_acquires(timepix)
+        timepix.backend_done[1].set_finished()
+        status.wait(timeout=2)
+
+    assert status.success is True
+    xes_put.assert_called_once()
+    assert "cannot be repeated" in logger.error.call_args.args[0]
+
+
+def test_timepix_frame_without_tdc_after_exposure_is_published(xes_timepix):
+    """No TDC pulse during a whole exposure is a detector/TDC problem: publish, do not repeat."""
+    timepix = xes_timepix
+    timepix._exposure_per_trigger = 0.0  # the exposure has ended when the frame arrives
+    with (
+        mock.patch.object(timepix.xes_data, "put") as xes_put,
+        mock.patch("superxas_bec.devices.timepix.timepix.logger") as logger,
+    ):
+        status = timepix.on_trigger()
+        timepix.msg_buffer_callback(*_empty_frames())
+        _camera_acquires(timepix)
+        timepix.backend_done[0].set_finished()
+        status.wait(timeout=1)
+
+    assert status.success is True
+    xes_put.assert_called_once()
+    assert timepix.backend.on_trigger_finished.call_count == 1
+    assert "no TDC pulse" in logger.error.call_args.args[0]
+
+
+def test_timepix_lost_acquisition_not_repeated_while_writing_images(xes_timepix):
+    """With HDF5 image writing on, a repetition would add an image: keep the point as it is."""
+    timepix = xes_timepix
+    _force_signal_value(timepix.hdf.enable, "Enable")
+    with mock.patch.object(timepix.xes_data, "put") as xes_put:
+        status = timepix.on_trigger()
+        timepix.msg_buffer_callback(*_empty_frames())
+        _camera_acquires(timepix)
+        timepix.backend_done[0].set_finished()
+        status.wait(timeout=1)
+
+    assert status.success is True
+    xes_put.assert_called_once()
+    assert timepix.backend.on_trigger_finished.call_count == 1
+
+
+def test_timepix_stop_cancels_repetition_of_lost_acquisition(xes_timepix):
+    """A stop while the lost acquisition is still running cancels the trigger, no repetition."""
+    timepix = xes_timepix
+    status = timepix.on_trigger()
+    timepix.msg_buffer_callback(*_empty_frames())
+
+    timepix.stop()
+
+    with pytest.raises(DeviceStoppedError):
+        status.wait(timeout=1)
+    _camera_acquires(timepix)
+    time.sleep(0.2)
+    assert timepix.backend.on_trigger_finished.call_count == 1
+    assert not timepix._xes_acquisitions
+    assert timepix._images_to_skip == 0