feat: add sequence to Position object + update MDASequence iter_sequence + test [WIP]

src/useq/_base_model.py

@@ -169,12 +169,16 @@ def yaml(
             YAML output ... If `stream` is provided, returns `None`.
         """
         from datetime import timedelta
+        from enum import Enum
 
         import yaml
 
         yaml.SafeDumper.add_multi_representer(
             timedelta, lambda dumper, data: dumper.represent_str(str(data))
         )
+        yaml.SafeDumper.add_multi_representer(
+            Enum, lambda dumper, data: dumper.represent_str(str(data.value))
+        )
 
         data = self.dict(
             include=include,

src/useq/_mda_event.py

@@ -140,6 +140,39 @@ def __repr_args__(self) -> ReprArgs:
         d.pop("sequence")
         return list(d.items())
 
+    def shifted(
+        self,
+        x_pos: float | None = None,
+        y_pos: float | None = None,
+        z_pos: float | None = None,
+    ) -> MDAEvent:
+        """Return a new event instance, shifted relative to this one."""
+        new_z = (
+            None
+            if z_pos is None and self.z_pos is None
+            else (z_pos or 0) + self.z_pos
+            if self.z_pos is not None
+            else z_pos
+        )
+
+        new_x = (
+            None
+            if x_pos is None and self.x_pos is None
+            else (x_pos or 0) + self.x_pos
+            if self.x_pos is not None
+            else x_pos
+        )
+
+        new_y = (
+            None
+            if y_pos is None and self.y_pos is None
+            else (y_pos or 0) + self.y_pos
+            if self.y_pos is not None
+            else y_pos
+        )
+
+        return self.copy(update={"x_pos": new_x, "y_pos": new_y, "z_pos": new_z})
+
     def to_pycromanager(self) -> dict:
         """Convenience method to convert this event to a pycro-manager events.
 

src/useq/_mda_sequence.py

@@ -1,13 +1,24 @@
 from __future__ import annotations
 
 from itertools import product
-from typing import Any, Dict, Iterator, Optional, Sequence, Tuple, Union, no_type_check
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+    no_type_check,
+)
 from uuid import UUID, uuid4
 from warnings import warn
 
 import numpy as np
 from pydantic import Field, PrivateAttr, root_validator, validator
 
+from . import _mda_event
 from ._base_model import UseqModel
 from ._channel import Channel
 from ._grid import AnyGridPlan, GridPosition, NoGrid
@@ -43,7 +54,7 @@ class MDASequence(UseqModel):
         The order of the axes in the sequence. Must be a permutation of `"tpgcz"`. The
         default is `"tpgcz"`.
     stage_positions : tuple[Position, ...]
-        The stage positions to visit. (each with `x`, `y`, `z`, `name`, and `z_plan`,
+        The stage positions to visit. (each with `x`, `y`, `z`, `name`, and `sequence`,
         all of which are optional).
     grid_plan : GridFromCorners, GridRelative, NoGrid
         The grid plan to follow. One of `GridFromCorners`, `GridRelative` or `NoGrid`.
@@ -81,13 +92,13 @@ class MDASequence(UseqModel):
     channels:
     - config: DAPI
       exposure: 1.0
+    grid_plan:
+      columns: 2
+      rows: 2
     stage_positions:
     - x: 1.0
       y: 1.0
       z: 1.0
-    grid_plan:
-      cols: 2
-      rows: 2
     time_plan:
       interval: '0:00:00.100000'
       loops: 2
@@ -135,7 +146,7 @@ def replace(
 
         MDASequences are immutable, so this method is useful for creating a new
         sequence with only a few fields changed.  The uid of the new sequence will
-        be different from the original
+        be different from the original.
         """
         kwargs = {
             k: v for k, v in locals().items() if v is not Undefined and k != "self"
@@ -186,6 +197,7 @@ def validate_mda(cls, values: Dict[str, Any]) -> Dict[str, Any]:
                 z_plan=values.get("z_plan"),
                 stage_positions=values.get("stage_positions", ()),
                 channels=values.get("channels", ()),
+                grid_plan=values.get("grid_plan"),
             )
         return values
 
@@ -202,13 +214,14 @@ def _check_order(
         z_plan: Optional[AnyZPlan] = None,
         stage_positions: Sequence[Position] = (),
         channels: Sequence[Channel] = (),
+        grid_plan: Optional[AnyGridPlan] = None,
     ) -> str:
         if (
             Z in order
             and POSITION in order
             and order.index(Z) < order.index(POSITION)
             and z_plan
-            and any(p.z_plan for p in stage_positions)
+            and any(p.sequence.z_plan for p in stage_positions if p.sequence)
         ):
             raise ValueError(
                 f"{Z!r} cannot precede {POSITION!r} in acquisition order if "
@@ -226,6 +239,31 @@ def _check_order(
                 "{TIME!r} in the acquisition order: may not yield intended results."
             )
 
+        if (
+            GRID in order
+            and POSITION in order
+            and grid_plan
+            and not grid_plan.is_relative
+            and len(stage_positions) > 1
+        ):
+            sub_position_grid_plans = [
+                p for p in stage_positions if p.sequence and p.sequence.grid_plan
+            ]
+            if len(stage_positions) - len(sub_position_grid_plans) > 1:
+                raise ValueError(
+                    "A MDASequence with an absloute grid_plan "
+                    "cannot have multiple stage positions!"
+                )
+
+        if (
+            POSITION in order
+            and stage_positions
+            and any(p.sequence.stage_positions for p in stage_positions if p.sequence)
+        ):
+            raise ValueError(
+                "Currently, a Position sequence cannot have multiple stage positions!"
+            )
+
         return order
 
     def __str__(self) -> str:
@@ -245,7 +283,6 @@ def shape(self) -> Tuple[int, ...]:
         """Return the shape of this sequence.
 
         !!! note
-
             This doesn't account for jagged arrays, like skipped Z or channel frames.
         """
         return tuple(s for s in self.sizes.values() if s)
@@ -321,13 +358,13 @@ def to_pycromanager(self) -> list[dict]:
 
 
 MDAEvent.update_forward_refs(MDASequence=MDASequence)
+Position.update_forward_refs(MDASequence=MDASequence)
 
 
 def iter_sequence(sequence: MDASequence) -> Iterator[MDAEvent]:
     """Iterate over all events in the MDA sequence.
 
     !!! note
-
         This method will usually be used via [`useq.MDASequence.iter_events`][], or by
         simply iterating over the sequence.
 
@@ -344,24 +381,52 @@ def iter_sequence(sequence: MDASequence) -> Iterator[MDAEvent]:
     MDAEvent
         Each event in the MDA sequence.
     """
-    order = sequence.used_axes
+    global_index = 0
 
-    event_iterator = (enumerate(sequence.iter_axis(ax)) for ax in order)
-    for global_index, item in enumerate(product(*event_iterator)):
+    event_iterator = (enumerate(sequence.iter_axis(ax)) for ax in sequence.used_axes)
+    for item in product(*event_iterator):
         if not item:  # the case with no events
             continue  # pragma: no cover
 
-        _ev = dict(zip(order, item))
+        _ev = dict(zip(sequence.used_axes, item))
         index = {k: _ev[k][0] for k in INDICES if k in _ev}
 
-        position: Optional[Position] = _ev[POSITION][1] if POSITION in _ev else None
-        channel: Optional[Channel] = _ev[CHANNEL][1] if CHANNEL in _ev else None
-        time: Optional[int] = _ev[TIME][1] if TIME in _ev else None
-        grid: Optional[GridPosition] = _ev[GRID][1] if GRID in _ev else None
+        position = cast(
+            "Position | None", _ev[POSITION][1] if POSITION in _ev else None
+        )
+        channel = cast("Channel | None", _ev[CHANNEL][1] if CHANNEL in _ev else None)
+        time = cast("int | None", _ev[TIME][1] if TIME in _ev else None)
+        grid = cast("GridPosition | None", _ev[GRID][1] if GRID in _ev else None)
 
         # skip channels
         if channel and TIME in index and index[TIME] % channel.acquire_every:
             continue
+        # skip if also in position.sequence
+        if position and position.sequence:
+            if CHANNEL in index and index[CHANNEL] != 0:
+                continue
+            if Z in index and index[Z] != 0 and position.sequence.z_plan:
+                continue
+            if GRID in index and index[GRID] != 0 and position.sequence.grid_plan:
+                continue
+
+        _channel = (
+            {"config": channel.config, "group": channel.group} if channel else None
+        )
+        _exposure = getattr(channel, "exposure", None)
+
+        pos_name = getattr(position, "name", None)
+
+        try:
+            z_pos = (
+                sequence._combine_z(_ev[Z][1], index[Z], channel, position)
+                if Z in _ev
+                else position.z
+                if position
+                else None
+            )
+        except sequence._SkipFrame:
+            continue
 
         if grid:
             x_pos: Optional[float] = grid.x
@@ -375,29 +440,59 @@ def iter_sequence(sequence: MDASequence) -> Iterator[MDAEvent]:
             x_pos = getattr(position, "x", None)
             y_pos = getattr(position, "y", None)
 
-        try:
-            z_pos = (
-                sequence._combine_z(_ev[Z][1], index[Z], channel, position)
-                if Z in _ev
-                else position.z
-                if position
-                else None
-            )
-        except sequence._SkipFrame:
+        if position and position.sequence:
+            for sub_event in iter_sequence(position.sequence):
+                update: dict[str, Any] = {
+                    "global_index": global_index,
+                    "index": {**index, **sub_event.index},
+                    "sequence": sequence,
+                    "pos_name": position.name or pos_name,
+                }
+
+                if not sub_event.channel and channel:
+                    update["channel"] = _mda_event.Channel(
+                        config=channel.config, group=channel.group
+                    )
+
+                if sub_event.exposure is None:
+                    update["exposure"] = _exposure
+
+                if sub_event.min_start_time is None:
+                    update["min_start_time"] = time
+
+                # z_plan
+                if position.sequence.z_plan and position.sequence.z_plan.is_relative:
+                    sub_event = sub_event.shifted(z_pos=position.z)
+                elif not position.sequence.z_plan:
+                    update["z_pos"] = z_pos
+
+                # grid_plan
+                if (
+                    position.sequence.grid_plan
+                    and position.sequence.grid_plan.is_relative
+                ):
+                    sub_event = sub_event.shifted(x_pos=position.x, y_pos=position.y)
+
+                elif not position.sequence.grid_plan:
+                    update["x_pos"] = x_pos
+                    update["y_pos"] = y_pos
+
+                yield sub_event.copy(update=update)
+
+                global_index += 1
+
             continue
 
-        _channel = (
-            {"config": channel.config, "group": channel.group} if channel else None
-        )
         yield MDAEvent(
             index=index,
             min_start_time=time,
-            pos_name=getattr(position, "name", None),
+            pos_name=pos_name,
             x_pos=x_pos,
             y_pos=y_pos,
             z_pos=z_pos,
-            exposure=getattr(channel, "exposure", None),
+            exposure=_exposure,
             channel=_channel,
             sequence=sequence,
             global_index=global_index,
         )
+        global_index += 1

src/useq/_position.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Generator, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Generator, Optional
 
 import numpy as np
-from pydantic import Field
 
 from ._base_model import FrozenModel
-from ._grid import GridRelative, NoGrid
-from ._z import AnyZPlan, NoZ
+
+if TYPE_CHECKING:
+    from useq import MDASequence
 
 
 class Position(FrozenModel):
@@ -27,21 +27,16 @@ class Position(FrozenModel):
         Z position in microns.
     name : str | None
         Optional name for the position.
-    z_plan : ZTopBottom | ZRangeAround | ZAboveBelow | ZRelativePositions | \
-        ZAbsolutePositions | NoZ | None
-        Z plan to execute at this position specifically. By default, [`NoZ`][useq.NoZ].
-    grid_plan : GridRelative, NoGrid
-        GridRelative plan execute at this position specifically.
-        By default, [`NoGrid`][useq.NoGrid].
+    sequence : MDASequence | None
+        Optional MDASequence relative this position.
     """
 
     # if None, implies 'do not move this axis'
     x: Optional[float] = None
     y: Optional[float] = None
     z: Optional[float] = None
     name: Optional[str] = None
-    z_plan: AnyZPlan = Field(default_factory=NoZ)
-    grid_plan: Union[GridRelative, NoGrid] = Field(default_factory=NoGrid)
+    sequence: Optional[MDASequence] = None
 
     @classmethod
     def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:

tests/conftest.py

@@ -15,7 +15,10 @@ def mda1() -> MDASequence:
                 "y": 20,
                 "z": 50,
                 "name": "test_name",
-                "z_plan": {"above": 10, "below": 0, "step": 1},
+                "sequence": MDASequence(
+                    z_plan={"above": 10, "below": 0, "step": 1},
+                    grid_plan={"rows": 2, "columns": 3},
+                ),
             },
         ],
         channels=[