B1: stack every z-plane before 3D processing

openhcs/core/aligned_image_payload.py

@@ -425,6 +425,26 @@ def validate_main_flow_cohort(
                 "One main-flow cohort cannot combine different declared image axes."
             )
 
+    @staticmethod
+    def keeps_single_member_whole(
+        payload: RuntimeArrayData,
+        *,
+        declared_plane_axis: RuntimePlaneAxis | None,
+    ) -> bool:
+        """Whether a one-member cohort is passed on without a new leading axis.
+
+        Only a member that already spans the cohort axis stays whole: a
+        persisted whole image (volume), or a member saved with its declared
+        plane axis. A member that declares no plane axis is composed on the
+        runtime-slice axis exactly as it is in a cohort of several, so the
+        rank a callable receives never depends on how many files matched.
+        """
+        metadata = image_payload_metadata(payload)
+        return metadata.persists_whole_image() or (
+            declared_plane_axis is not None
+            and declared_plane_axis is metadata.plane_axis
+        )
+
     @staticmethod
     def from_loaded_images(
         payloads: Sequence[RuntimeArrayData],
@@ -435,14 +455,13 @@ def from_loaded_images(
         workspace_source_lookups: Sequence[VirtualWorkspacePathLookup],
     ) -> RuntimeArrayData:
         """Compose a selected admissible input cohort in its declared image domain."""
-        if len(payloads) == 1 and (
-            image_payload_metadata(payloads[0]).persists_whole_image()
-            or (
-                producer_records
-                and len(producer_records) == 1
-                and producer_records[0].main_flow_plane_axis
-                is image_payload_metadata(payloads[0]).plane_axis
-            )
+        if len(payloads) == 1 and ImagePayloadStackComposition.keeps_single_member_whole(
+            payloads[0],
+            declared_plane_axis=(
+                producer_records[0].main_flow_plane_axis
+                if producer_records and len(producer_records) == 1
+                else None
+            ),
         ):
             main_data_stack = ImagePayloadStackComposition.copy_whole_image(
                 payloads[0],
@@ -529,9 +548,17 @@ def with_saved_output_context(
             and single_output_plane_axis is metadata[0].plane_axis
             and np.shape(data) == np.shape(image_payload_data(payloads[0]))
         ):
-            return metadata[0].with_current_intensity_from(current_intensity).payload_with(
+            member = metadata[0].with_current_intensity_from(current_intensity).payload_with(
                 data, image_payload_mask(stack_payload),
             )
+            if ImagePayloadStackComposition.keeps_single_member_whole(
+                member, declared_plane_axis=single_output_plane_axis,
+            ):
+                return member
+            # Cache the cohort the next step composes from this one saved member.
+            return stack_image_payloads(
+                (member,), metadata_mode=ImagePayloadMetadataCompositionMode.STACK,
+            )
         if np.shape(data)[:1] != (len(payloads),):
             raise ValueError(
                 "Output stack must match its declared output slice count: "

openhcs/core/callable_contract.py

@@ -1200,6 +1200,8 @@ def group_scope_inputs(self) -> ArtifactSpecCollection:
         input_specs = self.artifact_inputs
         source_refs = self.output_group_scope_sources
         if not source_refs:
+            if self.main_flow_supplies_primary_image:
+                return ArtifactSpecCollection(())
             return input_specs
         sources = ArtifactSpecCollection(
             source_spec
@@ -1219,6 +1221,18 @@ def group_scope_inputs(self) -> ArtifactSpecCollection:
             )
         return sources
 
+    @property
+    def main_flow_supplies_primary_image(self) -> bool:
+        """Whether the runtime passes main-flow pixels as the primary argument.
+
+        Such a callable runs once per main-flow cohort; its artifact inputs are
+        context for that cohort, not the owner of its group scope.
+        """
+        adapter = self.runtime_adapter
+        return self.accepts_implicit_main_flow_input and not (
+            adapter is not None and adapter.manages_artifact_inputs
+        )
+
     def preserves_input_main_flow(self) -> bool:
         """Return whether declared artifact outputs leave main flow unchanged."""
 

tests/integration/test_main.py

@@ -390,6 +390,11 @@ def _get_materialization_subdir() -> str:
 def _validate_subdirectory_fields(metadata: Dict) -> None:
     """Validate required fields in each subdirectory metadata."""
     materialization_subdir = _get_materialization_subdir()
+    # A saved-result destination is published with an empty raster inventory.
+    result_subdirs = {
+        subdir_metadata.get("results_dir")
+        for subdir_metadata in metadata[CONSTANTS.SUBDIRECTORIES_FIELD].values()
+    }
 
     for subdir_name, subdir_metadata in metadata[
         CONSTANTS.SUBDIRECTORIES_FIELD
@@ -402,10 +407,11 @@ def _validate_subdirectory_fields(metadata: Dict) -> None:
                 f"Subdirectory '{subdir_name}' missing fields: {missing_fields}"
             )
 
-        # Validate image_files (allow empty for materialization subdirectory)
+        # Validate image_files (allow empty for materialization and result subdirectories)
         if (
             not subdir_metadata.get("image_files")
             and subdir_name != materialization_subdir
+            and subdir_name not in result_subdirs
         ):
             raise ValueError(f"Subdirectory '{subdir_name}' has empty image_files list")
 

tests/unit/test_function_patterns.py

@@ -600,6 +600,47 @@ def apply_illumination(image, *, runtime):
     assert group.runtime_domain is RuntimeInvocationDomain.ARTIFACT_MANAGED
 
 
+def test_stored_context_input_beside_main_flow_image_keeps_every_anchor():
+    """A plain callable fed main-flow pixels runs once per anchor (B1).
+
+    Stitching consumed stored ``positions`` beside its image tiles; treating
+    that input as the group-scope owner kept one anchor, so only the first
+    z-plane was ever assembled.
+    """
+    positions = ArtifactSpec.input("positions", SpecialArtifactType)
+
+    @artifact_inputs(positions)
+    def assemble(image_tiles, positions):
+        del positions
+        return image_tiles
+
+    positions_plan = ArtifactInputPlan(
+        name="positions",
+        path="/memory/positions.pkl",
+        artifact_type=SpecialArtifactType,
+    )
+    compiled = compile_function_pattern(
+        assemble, {positions_plan.ref(): positions_plan}, {},
+    )
+    invocation = compiled.default_group.invocations[0]
+    invocation = invocation.with_artifact_input_edges((
+        exact_input_edge(
+            invocation,
+            input_index=0,
+            spec=positions,
+            storage_plan=positions_plan,
+            parameter_name="positions",
+        ),
+    ))
+    group = replace(compiled.default_group, invocations=(invocation,))
+
+    assert invocation.contract.main_flow_supplies_primary_image
+    assert invocation.contract.group_scope_inputs == ArtifactSpecCollection(())
+    assert group.runtime_domain is RuntimeInvocationDomain.SOURCE_ANCHORED
+    anchors = ["z001", "z002", "z003"]
+    assert group.runtime_domain.select_lifecycle_anchors(anchors) == anchors
+
+
 def test_source_bound_input_edge_keeps_source_anchored_runtime_domain():
     source_spec = ArtifactSpec.input("source", ImageArtifactType)
 

tests/unit/test_opaque_main_flow_image_domain.py

@@ -220,10 +220,14 @@ def test_save_and_next_load_preserve_domain_and_independent_cache(tmp_path, monk
     )
     monkeypatch.setattr(PatternGroupExecutionRequest, "source_workspace_projection_authority", lambda self: SimpleNamespace(projection_if_available=lambda: None))
     loaded = consumer.load_input_stack()[1]
-    assert image_payload_data(loaded).shape == pixels.shape
-    assert image_payload_metadata(loaded).plane_axis is axis
-    np.testing.assert_array_equal(image_payload_data(loaded), pixels)
-    np.testing.assert_array_equal(image_payload_mask(loaded), mask)
+    # A member saved without a plane axis is a one-member runtime-slice cohort,
+    # exactly as it is when several members match.
+    expected_pixels = pixels if axis is not None else pixels[None]
+    expected_mask = mask if axis is not None or mask is None else mask[None]
+    assert image_payload_data(loaded).shape == expected_pixels.shape
+    assert image_payload_metadata(loaded).plane_axis is (axis or RuntimePlaneAxis.RUNTIME_SLICE)
+    np.testing.assert_array_equal(image_payload_data(loaded), expected_pixels)
+    np.testing.assert_array_equal(image_payload_mask(loaded), expected_mask)
     assert not np.shares_memory(image_payload_data(loaded), pixels)
     if mask is not None:
         assert not np.shares_memory(image_payload_mask(loaded), mask)
@@ -379,7 +383,7 @@ def consumer_for(compiled):
 
     with pytest.raises(ValueError, match="different declared image axes"):
         consumer_for(pattern).load_input_stack()
-    for name, expected, axis in (("Planes", image_payload_data(planes), RuntimePlaneAxis.RUNTIME_SLICE), ("Volume", opaque_pixels, None)):
+    for name, expected, axis in (("Planes", image_payload_data(planes), RuntimePlaneAxis.RUNTIME_SLICE), ("Volume", opaque_pixels[None], RuntimePlaneAxis.RUNTIME_SLICE)):
         selected = consumer_for(_selected_pattern(name))
         actual = selected.load_input_stack()[1]
         np.testing.assert_array_equal(image_payload_data(actual), expected)
@@ -420,7 +424,9 @@ def test_opaque_named_bundle_uses_existing_same_slice_mask_composition():
     np.testing.assert_array_equal(image_payload_mask(stack_payload), image_payload_mask(expected))
 
 
-def test_scalar_rgb_occurrences_reload_on_runtime_slice_axis(tmp_path):
+@pytest.mark.parametrize("sites", [(1,), (1, 2, 3)], ids=("one-member", "three-members"))
+def test_scalar_rgb_occurrences_reload_on_runtime_slice_axis(tmp_path, sites):
+    """A cohort's rank does not depend on how many members matched (B1)."""
     from openhcs.core.aligned_image_payload import ImagePayloadStackComposition
     from openhcs.core.steps.function_output_manifest import ProducedOutputSemantics
     from openhcs.core.steps.function_output_identity import FunctionOutputIdentity
@@ -434,10 +440,10 @@ def test_scalar_rgb_occurrences_reload_on_runtime_slice_axis(tmp_path):
             extension=".tif", source="test",
         ),
         output_context=context, main_flow_plane_axis=None,
-    ) for site in (1, 2, 3))
+    ) for site in sites)
     payloads = tuple(ImagePayloadMetadata(
         source_image_names=("RGBImage",), source_channel_axis=-1,
-    ).payload_with(np.full((4, 5, 3), site, dtype=np.float32)) for site in (1, 2, 3))
+    ).payload_with(np.full((4, 5, 3), site, dtype=np.float32)) for site in sites)
     output = ImagePayloadStackComposition.from_loaded_images(
         payloads, producer_records=records, execution_plan=plan,
         source_projection=None, workspace_source_lookups=(),