fix: fix skip channel in iter_sequence + tests

src/useq/_mda_sequence.py

@@ -401,7 +401,22 @@ def iter_sequence(sequence: MDASequence) -> Iterator[MDAEvent]:  # noqa: C901
             continue
         # skip if also in position.sequence
         if position and position.sequence:
-            if CHANNEL in index and index[CHANNEL] != 0:
+            if (
+                # if a position specifies channels, then the *global* channel index
+                # is no longer relevant... so we skip all but the first "global" channel
+                CHANNEL in index
+                and index[CHANNEL] != 0
+                # UNLESS the position specifies any other plan.
+                # NOTE: if we ever add more plans, they will need to be explicitly added
+                # https://github.com/pymmcore-plus/useq-schema/pull/85
+                and not any(
+                    (
+                        position.sequence.grid_plan,
+                        position.sequence.z_plan,
+                        position.sequence.time_plan,
+                    )
+                )
+            ):
                 continue
             if Z in index and index[Z] != 0 and position.sequence.z_plan:
                 continue

tests/test_position_sequence.py

@@ -557,6 +557,7 @@ def test_order():
         ],
         z_plan={"range": 2, "step": 1},
     )
+
     assert [(i.global_index, i.index, i.z_pos, i.channel.config) for i in mda] == [
         (0, {"p": 0, "c": 0, "z": 0}, -1.0, "FITC"),
         (1, {"p": 0, "c": 0, "z": 1}, 0.0, "FITC"),
@@ -574,3 +575,104 @@ def test_order():
         (10, {"p": 1, "c": 0, "z": 2}, 51.0, "488"),
         (11, {"p": 1, "c": 1, "z": 2}, 51.0, "561"),
     ]
+
+
+def test_channels_and_pos_grid_plan():
+    # test that all channels are acquired for each grid position
+    mda = MDASequence(
+        axis_order="tpgcz",
+        channels=[
+            {"config": "Cy5", "exposure": 10},
+            {"config": "FITC", "exposure": 10},
+        ],
+        stage_positions=[
+            {
+                "x": 0,
+                "y": 0,
+                "sequence": {"grid_plan": {"rows": 2, "columns": 1}},
+            },
+        ],
+    )
+
+    assert [
+        (i.global_index, i.index, i.x_pos, i.y_pos, i.channel.config) for i in mda
+    ] == [
+        (0, {"p": 0, "c": 0, "g": 0}, 0.0, 0.5, "Cy5"),
+        (1, {"p": 0, "c": 0, "g": 1}, 0.0, -0.5, "Cy5"),
+        (2, {"p": 0, "c": 1, "g": 0}, 0.0, 0.5, "FITC"),
+        (3, {"p": 0, "c": 1, "g": 1}, 0.0, -0.5, "FITC"),
+    ]
+
+
+def test_channels_and_pos_z_plan():
+    # test that all channels are acquired for each z position
+    mda = MDASequence(
+        axis_order="tpgcz",
+        channels=[
+            {"config": "Cy5", "exposure": 10},
+            {"config": "FITC", "exposure": 10},
+        ],
+        stage_positions=[
+            {"x": 0, "y": 0, "z": 0, "sequence": {"z_plan": {"range": 2, "step": 1}}}
+        ],
+    )
+
+    assert [(i.global_index, i.index, i.z_pos, i.channel.config) for i in mda] == [
+        (0, {"p": 0, "c": 0, "z": 0}, -1.0, "Cy5"),
+        (1, {"p": 0, "c": 0, "z": 1}, 0.0, "Cy5"),
+        (2, {"p": 0, "c": 0, "z": 2}, 1.0, "Cy5"),
+        (3, {"p": 0, "c": 1, "z": 0}, -1.0, "FITC"),
+        (4, {"p": 0, "c": 1, "z": 1}, 0.0, "FITC"),
+        (5, {"p": 0, "c": 1, "z": 2}, 1.0, "FITC"),
+    ]
+
+
+def test_channels_and_pos_time_plan():
+    # test that all channels are acquired for each timepoint
+    mda = MDASequence(
+        axis_order="tpgcz",
+        channels=[
+            {"config": "Cy5", "exposure": 10},
+            {"config": "FITC", "exposure": 10},
+        ],
+        stage_positions=[
+            {"x": 0, "y": 0, "sequence": {"time_plan": [{"interval": 1, "loops": 3}]}},
+        ],
+    )
+
+    assert [
+        (i.global_index, i.index, i.min_start_time, i.channel.config) for i in mda
+    ] == [
+        (0, {"p": 0, "c": 0, "t": 0}, 0.0, "Cy5"),
+        (1, {"p": 0, "c": 0, "t": 1}, 1.0, "Cy5"),
+        (2, {"p": 0, "c": 0, "t": 2}, 2.0, "Cy5"),
+        (3, {"p": 0, "c": 1, "t": 0}, 0.0, "FITC"),
+        (4, {"p": 0, "c": 1, "t": 1}, 1.0, "FITC"),
+        (5, {"p": 0, "c": 1, "t": 2}, 2.0, "FITC"),
+    ]
+
+
+def test_channels_and_pos_z_grid_and_time_plan():
+    # test that all channels are acquired for each z and grid positions
+    mda = MDASequence(
+        axis_order="tpgcz",
+        channels=[
+            {"config": "Cy5", "exposure": 10},
+            {"config": "FITC", "exposure": 10},
+        ],
+        stage_positions=[
+            {
+                "x": 0,
+                "y": 0,
+                "z": 0,
+                "sequence": {
+                    "z_plan": {"range": 2, "step": 1},
+                    "grid_plan": {"rows": 2, "columns": 1},
+                    "time_plan": [{"interval": 1, "loops": 2}],
+                },
+            }
+        ],
+    )
+
+    chs = {i.channel.config for i in mda}
+    assert chs == {"Cy5", "FITC"}