added number_of_frames instead of iterations (#581)

* Added more flexible multi-run specs

* Refactor StandardDetector kickoff method to track the number of completable frames

This commit refactors the kickoff method in the StandardDetector class to track the number of frames that can be completed after kickoff. The frames_completed attribute has been replaced with the completable_frames attribute, which represents the number of frames that can be completed. Additionally, the frames_completed attribute has been removed from the complete method and replaced with the completable_frames attribute. This change ensures that the completable_frames attribute is reset to 0 once the target number of iterations is reached and the detector is disarmed.

* Refactor test_hardware_triggered_flyable to use number_of_triggers instead of number_of_frames

* Refactor test_hardware_triggered_flyable to use number_of_triggers and invoke_extra_kickoff_before_complete

* Refactor kickoff method to track total number of frames completed

src/ophyd_async/core/_detector.py

@@ -3,8 +3,9 @@
 import asyncio
 import time
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence
 from enum import Enum
+from functools import cached_property
 from typing import (
     Generic,
 )
@@ -20,7 +21,7 @@
     WritesStreamAssets,
 )
 from event_model import DataKey
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, NonNegativeInt, computed_field
 
 from ._device import Device
 from ._protocol import AsyncConfigurable, AsyncReadable
@@ -45,8 +46,16 @@ class DetectorTrigger(str, Enum):
 class TriggerInfo(BaseModel):
     """Minimal set of information required to setup triggering on a detector"""
 
-    #: Number of triggers that will be sent, 0 means infinite
-    number: int = Field(ge=0)
+    #: Number of triggers that will be sent, (0 means infinite) Can be:
+    #  - A single integer or
+    #  - A list of integers for multiple triggers
+    # Example for tomography: TriggerInfo(number=[2,3,100,3])
+    #:     This would trigger:
+    #:     - 2 times for dark field images
+    #:     - 3 times for initial flat field images
+    #:     - 100 times for projections
+    #:     - 3 times for final flat field images
+    number_of_triggers: NonNegativeInt | list[NonNegativeInt]
     #: Sort of triggers that will be sent
     trigger: DetectorTrigger = Field(default=DetectorTrigger.internal)
     #: What is the minimum deadtime between triggers
@@ -60,10 +69,15 @@ class TriggerInfo(BaseModel):
     #: e.g. if num=10 and multiplier=5 then the detector will take 10 frames,
     #: but publish 2 indices, and describe() will show a shape of (5, h, w)
     multiplier: int = 1
-    #: The number of times the detector can go through a complete cycle of kickoff and
-    #: complete without needing to re-arm. This is important for detectors where the
-    #: process of arming is expensive in terms of time
-    iteration: int = 1
+
+    @computed_field
+    @cached_property
+    def total_number_of_triggers(self) -> int:
+        return (
+            sum(self.number_of_triggers)
+            if isinstance(self.number_of_triggers, list)
+            else self.number_of_triggers
+        )
 
 
 class DetectorControl(ABC):
@@ -192,10 +206,14 @@ def __init__(
         # For kickoff
         self._watchers: list[Callable] = []
         self._fly_status: WatchableAsyncStatus | None = None
-        self._fly_start: float
-        self._iterations_completed: int = 0
-        self._initial_frame: int
-        self._last_frame: int
+        self._fly_start: float | None = None
+        self._frames_to_complete: int = 0
+        # Represents the total number of frames that will have been completed at the
+        # end of the next `complete`.
+        self._completable_frames: int = 0
+        self._number_of_triggers_iter: Iterator[int] | None = None
+        self._initial_frame: int = 0
+
         super().__init__(name)
 
     @property
@@ -251,7 +269,7 @@ async def trigger(self) -> None:
         if self._trigger_info is None:
             await self.prepare(
                 TriggerInfo(
-                    number=1,
+                    number_of_triggers=1,
                     trigger=DetectorTrigger.internal,
                     deadtime=None,
                     livetime=None,
@@ -301,8 +319,12 @@ async def prepare(self, value: TriggerInfo) -> None:
                 f"but trigger logic provides only {value.deadtime}s"
             )
         self._trigger_info = value
+        self._number_of_triggers_iter = iter(
+            self._trigger_info.number_of_triggers
+            if isinstance(self._trigger_info.number_of_triggers, list)
+            else [self._trigger_info.number_of_triggers]
+        )
         self._initial_frame = await self.writer.get_indices_written()
-        self._last_frame = self._initial_frame + self._trigger_info.number
         self._describe, _ = await asyncio.gather(
             self.writer.open(value.multiplier), self.controller.prepare(value)
         )
@@ -312,39 +334,50 @@ async def prepare(self, value: TriggerInfo) -> None:
 
     @AsyncStatus.wrap
     async def kickoff(self):
-        assert self._trigger_info, "Prepare must be called before kickoff!"
-        if self._iterations_completed >= self._trigger_info.iteration:
-            raise Exception(
+        if self._trigger_info is None or self._number_of_triggers_iter is None:
+            raise RuntimeError("Prepare must be called before kickoff!")
+        try:
+            self._frames_to_complete = next(self._number_of_triggers_iter)
+            self._completable_frames += self._frames_to_complete
+        except StopIteration as err:
+            raise RuntimeError(
                 f"Kickoff called more than the configured number of "
-                f"{self._trigger_info.iteration} iteration(s)!"
-            )
-        self._iterations_completed += 1
+                f"{self._trigger_info.total_number_of_triggers} iteration(s)!"
+            ) from err
 
     @WatchableAsyncStatus.wrap
     async def complete(self):
         assert self._trigger_info
-        async for index in self.writer.observe_indices_written(
+        indices_written = self.writer.observe_indices_written(
             self._trigger_info.frame_timeout
             or (
                 DEFAULT_TIMEOUT
                 + (self._trigger_info.livetime or 0)
                 + (self._trigger_info.deadtime or 0)
             )
-        ):
-            yield WatcherUpdate(
-                name=self.name,
-                current=index,
-                initial=self._initial_frame,
-                target=self._trigger_info.number,
-                unit="",
-                precision=0,
-                time_elapsed=time.monotonic() - self._fly_start,
-            )
-            if index >= self._trigger_info.number:
-                break
-        if self._iterations_completed == self._trigger_info.iteration:
-            self._iterations_completed = 0
-            await self.controller.wait_for_idle()
+        )
+        try:
+            async for index in indices_written:
+                yield WatcherUpdate(
+                    name=self.name,
+                    current=index,
+                    initial=self._initial_frame,
+                    target=self._frames_to_complete,
+                    unit="",
+                    precision=0,
+                    time_elapsed=time.monotonic() - self._fly_start
+                    if self._fly_start
+                    else None,
+                )
+                if index >= self._frames_to_complete:
+                    break
+        finally:
+            await indices_written.aclose()
+            if self._completable_frames >= self._trigger_info.total_number_of_triggers:
+                self._completable_frames = 0
+                self._frames_to_complete = 0
+                self._number_of_triggers_iter = None
+                await self.controller.wait_for_idle()
 
     async def describe_collect(self) -> dict[str, DataKey]:
         return self._describe

src/ophyd_async/epics/adaravis/_aravis_controller.py

@@ -30,7 +30,7 @@ def get_deadtime(self, exposure: float | None) -> float:
         return _HIGHEST_POSSIBLE_DEADTIME
 
     async def prepare(self, trigger_info: TriggerInfo):
-        if (num := trigger_info.number) == 0:
+        if trigger_info.total_number_of_triggers == 0:
             image_mode = adcore.ImageMode.continuous
         else:
             image_mode = adcore.ImageMode.multiple
@@ -43,7 +43,7 @@ async def prepare(self, trigger_info: TriggerInfo):
 
         await asyncio.gather(
             self._drv.trigger_source.set(trigger_source),
-            self._drv.num_images.set(num),
+            self._drv.num_images.set(trigger_info.total_number_of_triggers),
             self._drv.image_mode.set(image_mode),
         )
 

src/ophyd_async/epics/adkinetix/_kinetix_controller.py

@@ -29,7 +29,7 @@ def get_deadtime(self, exposure: float | None) -> float:
     async def prepare(self, trigger_info: TriggerInfo):
         await asyncio.gather(
             self._drv.trigger_mode.set(KINETIX_TRIGGER_MODE_MAP[trigger_info.trigger]),
-            self._drv.num_images.set(trigger_info.number),
+            self._drv.num_images.set(trigger_info.total_number_of_triggers),
             self._drv.image_mode.set(adcore.ImageMode.multiple),
         )
         if trigger_info.livetime is not None and trigger_info.trigger not in [

src/ophyd_async/epics/adpilatus/_pilatus_controller.py

@@ -40,7 +40,9 @@ async def prepare(self, trigger_info: TriggerInfo):
         await asyncio.gather(
             self._drv.trigger_mode.set(self._get_trigger_mode(trigger_info.trigger)),
             self._drv.num_images.set(
-                999_999 if trigger_info.number == 0 else trigger_info.number
+                999_999
+                if trigger_info.total_number_of_triggers == 0
+                else trigger_info.total_number_of_triggers
             ),
             self._drv.image_mode.set(adcore.ImageMode.multiple),
         )

src/ophyd_async/epics/adsimdetector/_sim_controller.py

@@ -32,7 +32,7 @@ async def prepare(self, trigger_info: TriggerInfo):
             DEFAULT_TIMEOUT + await self.driver.acquire_time.get_value()
         )
         await asyncio.gather(
-            self.driver.num_images.set(trigger_info.number),
+            self.driver.num_images.set(trigger_info.total_number_of_triggers),
             self.driver.image_mode.set(adcore.ImageMode.multiple),
         )
 

src/ophyd_async/epics/advimba/_vimba_controller.py

@@ -37,7 +37,7 @@ async def prepare(self, trigger_info: TriggerInfo):
         await asyncio.gather(
             self._drv.trigger_mode.set(TRIGGER_MODE[trigger_info.trigger]),
             self._drv.exposure_mode.set(EXPOSE_OUT_MODE[trigger_info.trigger]),
-            self._drv.num_images.set(trigger_info.number),
+            self._drv.num_images.set(trigger_info.total_number_of_triggers),
             self._drv.image_mode.set(adcore.ImageMode.multiple),
         )
         if trigger_info.livetime is not None and trigger_info.trigger not in [

src/ophyd_async/epics/eiger/_eiger_controller.py

@@ -41,7 +41,7 @@ async def prepare(self, trigger_info: TriggerInfo):
             self._drv.trigger_mode.set(
                 EIGER_TRIGGER_MODE_MAP[trigger_info.trigger].value
             ),
-            self._drv.num_images.set(trigger_info.number),
+            self._drv.num_images.set(trigger_info.total_number_of_triggers),
         ]
         if trigger_info.livetime is not None:
             coros.extend(

src/ophyd_async/plan_stubs/_fly.py

@@ -44,7 +44,6 @@ def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
     repeats: int = 1,
     period: float = 0.0,
     frame_timeout: float | None = None,
-    iteration: int = 1,
 ):
     """Prepare a hardware triggered flyable and one or more detectors.
 
@@ -62,12 +61,11 @@ def prepare_static_seq_table_flyer_and_detectors_with_same_trigger(
     deadtime = max(det.controller.get_deadtime(exposure) for det in detectors)
 
     trigger_info = TriggerInfo(
-        number=number_of_frames * repeats,
+        number_of_triggers=number_of_frames * repeats,
         trigger=DetectorTrigger.constant_gate,
         deadtime=deadtime,
         livetime=exposure,
         frame_timeout=frame_timeout,
-        iteration=iteration,
     )
     trigger_time = number_of_frames * (exposure + deadtime)
     pre_delay = max(period - 2 * shutter_time - trigger_time, 0)

src/ophyd_async/sim/demo/_pattern_detector/_pattern_detector_controller.py

@@ -32,7 +32,9 @@ async def arm(self):
         assert self.period
         self.task = asyncio.create_task(
             self._coroutine_for_image_writing(
-                self._trigger_info.livetime, self.period, self._trigger_info.number
+                self._trigger_info.livetime,
+                self.period,
+                self._trigger_info.total_number_of_triggers,
             )
         )
 

system_tests/epics/eiger/test_eiger_system.py

@@ -71,7 +71,7 @@ async def test_eiger(RE, ioc_prefixes) -> EigerDetector:
 async def test_trigger_saves_file(test_eiger: EigerDetector, setup_device: SetupDevice):
     single_shot = EigerTriggerInfo(
         frame_timeout=None,
-        number=1,
+        number_of_triggers=1,
         trigger=DetectorTrigger.internal,
         deadtime=None,
         livetime=None,

tests/conftest.py

@@ -230,7 +230,7 @@ def static_path_provider(
 def one_shot_trigger_info() -> TriggerInfo:
     return TriggerInfo(
         frame_timeout=None,
-        number=1,
+        number_of_triggers=1,
         trigger=DetectorTrigger.internal,
         deadtime=None,
         livetime=None,