diff --git a/openhcs/core/aligned_image_payload.py b/openhcs/core/aligned_image_payload.py index 3d464b51c2..04c4be43c9 100644 --- a/openhcs/core/aligned_image_payload.py +++ b/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: " diff --git a/openhcs/core/callable_contract.py b/openhcs/core/callable_contract.py index a437eaea5c..d50d686e5b 100644 --- a/openhcs/core/callable_contract.py +++ b/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.""" diff --git a/tests/integration/test_main.py b/tests/integration/test_main.py index f7ea466a24..b726853983 100644 --- a/tests/integration/test_main.py +++ b/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") diff --git a/tests/unit/test_function_patterns.py b/tests/unit/test_function_patterns.py index fb4664921f..d7e5331d90 100644 --- a/tests/unit/test_function_patterns.py +++ b/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) diff --git a/tests/unit/test_opaque_main_flow_image_domain.py b/tests/unit/test_opaque_main_flow_image_domain.py index 77e68be9bc..238732519f 100644 --- a/tests/unit/test_opaque_main_flow_image_domain.py +++ b/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=(),