added number_of_frames instead of iterations

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,13 @@ 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 number of frames that can be completed after kickoff
+        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 +268,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 +318,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)
         )
@@ -313,38 +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(
+        try:
+            if self._number_of_triggers_iter is None:
+                raise StopIteration
+            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,