diff --git a/docs/refactor/consolidation-20261009/R1e-no-jit-after-readiness.md b/docs/refactor/consolidation-20261009/R1e-no-jit-after-readiness.md new file mode 100644 index 0000000000..f9e02b4a08 --- /dev/null +++ b/docs/refactor/consolidation-20261009/R1e-no-jit-after-readiness.md @@ -0,0 +1,226 @@ +# R1e: nothing is JIT-compiled after readiness + +Branch `refactor/R1e-no-jit-after-readiness`. Follows R1d (#1225). + +## Owner requirement + +Every one-time cost is paid during function-catalog preparation, before the server +reports ready. After readiness no process (the server, its execution workers, the client +that converts and submits a pipeline) compiles a Numba kernel, and pipeline compilation +contains no one-time cost, because benchmarks time compilation and execution. +Compilation stays in preparation; a kernel found compiled late is moved into +preparation, never to first use. + +## How it was measured + +A `sitecustomize` module on `PYTHONPATH`, enabled by an environment variable, with no +product change: it registers a listener for Numba's `numba:compile` event (fired only +for a real compilation) and wraps `numba.core.caching.Cache.load_overload` (a disk-cache +lookup, hit or miss), and logs pid, role, function, signature and seconds per process. +The server's readiness and each request's compile and execute windows are marked by +wrapping `ZMQExecutionServer.prepare_runtime_capabilities`, +`ZMQExecutionServer._execute_pipeline` and `ZMQCompilationRequest.resolve`. Imports are +timed by wrapping each module loader's `exec_module`. Every run starts from an empty +`NUMBA_CACHE_DIR` and an empty catalog cache (`XDG_DATA_HOME`) with `OPENHCS_CPU_ONLY=1`. +The scratch version lives in `.claude/` (not committed); the counting part is committed +as `tests/integration/helpers/numba_trace/sitecustomize.py` for the guard. + +## What is wrong (measured on main `557eb99b4`, Linux, Ryzen 5 5500, 12 threads) + +Each row is one cold `run-well-throughput` run of `official30_portable_axis1.json` +(last row: `tests/integration/test_main.py`, ImageXpress 3D over ZMQ). + +| Pipeline | Compiled after readiness | Disk-cache loads after readiness | +| --- | --- | --- | +| ExampleHuman, 1 worker | 0 in server, 12 workers, client | client 107 kernels, 0.74 s | +| ExampleHuman, 4 workers | 0 | client 107, 0.79 s | +| ExampleFly | 7 in the server process during execution, 0.91 s | client 178, 0.69 s; with 2 workers each worker loads the late kernel, 0.01 s | +| ExampleColocalization | 7 in the server process during execution, 0.77 s | client 92, 0.96 s | +| ExampleCometAssay | 0 | client 116, 0.50 s | +| cp_tutorial_3d_monolayer | 0 | client 112, 0.53 s | +| ExampleSpeckles | 0 | client 60, 0.59 s | +| stock OpenHCS pipeline (not CellProfiler) | 0 | 0 in every process | + +1. **A kernel no preparation reaches.** The 7 compilations are + `_linear_costes_sorted_events_numba`, its 5 callees and Numba's `max`. + `NumbaNumpyColocalizationCostesBackendStrategy.linear_costes` branches on the data + (`slope > 0`), and its `prepare_backend` passed only anti-correlated arrays, so the + other branch was first compiled inside timed execution. A warm-up that calls a + public method covers one data-dependent path. +2. **It is not the only one.** Asking every first-party `cache=True` dispatcher, after a + full preparation, whether it has any compiled signature in memory or in Numba's disk + index: 52 of 298 had none (8 more are `inline="always"` and never have one), in 22 + modules. Some are dead (never called); the rest are branches, 3D variants or dtype + variants no warm-up reaches, and each would compile at first use. +3. **The client prepares kernels it never runs.** Converting a `.cppipe` built each + callable's contract with `CallableContract.from_prepared_callable`, and so did + `normalize_function_pattern`. In a process that has not run catalog preparation + that runs every preparation hook of the callable: 60 to 178 kernels loaded from the + disk cache (0.5 to 1.0 s), or compiled if the client gets there before the server + has filled the cache. `run-well-throughput` does this in workspace preparation, + before its pipeline clock; the GUI's pipeline editor calls the same function. + Preparation changes no signature: computed before and after a full preparation, the + canonical and runtime signatures of every catalogued callable are identical. +4. **Workers are forked before readiness and inherit what the server had loaded**, so + R1c's observation ("each of 4 workers compiled the same 2 thresholding kernels") + does not reproduce on main: the 12 workers are forked 0.25 s before readiness and + compile and load nothing. They do not inherit what the server first imports later. +5. **Late imports.** After readiness the server first-imports 44 modules setting up the + first request (0.06 s), 2 during compilation (0.01 s) and 60 during first execution + (0.04 s); every worker repeats them on its first task (45 to 102 modules, 0.09 to + 0.13 s per worker). The first pipeline compile in a server session took 0.30 s and + the second 0.28 s (ExampleHuman; 0.35 s and 0.44 s for ExampleFly, so within noise). + +Not measurable here: pyclesperanto and CuPy program builds (no GPU on this machine). +`matplotlib.font_manager` did not appear among the late imports. + +## Done + +- **Kernels no preparation reached.** Each of the 52, plus the 6 Costes kernels the + first measurement had already compiled late, was either deleted because nothing calls + it or is now reached from its owner's existing preparation, through the public method + real execution uses and with the dtypes and layouts its callers pass. No list of + kernel names was added and no kernel body changed. + - Covered: linear Costes with positively correlated data (`slope > 0`); alignment + with a partial mask; granularity with `subsample_size=0.25` (the default) for + float32 and float64; dense 3D object intensity (labels filling the volume); sparse + parent/child relationships in the four layout pairs callers produce; `filter_objects` + "parent with most overlap" (new `@processing_prepare(filter_objects)`); scikit-image + region orientations (new `prepare_backend`, strided and contiguous columns); + grayscale closing/opening for float32 and float64, writable and read-only; the Numba + smoothing provider, masked and unmasked; weighted Otsu with repeated values (the + counted branch, 7 kernels); the minimum threshold; the worm thinning kernel (new + `@processing_prepare` for the three `untangle_worms` callables). + - `label_geometry._numpy_124_svml_arccos` (added by #1224) exists only where the CPU + has the SVML symbol; `prepare_arccos_kernels`, called by the size/shape kernel + preparation, compiles the kernel this CPU uses and the portable one (a process + started with AVX-512 disabled shares the cache and uses it). + - Deleted: see certificates. +- **Authoring prepares nothing.** `CallableProjection.declared_contract()` reads the + contract from the declaration without running a hook; `.cppipe` conversion + (`module_artifact_contracts`, `module_artifact_declarations`, `primary_objects`) and + `NormalizedFunctionGroup.from_chain` use it. The compiler's entry is + `normalize_prepared_function_pattern` (used by `CompilationSession`): it prepares + every enabled callable and then lowers, so a process with no catalog preparation (an + in-process orchestrator) still prepares at compile time, before execution, as before. + The client now compiles 0 and loads 0 kernels. It still imports the modules of the + callables it converts (1.0 to 1.6 s, before the benchmark's clock): it needs the + callable objects to build the pipeline it submits, and a separate process cannot + inherit them. Published signatures are not sent from the server: the client reads the + same signature from the annotations, which the census showed is identical. +- **Late imports.** `FunctionCatalogExecutionServer.prepare_runtime_capabilities` now + also builds the storage backend registry (`polystore.base.ensure_storage_registry`) + and completes the domain-extended family registries + (`dataset_sources.discovery.discover_domain_registries`, which follows the registries + as they are created; no module list), before the worker fork. Function-local imports + with no cycle in `zmq_execution_signature`, `zmq_orchestrator_environment` and + `zmq_execution_server` became module-level. + +## Measurements (same machine; other agents' jobs were running, so wall times are noisy; counts are exact) + +| | main | this branch | +| --- | --- | --- | +| Compilations after readiness, ExampleFly / ExampleColocalization | 7 / 7 (0.91 s / 0.77 s, in execution) | 0 / 0 | +| Compilations after readiness, ExampleHuman 1 and 4 workers, CometAssay, 3D monolayer | 0 | 0 | +| Kernel loads in the client | 60 to 178 (0.5 to 1.0 s) | 0 | +| Kernel loads in workers (ExampleFly, 2 workers) | 1 per worker | 0 | +| First-party kernels with no compiled signature after preparation | 52 of 298 | 0 of 278 (13 of them are `inline="always"` and have none by design) | +| Server imports after readiness: request setup / compile / execute (ExampleHuman) | 44 (0.059 s) / 2 (0.010 s) / 60 (0.043 s) | 0 / 2 (0.002 s) / 57 (0.029 s) | +| Imports per worker on its first task | 45 to 102 (0.09 to 0.13 s; 0.44 s when PolyStore's registry cache misses) | 1 to 58 (0.002 to 0.034 s) | +| Pipeline compile, first / second in one server session, ExampleHuman | 0.302 / 0.282 s | 0.181 / 0.176 s | +| Same, ExampleFly | 0.350 / 0.436 s | 0.296 / 0.273 s | +| Cold preparation, traced, six runs | 48.7 to 55.8 s, mean 52.3 (one 72.4 s outlier excluded) | 47.3 to 53.3 s, mean 51.0 | + +Cold preparation is not measurably slower: the added warm-ups cost about 40 CPU-seconds +of compilation when timed alone and cold, spread over the cache children, and 17 dead +kernels fewer are compiled. The before/after compile times were taken hours apart under +different load; only "first equals second" within a session is a claim. + +## Guards + +`tests/integration/test_no_jit_after_readiness.py` (integration tier: a cold preparation +is the thing under test, about 90 s here; counts only, no wall-clock). One cold run, in a +child process, of the stock pipeline (`test_main.py`, ImageXpress 3D) and an imported +`.cppipe`, each over a real execution server, with +`tests/integration/helpers/numba_trace/sitecustomize.py` on `PYTHONPATH`. + +- `test_no_kernel_is_compiled_or_loaded_after_readiness`: preparation compiled at least + one kernel (so the counter works), and no process compiled or loaded a kernel after + its server's readiness; a client may not at any time. +- `test_preparation_compiles_every_kernel`: derived from the code, not from a pipeline. + At readiness every `openhcs` Numba dispatcher that is not `inline="always"` has a + signature in memory or in Numba's disk index. Seen failing: with the worm hook's call + removed it names `worm_geometry._skeletonize_loop_numba`. + +What they cannot see: a kernel that has one prepared signature but is reached at run +time with another (dtype, rank, layout, read-only); only a pipeline that does so shows +it, and the two guard pipelines are small. Known candidates, not established as +reachable: grayscale closing/opening on integer images with the Numba provider; the +dense 3D intensity kernels and the threshold primitives are prepared for float32 and +writable float64 only. Kernels in a module nothing has imported by readiness are not +enumerated. The first guard was not run against main; main's behaviour there is the +table above. GPU program builds are not covered. + +## Deletions (certificates) + +**unused** (no reference outside the definition and re-exports): +`_scaled_second_channel_costes_numba`, `_scaled_second_channel_costes_sorted_events_numba`, +`_scaled_second_channel_costes_grouped_events_numba`, `_prefix_pearson_from_group_count_numba`, +`quantized_unit_interval_event_summaries` and their re-exports; in +`intensity_object_quantiles_numba`: `_group_starts_numba`, `_median_from_sorted_numba`, +`_quartiles_from_sorted_numba`, `_quantile_from_sorted_group`, +`_quantile_from_dense_sorted_group`, `_object_intensity_std_3d_numba`; +`shape._maximum_position_of_labels_numba`, `shape._radius_features_from_distance_image_numba`; +`zernike._score_intensity_zernike_moments_batch_numba` with `intensity_zernike_moments_batch` +(an exported name) and its three strategy methods; +`distance_propagation_numba._propagate_labels_numba`; `illumination._rank_median_uint16_2d_numba`, +`_select_uint16`, `_partition_uint16`; `morphology._fill_labeled_holes_from_components_numba`, +`_first_adjacent_foreground_label_numba`, `_seed_points_from_components_numba`, +`_restore_removed_declump_basins_numba` with the three `restore_removed_declump_basins` +methods; `_threshold_weighted_variance_unmasked_finite_numba`; +`_minimum_cross_entropy_threshold_numba`. + +**tests-only** (a Numba override removed; the strategy now inherits the NumPy/SciPy +implementation for inputs production does not pass; tests of those inputs pass through +the inherited one): +- `morphology._block_labels_2d_numba`: the Numba `block_labels` override. Production + reaches it only for images of 4 or more dimensions. +- `morphology._fill_labeled_holes_single_label_components_numba`, + `_record_component_boundary_label_numba`: the non-boolean branch of + `fill_labeled_holes_below_size`; the one production caller passes a boolean mask. +- `morphology._relabel_sequential_3d_numba`: both production callers are 2D. +What the interface admits and tests do not vouch for: the speed of those inputs, and bit +equality of the inherited implementation with the removed kernel on them. + +## Not done, ranked + +1. PIL plugin and imagecodecs loading at the first image read or write: 57 modules, + 0.03 s in the server's first execution and in each worker. The owner is PolyStore's + disk backend: PolyStore PR #40 adds `BackendBase.prepare_process` and + `prepare_storage_backends`; after its release the server calls that in place of + `ensure_storage_registry`. +2. `openhcs.interop.cellprofiler.runtime.module_execution`: 0.002 s at first compile + and per worker; a function-local import in `interop/cellprofiler/runtime/adapter.py` + with a real import cycle. +3. First request setup is 0.014 to 0.021 s slower than the second with no import + involved; cause not traced. +4. An in-process orchestrator (no server, no readiness) prepares at compile time what + its pipeline uses; a function called directly, outside a pipeline, compiles at first + use. Unchanged. +5. Signature variants listed under "What they cannot see". + +## Gates (run on `4ff20c10d` plus this section) + +- `verify_branch.sh r1e` (heavy slot): 11656 passed, 34 failed, 2 errors, 25 skipped; + collected 11714, reported 11691, no crash lines; new failures against the known list: 0. +- Official30 parity (heavy slot, `official30.lock`, empty Numba and catalog caches, + traced): `1 passed in 716 s`, the one test that compiles, executes and compares all + 30 cases. The trace of that run: after the first server's cold preparation, 0 + compilations and 0 cache misses in any process across the 30 cases; each of the 29 + later servers loads exactly the 286 signatures its own preparation loads + (8294 = 29 x 286) and nothing after; the 360 workers and the client load nothing. +- `smoke_local.sh`: `smoke rc=0` (7 passed, GUI smoke exit 0). +- `guardrails_local.sh`: `guardrails rc=0`. +- Guards: `tests/integration/test_no_jit_after_readiness.py`, 2 passed (88 s). +- Touched unit tests: `test_callable_abi_preparation`, `test_processing_preparation`, + `test_preparation_admission`, `test_function_catalog_zmq`: 117 passed. diff --git a/openhcs/core/callable_contract.py b/openhcs/core/callable_contract.py index 2bad708b2b..1351786b8f 100644 --- a/openhcs/core/callable_contract.py +++ b/openhcs/core/callable_contract.py @@ -1534,14 +1534,7 @@ def raw_runtime_signature(self) -> inspect.Signature: @classmethod def from_prepared_callable(cls, func: Callable[..., object]) -> "CallableContract": """Construct a compiled view only after the declaration owner prepares it.""" - projection = CallableProjection.from_callable(func) - prepared = cls.from_callable(projection.prepared_callable()) - if prepared.prepared_canonical_signature is None: - prepared = prepared.with_prepared_signature() - return dataclasses.replace( - prepared, func=func, function_name=projection.name, - module_name=projection.module_name, - ) + return CallableProjection.from_callable(func).prepared_contract() def decode_public_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]: """Descend declared scalar enums at authoring, before source projection.""" @@ -2229,13 +2222,45 @@ def from_callable(cls, func: Any) -> "CallableProjection": ) return cls(func=func, name=name, module_name=module_name, namespace=namespace) - def prepared_callable(self) -> Any: - """Keep ready library references or prepare live authored targets before validation.""" - prepared = DeclaredSignature().read( + def has_published_signature(self) -> bool: + """Whether catalog preparation already published this callable's signature.""" + return DeclaredSignature().read( self.namespace.get(FunctionContractAttribute.canonical_signature), owner=f"{self.name!r}.{FunctionContractAttribute.canonical_signature}", + ) is not None + + def declared_callable(self) -> Any: + """Keep ready library references or resolve live authored targets, preparing nothing.""" + if self.has_published_signature(): + return self.func + return _resolve_declared_callable(self.func) + + def declared_contract(self) -> "CallableContract": + """The contract read from the declaration alone, preparing nothing. + + A signature is read from the callable's annotations and preparation never + changes it. Authoring a pipeline (converting a ``.cppipe``, lowering a + function pattern) runs no kernel, so it must neither compile nor load + one: kernels belong to the catalog preparation of the process that executes. + """ + return self._contract_of(self.declared_callable()) + + def prepared_contract(self) -> "CallableContract": + """The same contract, after the declaration owner prepared the callable.""" + return self._contract_of(self.prepared_callable()) + + def _contract_of(self, resolved: Any) -> "CallableContract": + contract = CallableContract.from_callable(resolved) + if contract.prepared_canonical_signature is None: + contract = contract.with_prepared_signature() + return dataclasses.replace( + contract, func=self.func, function_name=self.name, + module_name=self.module_name, ) - if prepared is not None: + + def prepared_callable(self) -> Any: + """Keep ready library references or prepare live authored targets before validation.""" + if self.has_published_signature(): return self.func resolved = _resolve_declared_callable(self.func) for target in type(self).from_callable(resolved).prepare_targets(): diff --git a/openhcs/core/dataset_sources/discovery.py b/openhcs/core/dataset_sources/discovery.py index 6fe138ba20..07877ff591 100644 --- a/openhcs/core/dataset_sources/discovery.py +++ b/openhcs/core/dataset_sources/discovery.py @@ -18,6 +18,7 @@ from openhcs.core.axes import AxisFamily _KERNEL_PACKAGE = __name__.rpartition(".")[0] +_domain_registries: list[LazyDiscoveryDict] = [] def load_domain_extensions() -> None: @@ -45,8 +46,10 @@ def domain_registry_config( ) -> RegistryConfig: """Registry configuration for a kernel family that domains extend.""" + registry = LazyDiscoveryDict(enable_cache=False) + _domain_registries.append(registry) return RegistryConfig( - registry_dict=LazyDiscoveryDict(enable_cache=False), + registry_dict=registry, key_attribute=key_attribute, skip_if_no_key=skip_if_no_key, registry_name=registry_name, @@ -56,4 +59,23 @@ def domain_registry_config( ) -__all__ = ["domain_registry_config", "load_domain_extensions"] +def discover_domain_registries() -> None: + """Complete every domain-extended family registry now, as first access would. + + A server calls this before it forks workers and reports ready, so neither + it nor a worker imports kernel or domain members while serving a request. + Discovery imports the family roots that were not yet defined, so the loop + follows the list as it grows. + """ + + discovered = 0 + while discovered < len(_domain_registries): + _domain_registries[discovered].values() + discovered += 1 + + +__all__ = [ + "discover_domain_registries", + "domain_registry_config", + "load_domain_extensions", +] diff --git a/openhcs/core/function_patterns.py b/openhcs/core/function_patterns.py index be0a0d8b30..31beafd508 100644 --- a/openhcs/core/function_patterns.py +++ b/openhcs/core/function_patterns.py @@ -43,6 +43,7 @@ ) from openhcs.core.aligned_image_payload import ImagePayloadStackComposition from openhcs.core.callable_contract import ( + CallableProjection, CallableContract, FunctionStepExecutionScope, ImagePayloadConsumption, @@ -409,7 +410,7 @@ def from_chain( kwargs = {} if kwargs is None else kwargs if RUNTIME_CALLABLE_KWARG_POLICY.item_is_disabled(kwargs): continue - contract = CallableContract.from_prepared_callable(func) + contract = CallableProjection.from_callable(func).declared_contract() normalized_items.append( NormalizedFunctionItem( key=FunctionInvocationKey.from_contract( @@ -1553,6 +1554,28 @@ def normalize_function_pattern( ) +def normalize_prepared_function_pattern( + pattern: FunctionPatternSyntax | NormalizedFunctionPattern, +) -> NormalizedFunctionPattern: + """Prepare every enabled callable, then lower the syntax: the compiler's entry. + + Lowering alone prepares nothing, so authoring and validation never compile a + kernel. A pipeline about to be compiled for execution prepares first, and the + contracts are read afterwards. A process that ran catalog preparation has + nothing left to prepare here. + """ + if isinstance(pattern, NormalizedFunctionPattern): + for item in pattern.iter_items(): + CallableProjection.from_callable(item.func).prepared_callable() + return pattern + for _, chain in AuthoredFunctionPattern.groups(pattern): + for leaf in chain: + func, kwargs = AuthoredFunctionPattern.split_leaf(leaf) + if not RUNTIME_CALLABLE_KWARG_POLICY.item_is_disabled(kwargs or {}): + CallableProjection.from_callable(func).prepared_callable() + return normalize_function_pattern(pattern) + + def compile_function_pattern( graph: "ArtifactGraph", input_plans: Mapping[ArtifactSpecRef, ArtifactInputPlan], diff --git a/openhcs/core/pipeline/compilation_session.py b/openhcs/core/pipeline/compilation_session.py index 776357a4ed..9c341d56bc 100644 --- a/openhcs/core/pipeline/compilation_session.py +++ b/openhcs/core/pipeline/compilation_session.py @@ -15,7 +15,7 @@ from openhcs.core.steps.abstract import AbstractStep from openhcs.core.steps.function_step import FunctionStep from openhcs.core.function_patterns import ( - normalize_function_pattern, + normalize_prepared_function_pattern, NormalizedFunctionItem, ) from openhcs.core.invocation_artifacts import ( @@ -389,7 +389,7 @@ def __post_init__(self) -> None: "steps", tuple( ( - step.with_function_spec(normalize_function_pattern(step.func or [])) + step.with_function_spec(normalize_prepared_function_pattern(step.func or [])) if isinstance(step, FunctionStep) else step ) diff --git a/openhcs/interop/cellprofiler/module_artifact_contracts.py b/openhcs/interop/cellprofiler/module_artifact_contracts.py index 837b72448f..00609a4239 100644 --- a/openhcs/interop/cellprofiler/module_artifact_contracts.py +++ b/openhcs/interop/cellprofiler/module_artifact_contracts.py @@ -22,7 +22,7 @@ ImageArtifactType, MeasurementsArtifactType, ) -from openhcs.core.callable_contract import CallableContract +from openhcs.core.callable_contract import CallableContract, CallableProjection from openhcs.core.function_patterns import DEFAULT_GROUP_KEY from openhcs.core.invocation_artifacts import ArtifactDeclarationStepContext from openhcs.interop.cellprofiler.settings_binder import SettingToKeywordBinding @@ -1050,9 +1050,9 @@ def artifact_output_relations( del module, binding, name, step_context, output_position primary_images = cls.primary_image_inputs( - CallableContract.from_prepared_callable( + CallableProjection.from_callable( cls.require_callable(invocation_key.function_name) - ), + ).declared_contract(), artifact_inputs.specs, ) contextual_inputs = tuple( @@ -1183,7 +1183,7 @@ def callable_contract( group_key=invocation_key.group_key, ) func = cls.require_callable(invocation_key.function_name) - callable_contract = CallableContract.from_prepared_callable(func) + callable_contract = CallableProjection.from_callable(func).declared_contract() inputs = cls.artifact_contract_inputs( module, invocation_key=invocation_key, diff --git a/openhcs/interop/cellprofiler/module_artifact_declarations.py b/openhcs/interop/cellprofiler/module_artifact_declarations.py index 2a82729330..c1a5177c02 100644 --- a/openhcs/interop/cellprofiler/module_artifact_declarations.py +++ b/openhcs/interop/cellprofiler/module_artifact_declarations.py @@ -40,6 +40,7 @@ ObjectLineageArtifactType, ) from openhcs.core.callable_contract import ( + CallableProjection, CallableContract, ImagePayloadConsumption, ) @@ -741,9 +742,9 @@ def measurement_output_relations( for artifact_input in artifact_inputs.specs ) invocation_domain_inputs = cls.invocation_domain_inputs( - CallableContract.from_prepared_callable( + CallableProjection.from_callable( cls.require_callable(invocation_key.function_name) - ), + ).declared_contract(), artifact_inputs.specs, ) return ( diff --git a/openhcs/processing/backends/analysis/region_properties.py b/openhcs/processing/backends/analysis/region_properties.py index c704331e3a..4e1972b8df 100644 --- a/openhcs/processing/backends/analysis/region_properties.py +++ b/openhcs/processing/backends/analysis/region_properties.py @@ -174,6 +174,14 @@ class SkimageNumpyLabelRegionPropertiesBackendStrategy( backend_provider = SkimageBackendProvider is_default_backend = False + def prepare_backend(self) -> None: + # Orientation is the one Numba kernel behind the scikit-image table, + # whose bounding-box columns are strided for several labels and + # contiguous for one. + labels = np.array([[0, 1, 1], [0, 1, 0], [2, 2, 0]], dtype=np.int32) + self.measure_2d(labels, include_advanced=False) + self.measure_2d(np.minimum(labels, 1), include_advanced=False) + def measure_2d( self, labels: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/alignment.py b/openhcs/processing/backends/cellprofiler/alignment.py index 26d5089256..dd4cb4a8cf 100644 --- a/openhcs/processing/backends/cellprofiler/alignment.py +++ b/openhcs/processing/backends/cellprofiler/alignment.py @@ -522,6 +522,9 @@ def prepare_backend(self) -> None: pixels = np.arange(16, dtype=np.float64).reshape((4, 4)) mask = np.ones(pixels.shape, dtype=np.bool_) self.mutual_information_offset(pixels, pixels, mask, mask) + partial_mask = mask.copy() + partial_mask[0, 0] = False + self.mutual_information_offset(pixels, pixels, partial_mask, mask) def mutual_information_offset( self, diff --git a/openhcs/processing/backends/cellprofiler/colocalization.py b/openhcs/processing/backends/cellprofiler/colocalization.py index 1cea5b180d..d549ebf6b1 100644 --- a/openhcs/processing/backends/cellprofiler/colocalization.py +++ b/openhcs/processing/backends/cellprofiler/colocalization.py @@ -1447,6 +1447,8 @@ def prepare_backend(self) -> None: ) thresholded_colocalization_metrics(first, second, 15.0, True, True, True) self.linear_costes(first, second, 255, False) + # Positively correlated channels take the sorted-events search. + self.linear_costes(first, np.sqrt(first), 255, False) quantized_codes = np.arange(64 * 64, dtype=np.uint16) % 512 + 1024 quantized = quantized_codes.astype(np.float32) / np.float32(65535) self.scaled_second_channel_costes(quantized, quantized.copy(), 255) diff --git a/openhcs/processing/backends/cellprofiler/colocalization_costes.py b/openhcs/processing/backends/cellprofiler/colocalization_costes.py index bab4be6bee..f2b11d0c6d 100644 --- a/openhcs/processing/backends/cellprofiler/colocalization_costes.py +++ b/openhcs/processing/backends/cellprofiler/colocalization_costes.py @@ -17,15 +17,10 @@ _max_pair_numba, _pearson_below_threshold_numba, _prefix_pearson_from_count_numba, - _prefix_pearson_from_group_count_numba, _prefix_pearson_numba, _prefix_value_numba, _regression_line_numba, - _scaled_second_channel_costes_grouped_events_numba, - _scaled_second_channel_costes_numba, - _scaled_second_channel_costes_sorted_events_numba, UnitIntervalDenseRankSemantics, - quantized_unit_interval_event_summaries, ) diff --git a/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix.py b/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix.py index 1e2bbb7797..a73713ae94 100644 --- a/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix.py +++ b/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix.py @@ -15,13 +15,9 @@ _max_pair_numba, _pearson_below_threshold_numba, _prefix_pearson_from_count_numba, - _prefix_pearson_from_group_count_numba, _prefix_pearson_numba, _prefix_value_numba, _regression_line_numba, - _scaled_second_channel_costes_grouped_events_numba, - _scaled_second_channel_costes_numba, - _scaled_second_channel_costes_sorted_events_numba, ) @@ -115,71 +111,3 @@ def dense_codes_and_values( dense = np.ascontiguousarray(lookup[codes], dtype=np.int64) values = np.flatnonzero(present).astype(np.float32, copy=False) return dense, (values / np.float32(scale)).astype(np.float64, copy=False) - - -def quantized_unit_interval_event_summaries( - first: np.ndarray, - second: np.ndarray, - slope: float, - intercept: float, - preferred_scale: int, -) -> ( - tuple[ - np.ndarray, - np.ndarray, - np.ndarray, - np.ndarray, - np.ndarray, - np.ndarray, - np.ndarray, - ] - | None -): - first_code_result = UnitIntervalDenseRankSemantics.integer_codes( - first, - preferred_scale=preferred_scale, - ) - second_code_result = UnitIntervalDenseRankSemantics.integer_codes( - second, - preferred_scale=preferred_scale, - ) - if first_code_result is None or second_code_result is None: - return None - first_codes, first_scale = first_code_result - second_codes, second_scale = second_code_result - first_dense, first_values = UnitIntervalDenseRankSemantics.dense_codes_and_values( - first_codes, - first_scale, - ) - second_dense, second_values = UnitIntervalDenseRankSemantics.dense_codes_and_values( - second_codes, - second_scale, - ) - second_value_count = second_values.size - combined = first_dense * second_value_count + second_dense - pair_counts = np.bincount( - combined, - minlength=first_values.size * second_value_count, - ) - active_pairs = np.flatnonzero(pair_counts) - if active_pairs.size == 0: - return None - first_index = active_pairs // second_value_count - second_index = active_pairs - first_index * second_value_count - counts = pair_counts[active_pairs].astype(np.float64, copy=False) - first_pair_values = first_values[first_index] - second_pair_values = second_values[second_index] - event_thresholds = np.minimum( - second_pair_values, - slope * first_pair_values + intercept, - ) - unique_events, inverse = np.unique(event_thresholds, return_inverse=True) - return ( - unique_events, - np.bincount(inverse, weights=counts).astype(np.int64, copy=False), - np.bincount(inverse, weights=counts * first_pair_values), - np.bincount(inverse, weights=counts * second_pair_values), - np.bincount(inverse, weights=counts * first_pair_values * first_pair_values), - np.bincount(inverse, weights=counts * second_pair_values * second_pair_values), - np.bincount(inverse, weights=counts * first_pair_values * second_pair_values), - ) diff --git a/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix_numba_pearson.py b/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix_numba_pearson.py index 1675a35307..b42a43f5f3 100644 --- a/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix_numba_pearson.py +++ b/openhcs/processing/backends/cellprofiler/colocalization_costes_prefix_numba_pearson.py @@ -425,183 +425,6 @@ def _integer_unit_interval_codes_for_scale_numba( return True, codes -@njit(cache=True) -def _scaled_second_channel_costes_numba( - first: np.ndarray, - second: np.ndarray, - scale_max: int, -) -> tuple[float, float]: - valid, slope, intercept = _regression_line_numba(first, second) - if not valid: - return 0.0, 0.0 - - minimum_scale_index = min(scale_max, 5) - selected_first_threshold = 0.0 - selected_second_threshold = minimum_scale_index / scale_max - selected_correlation = np.nan - - for second_scale_index in range(scale_max, minimum_scale_index - 1, -1): - second_threshold = second_scale_index / scale_max - if slope == 0.0: - first_threshold = 0.0 - else: - first_threshold = (second_threshold - intercept) / slope - if first_threshold < 0.0: - first_threshold = 0.0 - - _count, cost_regression = _pearson_below_threshold_numba( - first, - second, - first_threshold, - second_threshold, - ) - selected_first_threshold = first_threshold - selected_second_threshold = second_threshold - selected_correlation = cost_regression - if np.isfinite(cost_regression) and cost_regression <= 0.0: - break - - if ( - (not np.isfinite(selected_correlation)) - or selected_correlation > 0.0 - or selected_first_threshold <= 0.0 - ): - selected_first_threshold = 0.0 - - return selected_first_threshold, selected_second_threshold - - -@njit(cache=True) -def _scaled_second_channel_costes_sorted_events_numba( - costes_context: tuple, -) -> tuple[float, float]: - ( - sorted_event_threshold, - prefix_x, - prefix_y, - prefix_x2, - prefix_y2, - prefix_xy, - scale_max, - slope, - intercept, - ) = costes_context - minimum_scale_index = min(scale_max, 5) - selected_first_threshold = 0.0 - selected_second_threshold = minimum_scale_index / scale_max - selected_correlation = np.nan - - for second_scale_index in range(scale_max, minimum_scale_index - 1, -1): - second_threshold = second_scale_index / scale_max - first_threshold = (second_threshold - intercept) / slope - if first_threshold < 0.0: - first_threshold = 0.0 - count = _event_count_for_threshold_numba( - sorted_event_threshold, - second_threshold, - ) - cost_regression = _prefix_pearson_from_count_numba( - count, - prefix_x, - prefix_y, - prefix_x2, - prefix_y2, - prefix_xy, - ) - selected_first_threshold = first_threshold - selected_second_threshold = second_threshold - selected_correlation = cost_regression - if np.isfinite(cost_regression) and cost_regression <= 0.0: - break - - if ( - (not np.isfinite(selected_correlation)) - or selected_correlation > 0.0 - or selected_first_threshold <= 0.0 - ): - selected_first_threshold = 0.0 - - return selected_first_threshold, selected_second_threshold - - -@njit(cache=True) -def _scaled_second_channel_costes_grouped_events_numba( - costes_context: tuple, - prefix_count: np.ndarray, -) -> tuple[float, float]: - ( - sorted_event_threshold, - prefix_x, - prefix_y, - prefix_x2, - prefix_y2, - prefix_xy, - scale_max, - slope, - intercept, - ) = costes_context - minimum_scale_index = min(scale_max, 5) - selected_first_threshold = 0.0 - selected_second_threshold = minimum_scale_index / scale_max - selected_correlation = np.nan - - for second_scale_index in range(scale_max, minimum_scale_index - 1, -1): - second_threshold = second_scale_index / scale_max - first_threshold = (second_threshold - intercept) / slope - if first_threshold < 0.0: - first_threshold = 0.0 - group_count = _event_count_for_threshold_numba( - sorted_event_threshold, - second_threshold, - ) - cost_regression = _prefix_pearson_from_group_count_numba( - group_count, - prefix_count, - prefix_x, - prefix_y, - prefix_x2, - prefix_y2, - prefix_xy, - ) - selected_first_threshold = first_threshold - selected_second_threshold = second_threshold - selected_correlation = cost_regression - if np.isfinite(cost_regression) and cost_regression <= 0.0: - break - - if ( - (not np.isfinite(selected_correlation)) - or selected_correlation > 0.0 - or selected_first_threshold <= 0.0 - ): - selected_first_threshold = 0.0 - - return selected_first_threshold, selected_second_threshold - - -@njit(cache=True) -def _prefix_pearson_from_group_count_numba( - group_count: int, - prefix_count: np.ndarray, - prefix_x: np.ndarray, - prefix_y: np.ndarray, - prefix_x2: np.ndarray, - prefix_y2: np.ndarray, - prefix_xy: np.ndarray, -) -> float: - if group_count <= 0: - return np.nan - count = int(prefix_count[group_count - 1]) - return _prefix_pearson_numba( - count, - _prefix_value_numba(prefix_x, group_count), - _prefix_value_numba(prefix_y, group_count), - _prefix_value_numba(prefix_x2, group_count), - _prefix_value_numba(prefix_y2, group_count), - _prefix_value_numba(prefix_xy, group_count), - ) - - @njit(cache=True) def _costes_manders_numba( first: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/distance_propagation_numba.py b/openhcs/processing/backends/cellprofiler/distance_propagation_numba.py index a9bb05ccac..f562fbe975 100644 --- a/openhcs/processing/backends/cellprofiler/distance_propagation_numba.py +++ b/openhcs/processing/backends/cellprofiler/distance_propagation_numba.py @@ -283,30 +283,6 @@ def _propagation_heap_pop( return size, value, label, y, x -@njit(cache=True) -def _propagate_labels_numba( - image: np.ndarray, - seed_labels: np.ndarray, - mask: np.ndarray, - weight: float, - max_distance: float, -) -> np.ndarray: - output, distances = _propagate_labels_and_distances_numba( - image, - seed_labels, - mask, - weight, - max_distance, - ) - if max_distance >= 0.0: - height, width = output.shape - for y in range(height): - for x in range(width): - if distances[y, x] > max_distance and seed_labels[y, x] <= 0: - output[y, x] = 0 - return output - - @njit(cache=True) def _propagate_labels_and_distances_numba( image: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/granularity.py b/openhcs/processing/backends/cellprofiler/granularity.py index 08fd8dbe60..0514139e8d 100644 --- a/openhcs/processing/backends/cellprofiler/granularity.py +++ b/openhcs/processing/backends/cellprofiler/granularity.py @@ -1713,6 +1713,15 @@ def _prepare_granularity_backend() -> None: element_radius=10, spectrum_length=5, ) + # A subsampled spectrum takes object means from resampled values. + measure_granularity_objects.__wrapped__( + image, + labels, + subsample_size=0.25, + background_subsample_size=0.25, + element_radius=10, + spectrum_length=5, + ) __all__ = [ diff --git a/openhcs/processing/backends/cellprofiler/illumination.py b/openhcs/processing/backends/cellprofiler/illumination.py index 9b28e4566a..df88393b0c 100644 --- a/openhcs/processing/backends/cellprofiler/illumination.py +++ b/openhcs/processing/backends/cellprofiler/illumination.py @@ -2553,67 +2553,6 @@ def _fenwick_select_uint16(tree: np.ndarray, kth: int) -> np.uint16: return np.uint16(_fenwick_select_index(tree, kth, 32768)) -@njit(cache=True) -def _rank_median_uint16_2d_numba( - image: np.ndarray, mask: np.ndarray, offsets_y: np.ndarray, offsets_x: np.ndarray -) -> np.ndarray: - height, width = image.shape - output = np.empty((height, width), dtype=np.uint16) - footprint_size = offsets_y.shape[0] - for y in range(height): - values = np.empty(footprint_size, dtype=np.uint16) - for x in range(width): - count = 0 - for offset_index in range(footprint_size): - yy = y + offsets_y[offset_index] - xx = x + offsets_x[offset_index] - if 0 <= yy < height and 0 <= xx < width: - if mask[yy, xx]: - values[count] = image[yy, xx] - else: - values[count] = np.uint16(0) - count += 1 - output[y, x] = _select_uint16(values, count, count // 2) - return output - - -@njit(cache=True) -def _select_uint16(values: np.ndarray, count: int, kth: int) -> np.uint16: - left = 0 - right = count - 1 - while True: - if left == right: - return values[left] - pivot_index = (left + right) // 2 - pivot_index = _partition_uint16(values, left, right, pivot_index) - if kth == pivot_index: - return values[kth] - if kth < pivot_index: - right = pivot_index - 1 - else: - left = pivot_index + 1 - - -@njit(cache=True) -def _partition_uint16( - values: np.ndarray, left: int, right: int, pivot_index: int -) -> int: - pivot_value = values[pivot_index] - values[pivot_index] = values[right] - values[right] = pivot_value - store_index = left - for index in range(left, right): - if values[index] < pivot_value: - current = values[store_index] - values[store_index] = values[index] - values[index] = current - store_index += 1 - current = values[right] - values[right] = values[store_index] - values[store_index] = current - return store_index - - __all__ = public_names_from_objects( AutomaticSmoothingFilterSizeStrategy, AverageSmoothingPlaneStrategy, diff --git a/openhcs/processing/backends/cellprofiler/intensity.py b/openhcs/processing/backends/cellprofiler/intensity.py index b052b24ff3..6eb8d2a47e 100644 --- a/openhcs/processing/backends/cellprofiler/intensity.py +++ b/openhcs/processing/backends/cellprofiler/intensity.py @@ -1408,6 +1408,16 @@ def prepare_backend(self) -> None: self.measure_prepared_batch( (image_3d, np.ascontiguousarray(1.0 - image_3d)), prepared_3d ) + # Labels filling most of the volume take the dense scan, not the sparse one. + dense_labels_3d = np.ones(image_3d.shape, dtype=np.int32) + dense_labels_3d[4:] = 2 + self.measure( + image_3d, + ObjectLabelPayload( + variant_data=ObjectLabelVariantData(labels=dense_labels_3d), + domain=ObjectLabelDomain(declared_object_ids=(1, 2)), + ), + ) def object_intensity_backend( diff --git a/openhcs/processing/backends/cellprofiler/intensity_object_quantiles_numba.py b/openhcs/processing/backends/cellprofiler/intensity_object_quantiles_numba.py index 2cef276048..c7e66dd129 100644 --- a/openhcs/processing/backends/cellprofiler/intensity_object_quantiles_numba.py +++ b/openhcs/processing/backends/cellprofiler/intensity_object_quantiles_numba.py @@ -445,110 +445,6 @@ def _quantile_from_partitioned_rank_pair( return low_value * (1.0 - qfraction) + high_value * qfraction -@njit(cache=True) -def _quantile_from_dense_sorted_group( - sorted_values: np.ndarray, - fraction: float, -) -> float: - count = sorted_values.size - if count <= 0: - return 0.0 - qindex = count * fraction - low = int(qindex) - qfraction = qindex - low - last = count - 1 - if low < last: - return ( - sorted_values[low] * (1.0 - qfraction) + sorted_values[low + 1] * qfraction - ) - return sorted_values[last] - - -@njit(cache=True) -def _group_starts_numba( - sorted_labels: np.ndarray, - label_to_index: np.ndarray, - starts: np.ndarray, -) -> None: - for offset in range(sorted_labels.size): - label = sorted_labels[offset] - if label < 0 or label >= label_to_index.size: - continue - index = label_to_index[label] - if index >= 0 and starts[index] < 0: - starts[index] = offset - - -@njit(cache=True) -def _quartiles_from_sorted_numba( - sorted_values: np.ndarray, - starts: np.ndarray, - counts: np.ndarray, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - object_count = counts.size - lower = np.zeros(object_count, dtype=np.float64) - median = np.zeros(object_count, dtype=np.float64) - upper = np.zeros(object_count, dtype=np.float64) - for index in range(object_count): - lower[index] = _quantile_from_sorted_group( - sorted_values, - starts[index], - counts[index], - 0.25, - ) - median[index] = _quantile_from_sorted_group( - sorted_values, - starts[index], - counts[index], - 0.5, - ) - upper[index] = _quantile_from_sorted_group( - sorted_values, - starts[index], - counts[index], - 0.75, - ) - return lower, median, upper - - -@njit(cache=True) -def _median_from_sorted_numba( - sorted_values: np.ndarray, - starts: np.ndarray, - counts: np.ndarray, -) -> np.ndarray: - object_count = counts.size - output = np.zeros(object_count, dtype=np.float64) - for index in range(object_count): - output[index] = _quantile_from_sorted_group( - sorted_values, - starts[index], - counts[index], - 0.5, - ) - return output - - -@njit(cache=True) -def _quantile_from_sorted_group( - sorted_values: np.ndarray, - start: int, - count: int, - fraction: float, -) -> float: - if count <= 0: - return 0.0 - qindex = start + count * fraction - low = int(qindex) - qfraction = qindex - low - last = start + count - 1 - if low < last: - return ( - sorted_values[low] * (1.0 - qfraction) + sorted_values[low + 1] * qfraction - ) - return sorted_values[last] - - @njit(cache=True) def _object_intensity_scan_numba( image: np.ndarray, @@ -1001,60 +897,6 @@ def _is_inner_boundary_voxel( return False -@njit(cache=True) -def _object_intensity_std_3d_numba( - image: np.ndarray, - labels: np.ndarray, - label_to_index: np.ndarray, - means: np.ndarray, - edge_means: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - object_count = means.size - variance_sums = np.zeros(object_count, dtype=np.float64) - counts = np.zeros(object_count, dtype=np.float64) - edge_variance_sums = np.zeros(object_count, dtype=np.float64) - edge_counts = np.zeros(object_count, dtype=np.float64) - z_size, y_size, x_size = image.shape - for z_index in range(z_size): - for y_index in range(y_size): - for x_index in range(x_size): - label = labels[z_index, y_index, x_index] - if label <= 0 or label >= label_to_index.size: - continue - object_index = label_to_index[label] - if object_index < 0: - continue - value = image[z_index, y_index, x_index] - if not np.isfinite(value): - continue - diff = value - means[object_index] - variance_sums[object_index] += diff * diff - counts[object_index] += 1.0 - if _is_inner_boundary_voxel( - labels, - z_index, - y_index, - x_index, - label, - ): - edge_diff = value - edge_means[object_index] - edge_variance_sums[object_index] += edge_diff * edge_diff - edge_counts[object_index] += 1.0 - - stds = np.zeros(object_count, dtype=np.float64) - edge_stds = np.zeros(object_count, dtype=np.float64) - for object_index in range(object_count): - if counts[object_index] > 0.0: - stds[object_index] = np.sqrt( - variance_sums[object_index] / counts[object_index] - ) - if edge_counts[object_index] > 0.0: - edge_stds[object_index] = np.sqrt( - edge_variance_sums[object_index] / edge_counts[object_index] - ) - return stds, edge_stds - - @njit(cache=True) def _object_intensity_std_3d_batch_numba( images: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/label_geometry.py b/openhcs/processing/backends/cellprofiler/label_geometry.py index f2d4c136f1..b645c1616e 100644 --- a/openhcs/processing/backends/cellprofiler/label_geometry.py +++ b/openhcs/processing/backends/cellprofiler/label_geometry.py @@ -53,13 +53,15 @@ def codegen(context, builder, resolved_signature, arguments): return signature, codegen -@njit(cache=True) -def _numpy_124_svml_arccos(values: np.ndarray) -> np.ndarray: - """Return NumPy 1.24 standard-SVML arccos values for a flat array.""" - result = np.empty(values.size, dtype=np.float64) - for index in range(values.size): - result[index] = _numpy_124_scalar_arccos(values[index]) - return result +if _NUMPY_124_SVML_ACOS_AVAILABLE: + # Defined only where it can run: a kernel that exists is one preparation compiles. + @njit(cache=True) + def _numpy_124_svml_arccos(values: np.ndarray) -> np.ndarray: + """Return NumPy 1.24 standard-SVML arccos values for a flat array.""" + result = np.empty(values.size, dtype=np.float64) + for index in range(values.size): + result[index] = _numpy_124_scalar_arccos(values[index]) + return result _DEKKER_SPLITTER = 134217729.0 # 2**27 + 1 @@ -238,6 +240,17 @@ def _numpy_124_arccos(values: np.ndarray) -> np.ndarray: return _numpy_124_portable_arccos(value_array) +def prepare_arccos_kernels() -> None: + """Compile the arccos kernel this CPU uses, and the portable one. + + A process started with AVX-512 disabled shares the Numba cache and uses the + portable kernel, so it is compiled on every machine. + """ + values = np.zeros(1, dtype=np.float64) + _numpy_124_arccos(values) + _numpy_124_portable_arccos(values) + + def feret_diameters_from_labels( labels: np.ndarray, label_ids: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/morphology.py b/openhcs/processing/backends/cellprofiler/morphology.py index 01e7995e21..4f25bdc9e4 100644 --- a/openhcs/processing/backends/cellprofiler/morphology.py +++ b/openhcs/processing/backends/cellprofiler/morphology.py @@ -1988,15 +1988,6 @@ def fill_labeled_holes_below_size( ) -> np.ndarray: """Fill enclosed background components smaller than a size limit.""" - @abstractmethod - def restore_removed_declump_basins( - self, - pre_declump_labels: np.ndarray, - labels_before_size_filter: np.ndarray, - labels_after_size_filter: np.ndarray, - ) -> np.ndarray: - """Restore removed watershed basins with one surviving original identity.""" - @abstractmethod def local_maxima_by_label( self, image: np.ndarray, labels: np.ndarray, footprint: np.ndarray @@ -2243,18 +2234,6 @@ def _scipy_fill_labeled_holes( ) return output - def restore_removed_declump_basins( - self, - pre_declump_labels: np.ndarray, - labels_before_size_filter: np.ndarray, - labels_after_size_filter: np.ndarray, - ) -> np.ndarray: - return _restore_removed_declump_basins_numba( - np.ascontiguousarray(pre_declump_labels, dtype=np.int64), - np.ascontiguousarray(labels_before_size_filter, dtype=np.int64), - np.ascontiguousarray(labels_after_size_filter, dtype=np.int64), - ).astype(np.asarray(labels_after_size_filter).dtype, copy=False) - def local_maxima_by_label( self, image: np.ndarray, labels: np.ndarray, footprint: np.ndarray ) -> np.ndarray: @@ -2341,6 +2320,11 @@ def prepare_backend(self) -> None: for writable_image in (False, True): image = np.arange(9, dtype=dtype).reshape((3, 3)) image.flags.writeable = writable_image + # Named on this class: the OpenCV subclass replaces closing + # with cv2 and shares this preparation. + NumbaNumpyMorphologyBackendStrategy.grayscale_closing( + self, image, footprint + ) for full_mask in (False, True): for writable_mask in (False, True): smoothing_mask = ( @@ -2433,16 +2417,6 @@ def _grayscale_morphology( first_pass_is_dilation, ) - def block_labels( - self, image_shape: tuple[int, int], block_size: int - ) -> tuple[np.ndarray, np.ndarray]: - if len(image_shape) != 2: - raise NotImplementedError( - "Numba morphology backend currently supports 2-D block labels." - ) - height, width = image_shape - return _block_labels_2d_numba(int(height), int(width), max(1, int(block_size))) - def blockwise_minimum( self, image: np.ndarray, mask: np.ndarray | None, block_size: int ) -> np.ndarray: @@ -2628,33 +2602,14 @@ def _fill_labeled_holes_2d( size_predicate=size_predicate, ) - def restore_removed_declump_basins( - self, - pre_declump_labels: np.ndarray, - labels_before_size_filter: np.ndarray, - labels_after_size_filter: np.ndarray, - ) -> np.ndarray: - pre_declump = np.asarray(pre_declump_labels) - before = np.asarray(labels_before_size_filter) - after = np.asarray(labels_after_size_filter) - if pre_declump.ndim != 2 or before.ndim != 2 or after.ndim != 2: - raise NotImplementedError( - "Numba morphology backend currently supports 2-D declump basin restoration." - ) - if pre_declump.shape != before.shape or before.shape != after.shape: - raise ValueError( - f"Declump basin restoration inputs must have identical shapes; got {pre_declump.shape!r}, {before.shape!r}, and {after.shape!r}." - ) - return _restore_removed_declump_basins_numba( - np.ascontiguousarray(pre_declump, dtype=np.int64), - np.ascontiguousarray(before, dtype=np.int64), - np.ascontiguousarray(after, dtype=np.int64), - ).astype(after.dtype, copy=False) - def fill_labeled_holes_below_size( self, labels: np.ndarray, maximum_hole_size: int ) -> np.ndarray: labels_array = np.asarray(labels) + if labels_array.dtype != np.bool_: + return super().fill_labeled_holes_below_size( + labels_array, maximum_hole_size + ) if labels_array.ndim != 2: return self._fill_labeled_holes_below_size_planewise( labels_array, maximum_hole_size @@ -2685,11 +2640,7 @@ def _fill_labeled_holes_below_size_2d( fill_flags = _hole_fill_flags_below_size_numba( sizes, touches_border, component_count, int(maximum_hole_size) ) - if labels.dtype == np.bool_: - return _fill_binary_holes_from_components_numba( - np.ascontiguousarray(labels), components, fill_flags - ) - return _fill_labeled_holes_single_label_components_numba( + return _fill_binary_holes_from_components_numba( np.ascontiguousarray(labels), components, fill_flags ) @@ -2766,16 +2717,8 @@ def relabel_sequential(self, labels: np.ndarray) -> tuple[np.ndarray, int]: if np.issubdtype(labels_array.dtype, np.integer) else np.dtype(np.int64) ) - if labels_array.ndim > 2: - relabeled_planes, count = _relabel_sequential_3d_numba( - np.ascontiguousarray( - labels_array.reshape((-1, *labels_array.shape[-2:])), - dtype=label_dtype, - ) - ) - return (relabeled_planes.reshape(labels_array.shape), int(count)) if labels_array.ndim != 2: - raise ValueError("Relabeling requires at least two dimensions.") + return super().relabel_sequential(labels_array) return _relabel_sequential_numba( np.ascontiguousarray(labels_array, dtype=label_dtype) ) @@ -2911,28 +2854,6 @@ def _scipy_erode_labeled_objects( return labels_array * (contours == 0) -@njit(cache=True) -def _block_labels_2d_numba( - height: int, width: int, block_size: int -) -> tuple[np.ndarray, np.ndarray]: - row_blocks = max(1, int(np.floor(float(height) / float(block_size)))) - column_blocks = max(1, int(np.floor(float(width) / float(block_size)))) - labels = np.empty((height, width), dtype=np.int32) - indexes = np.empty(row_blocks * column_blocks, dtype=np.int32) - for row in range(row_blocks): - y_start = int(np.ceil(float(row * height) / float(row_blocks))) - y_stop = int(np.ceil(float((row + 1) * height) / float(row_blocks))) - for column in range(column_blocks): - x_start = int(np.ceil(float(column * width) / float(column_blocks))) - x_stop = int(np.ceil(float((column + 1) * width) / float(column_blocks))) - label = row * column_blocks + column - indexes[label] = label - for y in range(y_start, y_stop): - for x in range(x_start, x_stop): - labels[y, x] = label - return (labels, indexes) - - @njit(cache=True) def _blockwise_minimum_numba( image: np.ndarray, mask: np.ndarray, has_mask: bool, block_size: int @@ -4364,51 +4285,6 @@ def _fill_binary_holes_from_components_numba( return output -@njit(cache=True) -def _fill_labeled_holes_single_label_components_numba( - labels: np.ndarray, components: np.ndarray, fill_flags: np.ndarray -) -> np.ndarray: - height, width = labels.shape - component_labels = np.zeros(fill_flags.size, dtype=np.int64) - for y in range(height): - for x in range(width): - component = components[y, x] - if component <= 0 or not fill_flags[component]: - continue - if y > 0: - _record_component_boundary_label_numba( - component_labels, component, int(labels[y - 1, x]) - ) - if y + 1 < height: - _record_component_boundary_label_numba( - component_labels, component, int(labels[y + 1, x]) - ) - if x > 0: - _record_component_boundary_label_numba( - component_labels, component, int(labels[y, x - 1]) - ) - if x + 1 < width: - _record_component_boundary_label_numba( - component_labels, component, int(labels[y, x + 1]) - ) - has_fillable_component = False - for component in range(component_labels.size): - if component_labels[component] > 0: - has_fillable_component = True - break - if not has_fillable_component: - return labels - output = labels.copy() - for y in range(height): - for x in range(width): - component = components[y, x] - if component > 0 and fill_flags[component]: - label = component_labels[component] - if label > 0: - output[y, x] = label - return output - - @njit(cache=True) def _max_label_2d_numba(labels: np.ndarray) -> int: height, width = labels.shape @@ -4421,211 +4297,6 @@ def _max_label_2d_numba(labels: np.ndarray) -> int: return max_label -@njit(cache=True) -def _restore_removed_declump_basins_numba( - pre_declump_labels: np.ndarray, - labels_before_size_filter: np.ndarray, - labels_after_size_filter: np.ndarray, -) -> np.ndarray: - height, width = labels_before_size_filter.shape - output = labels_after_size_filter.copy() - max_pre_declump_label = _max_label_2d_numba(pre_declump_labels) - component_surviving_label = np.zeros(max_pre_declump_label + 1, dtype=np.int64) - for y in range(height): - for x in range(width): - pre_label = int(pre_declump_labels[y, x]) - after_label = int(labels_after_size_filter[y, x]) - if pre_label <= 0 or after_label <= 0: - continue - current_label = component_surviving_label[pre_label] - if current_label == 0: - component_surviving_label[pre_label] = after_label - elif current_label != after_label: - component_surviving_label[pre_label] = -1 - visited = np.zeros((height, width), dtype=np.bool_) - stack_y = np.empty(height * width, dtype=np.int64) - stack_x = np.empty(height * width, dtype=np.int64) - component_y = np.empty(height * width, dtype=np.int64) - component_x = np.empty(height * width, dtype=np.int64) - for start_y in range(height): - for start_x in range(width): - if visited[start_y, start_x]: - continue - if ( - labels_before_size_filter[start_y, start_x] <= 0 - or labels_after_size_filter[start_y, start_x] > 0 - ): - visited[start_y, start_x] = True - continue - stack_size = 1 - stack_y[0] = start_y - stack_x[0] = start_x - visited[start_y, start_x] = True - component_size = 0 - boundary_label = 0 - has_multiple_boundary_labels = False - pre_declump_label = 0 - has_multiple_pre_declump_labels = False - while stack_size: - stack_size -= 1 - y = stack_y[stack_size] - x = stack_x[stack_size] - component_y[component_size] = y - component_x[component_size] = x - component_size += 1 - current_pre_declump_label = int(pre_declump_labels[y, x]) - if current_pre_declump_label <= 0: - has_multiple_pre_declump_labels = True - elif pre_declump_label == 0: - pre_declump_label = current_pre_declump_label - elif pre_declump_label != current_pre_declump_label: - has_multiple_pre_declump_labels = True - for dy in range(-1, 2): - yy = y + dy - if yy < 0 or yy >= height: - continue - for dx in range(-1, 2): - if dy == 0 and dx == 0: - continue - xx = x + dx - if xx < 0 or xx >= width: - continue - neighbor_after = labels_after_size_filter[yy, xx] - if neighbor_after > 0: - if boundary_label == 0: - boundary_label = neighbor_after - elif boundary_label != neighbor_after: - has_multiple_boundary_labels = True - continue - if visited[yy, xx]: - continue - if ( - labels_before_size_filter[yy, xx] > 0 - and labels_after_size_filter[yy, xx] == 0 - ): - visited[yy, xx] = True - stack_y[stack_size] = yy - stack_x[stack_size] = xx - stack_size += 1 - if ( - boundary_label <= 0 - or has_multiple_boundary_labels - or pre_declump_label <= 0 - or has_multiple_pre_declump_labels - ): - continue - surviving_label = component_surviving_label[pre_declump_label] - if surviving_label <= 0 or surviving_label != boundary_label: - continue - for index in range(component_size): - output[component_y[index], component_x[index]] = boundary_label - return output - - -@njit(cache=True) -def _record_component_boundary_label_numba( - component_labels: np.ndarray, component: int, label: int -) -> None: - if label <= 0: - return - current = component_labels[component] - if current == 0: - component_labels[component] = label - elif current != label: - component_labels[component] = -1 - - -@njit(cache=True) -def _fill_labeled_holes_from_components_numba( - labels: np.ndarray, components: np.ndarray, fill_flags: np.ndarray -) -> np.ndarray: - height, width = labels.shape - capacity = height * width - output = labels.copy() - queue_y = np.empty(capacity, dtype=np.int64) - queue_x = np.empty(capacity, dtype=np.int64) - head = 0 - tail = 0 - for y in range(height): - for x in range(width): - component = components[y, x] - if component <= 0 or not fill_flags[component]: - continue - label = _first_adjacent_foreground_label_numba(labels, y, x) - if label == 0: - continue - output[y, x] = label - queue_y[tail] = y - queue_x[tail] = x - tail += 1 - while head < tail: - y = queue_y[head] - x = queue_x[head] - head += 1 - component = components[y, x] - label = output[y, x] - if ( - y > 0 - and components[y - 1, x] == component - and fill_flags[component] - and (output[y - 1, x] == 0) - ): - output[y - 1, x] = label - queue_y[tail] = y - 1 - queue_x[tail] = x - tail += 1 - if ( - y + 1 < height - and components[y + 1, x] == component - and fill_flags[component] - and (output[y + 1, x] == 0) - ): - output[y + 1, x] = label - queue_y[tail] = y + 1 - queue_x[tail] = x - tail += 1 - if ( - x > 0 - and components[y, x - 1] == component - and fill_flags[component] - and (output[y, x - 1] == 0) - ): - output[y, x - 1] = label - queue_y[tail] = y - queue_x[tail] = x - 1 - tail += 1 - if ( - x + 1 < width - and components[y, x + 1] == component - and fill_flags[component] - and (output[y, x + 1] == 0) - ): - output[y, x + 1] = label - queue_y[tail] = y - queue_x[tail] = x + 1 - tail += 1 - return output - - -@njit(cache=True) -def _first_adjacent_foreground_label_numba(labels: np.ndarray, y: int, x: int): - height, width = labels.shape - for dy in range(-1, 2): - neighbor_y = y + dy - if neighbor_y < 0 or neighbor_y >= height: - continue - for dx in range(-1, 2): - if dy == 0 and dx == 0: - continue - neighbor_x = x + dx - if neighbor_x < 0 or neighbor_x >= width: - continue - label = labels[neighbor_y, neighbor_x] - if label != 0: - return label - return labels[y, x] - - @njit(cache=True) def _binary_shrink_2d_numba( image: np.ndarray, @@ -4706,50 +4377,6 @@ def _binary_shrink_2d_numba( return output -@njit(cache=True) -def _seed_points_from_components_numba( - components: np.ndarray, component_count: int -) -> np.ndarray: - height, width = components.shape - counts = np.zeros(component_count + 1, dtype=np.int64) - sum_y = np.zeros(component_count + 1, dtype=np.float64) - sum_x = np.zeros(component_count + 1, dtype=np.float64) - for y in range(height): - for x in range(width): - component = components[y, x] - if component <= 0: - continue - counts[component] += 1 - sum_y[component] += y - sum_x[component] += x - best_distance = np.empty(component_count + 1, dtype=np.float64) - best_y = np.full(component_count + 1, -1, dtype=np.int64) - best_x = np.full(component_count + 1, -1, dtype=np.int64) - for component in range(component_count + 1): - best_distance[component] = np.inf - for y in range(height): - for x in range(width): - component = components[y, x] - if component <= 0: - continue - centroid_y = sum_y[component] / counts[component] - centroid_x = sum_x[component] / counts[component] - dy = y - centroid_y - dx = x - centroid_x - distance = dy * dy + dx * dx - if distance < best_distance[component]: - best_distance[component] = distance - best_y[component] = y - best_x[component] = x - seeds = np.zeros((height, width), dtype=np.bool_) - for component in range(1, component_count + 1): - y = best_y[component] - x = best_x[component] - if y >= 0 and x >= 0: - seeds[y, x] = True - return seeds - - @njit(cache=True) def _relabel_sequential_flat_numba(labels: np.ndarray) -> tuple[np.ndarray, int]: max_label = 0 @@ -4785,15 +4412,6 @@ def _relabel_sequential_numba(labels: np.ndarray) -> tuple[np.ndarray, int]: return (flat_output.reshape((height, width)), count) -@njit(cache=True) -def _relabel_sequential_3d_numba(labels: np.ndarray) -> tuple[np.ndarray, int]: - plane_count, height, width = labels.shape - flat_output, count = _relabel_sequential_flat_numba( - labels.reshape(plane_count * height * width) - ) - return (flat_output.reshape((plane_count, height, width)), count) - - class ExpandShrinkOperationStrategy( EnumKeyedStrategyMixin[ExpandShrinkMode], ABC, metaclass=AutoRegisterMeta ): diff --git a/openhcs/processing/backends/cellprofiler/object_filtering.py b/openhcs/processing/backends/cellprofiler/object_filtering.py index 14985b2454..8a85872bc0 100644 --- a/openhcs/processing/backends/cellprofiler/object_filtering.py +++ b/openhcs/processing/backends/cellprofiler/object_filtering.py @@ -23,7 +23,10 @@ from numba import njit from python_introspect import public_names_from_objects from openhcs.core.memory.decorators import numpy -from openhcs.core.callable_contract import KeywordRuntimeParameter +from openhcs.core.callable_contract import ( + KeywordRuntimeParameter, + processing_prepare, +) from openhcs.core.pipeline.function_contracts import ( runtime_bound_parameters, special_inputs, @@ -2377,6 +2380,20 @@ def filter_border_objects( ) +@processing_prepare(filter_objects) +def _prepare_filter_objects() -> None: + """Compile the most-overlap parent assignment before timed execution.""" + ParentWithMostOverlapAssignmentStrategy().indexes_to_keep( + PerObjectAssignmentRequest( + child_labels=np.array([[1, 1, 2], [0, 2, 2]], dtype=np.int32), + enclosing_labels=np.array([[1, 1, 1], [0, 2, 2]], dtype=np.int32), + measurement_values=np.array([0.0, 1.0, 2.0]), + child_count=2, + keep_max=True, + ) + ) + + __all__ = public_names_from_objects( FilterMethod, FilterMode, diff --git a/openhcs/processing/backends/cellprofiler/primary_objects.py b/openhcs/processing/backends/cellprofiler/primary_objects.py index 311e8b1392..d72f32b17a 100644 --- a/openhcs/processing/backends/cellprofiler/primary_objects.py +++ b/openhcs/processing/backends/cellprofiler/primary_objects.py @@ -91,7 +91,7 @@ from openhcs.processing.backends.cellprofiler.enum_attributes import ( CellProfilerEnumAttributeMixin, ) -from openhcs.core.callable_contract import CallableContract +from openhcs.core.callable_contract import CallableProjection from openhcs.core.artifacts.specs import ArtifactSpecCollection from openhcs.core.artifacts.kinds import ImageArtifactType, ObjectLabelsArtifactType from openhcs.core.runtime_image_values import ImagePayload @@ -657,9 +657,9 @@ def finalize_artifact_contract_outputs( artifact_outputs=artifact_outputs, ) source = cls.primary_image_inputs( - CallableContract.from_prepared_callable( + CallableProjection.from_callable( cls.require_callable(invocation_key.function_name) - ), + ).declared_contract(), artifact_inputs.specs, ) if len(source) != 1: diff --git a/openhcs/processing/backends/cellprofiler/relationships.py b/openhcs/processing/backends/cellprofiler/relationships.py index 86f7649941..d4f94dbcf6 100644 --- a/openhcs/processing/backends/cellprofiler/relationships.py +++ b/openhcs/processing/backends/cellprofiler/relationships.py @@ -856,6 +856,7 @@ def aggregate_child_feature_name( from openhcs.core.runtime_relationships import ( ObjectRelationship, ) +from openhcs.core.runtime_sparse_labels import SparseIJVLabelRows from openhcs.core.processing_contracts import ( Pure2DContract, ) @@ -1779,6 +1780,18 @@ def prepare_backend(self) -> None: self.label_centers(parent_labels) self.centroid_distances(parent_labels, child_labels, parents_of) self.minimum_distances(parent_labels, child_labels, parents_of) + # Sparse rows are C-ordered, or column-ordered once their slice column + # is dropped; the two endpoints need not agree. + plane_rows = SparseIJVLabelRows.from_dense_labels(child_labels) + layouts = tuple( + np.asarray(rows.as_yx_label_array(), dtype=np.int64) + for rows in (plane_rows, SparseIJVLabelRows.from_slices((plane_rows,))) + ) + for parent_rows in layouts: + for child_rows in layouts: + self.relate_sparse_ijv_children_to_parents( + parent_rows, child_rows, 3, 3 + ) def relate_children_to_parents( self, parent_labels: np.ndarray, child_labels: np.ndarray, child_count: int diff --git a/openhcs/processing/backends/cellprofiler/shape.py b/openhcs/processing/backends/cellprofiler/shape.py index bedea1969a..f42d96acc4 100644 --- a/openhcs/processing/backends/cellprofiler/shape.py +++ b/openhcs/processing/backends/cellprofiler/shape.py @@ -482,6 +482,7 @@ def postprocess_bound_settings( ) from openhcs.processing.backends.cellprofiler.label_geometry import ( feret_diameters_from_labels, + prepare_arccos_kernels, _numpy124_aquicksort_indices, _numpy124_ordered_label_maximum_indices, ) @@ -1018,6 +1019,7 @@ class ObjectSizeShapeKernelPreparation( def execute(self) -> None: """Compile AreaShape paths before benchmark execution.""" + prepare_arccos_kernels() image = np.linspace(0.0, 1.0, 32 * 32, dtype=np.float32).reshape((32, 32)) labels = np.zeros((32, 32), dtype=np.int32) labels[8:24, 8:24] = 1 @@ -1959,64 +1961,6 @@ def _first_scalar(value: object) -> float: return float(array.reshape(-1)[0]) -@njit(cache=True) -def _radius_features_from_distance_image_numba( - labels: np.ndarray, distances: np.ndarray, label_ids: np.ndarray -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - object_count = label_ids.size - max_label = 0 - for i in range(object_count): - label_id = int(label_ids[i]) - if label_id > max_label: - max_label = label_id - counts_by_label = np.zeros(max_label + 1, dtype=np.int64) - sums_by_label = np.zeros(max_label + 1, dtype=np.float64) - max_by_label = np.zeros(max_label + 1, dtype=np.float64) - rows, cols = labels.shape - for row in range(rows): - for col in range(cols): - label = int(labels[row, col]) - if label > 0 and label <= max_label: - value = distances[row, col] - counts_by_label[label] += 1 - sums_by_label[label] += value - if value > max_by_label[label]: - max_by_label[label] = value - offsets = np.zeros(max_label + 2, dtype=np.int64) - for label in range(max_label + 1): - offsets[label + 1] = offsets[label] + counts_by_label[label] - cursor = offsets.copy() - ordered = np.empty(offsets[max_label + 1], dtype=np.float64) - for row in range(rows): - for col in range(cols): - label = int(labels[row, col]) - if label > 0 and label <= max_label: - index = cursor[label] - ordered[index] = distances[row, col] - cursor[label] = index + 1 - max_radius = np.zeros(object_count, dtype=np.float64) - mean_radius = np.zeros(object_count, dtype=np.float64) - median_radius = np.zeros(object_count, dtype=np.float64) - for i in range(object_count): - label = int(label_ids[i]) - if label <= 0 or label > max_label: - continue - count = counts_by_label[label] - if count <= 0: - continue - start = offsets[label] - values = ordered[start : start + count].copy() - values.sort() - max_radius[i] = max_by_label[label] - mean_radius[i] = sums_by_label[label] / count - middle = count // 2 - if count % 2 == 1: - median_radius[i] = values[middle] - else: - median_radius[i] = 0.5 * (values[middle - 1] + values[middle]) - return (max_radius, mean_radius, median_radius) - - @njit(cache=True) def _radius_features_from_labels_numba( labels: np.ndarray, label_ids: np.ndarray @@ -2219,42 +2163,6 @@ def _distance_to_label_edge_numba(labels: np.ndarray) -> np.ndarray: return output -@njit(cache=True) -def _maximum_position_of_labels_numba( - image: np.ndarray, labels: np.ndarray, label_ids: np.ndarray -) -> tuple[np.ndarray, np.ndarray]: - object_count = label_ids.size - max_label = 0 - for index in range(object_count): - label = int(label_ids[index]) - if label > max_label: - max_label = label - best_values = np.full(max_label + 1, -np.inf, dtype=np.float64) - best_y = np.full(max_label + 1, -1, dtype=np.int64) - best_x = np.full(max_label + 1, -1, dtype=np.int64) - seen = np.zeros(max_label + 1, dtype=np.bool_) - height, width = labels.shape - for y in range(height): - for x in range(width): - label = int(labels[y, x]) - if label <= 0 or label > max_label: - continue - value = image[y, x] - if not seen[label] or value > best_values[label]: - seen[label] = True - best_values[label] = value - best_y[label] = y - best_x[label] = x - centers_i = np.zeros(object_count, dtype=np.float64) - centers_j = np.zeros(object_count, dtype=np.float64) - for index in range(object_count): - label = int(label_ids[index]) - if label > 0 and label <= max_label and seen[label]: - centers_i[index] = float(best_y[label]) - centers_j[index] = float(best_x[label]) - return (centers_i, centers_j) - - def _maximum_position_of_labels_scipy_select( image: np.ndarray, labels: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/smoothing.py b/openhcs/processing/backends/cellprofiler/smoothing.py index 4d604bc585..3d20adea58 100644 --- a/openhcs/processing/backends/cellprofiler/smoothing.py +++ b/openhcs/processing/backends/cellprofiler/smoothing.py @@ -88,7 +88,7 @@ class SmoothModule(CellProfilerModule): FlexibleContract, Pure2DContract, ) -from openhcs.core.runtime_image_values import ImagePayload +from openhcs.core.runtime_image_values import ImagePayload, ImagePayloadMetadata class SmoothingMethod(Enum): @@ -986,6 +986,20 @@ def prepare_smooth() -> None: """Compile default Gaussian smoothing before timed execution.""" image = np.linspace(0.0, 1.0, 64 * 64, dtype=np.float32).reshape((64, 64)) smooth.__wrapped__(image) + # The Numba provider keeps float32 planes and masks as they arrive, so the + # read-only and writable forms are separate compiled signatures. + numba_image = np.linspace(0.0, 1.0, 16 * 16, dtype=np.float32).reshape((16, 16)) + numba_mask = np.ones(numba_image.shape, dtype=np.bool_) + numba_mask[0, 0] = False + for image_writeable in (True, False): + numba_image.flags.writeable = image_writeable + smooth.__wrapped__(numba_image, smoothing_backend_provider=NumbaBackendProvider) + for mask_writeable in (True, False): + numba_mask.flags.writeable = mask_writeable + smooth.__wrapped__( + ImagePayloadMetadata().payload_with(numba_image, numba_mask), + smoothing_backend_provider=NumbaBackendProvider, + ) pure_2d_batch_executor(smooth_batch)(smooth) diff --git a/openhcs/processing/backends/cellprofiler/thresholding.py b/openhcs/processing/backends/cellprofiler/thresholding.py index 732700ac09..7e8a6216d8 100644 --- a/openhcs/processing/backends/cellprofiler/thresholding.py +++ b/openhcs/processing/backends/cellprofiler/thresholding.py @@ -1367,6 +1367,10 @@ def prepare_backend(self) -> None: self.mad(values) self.otsu_threshold(values) self.weighted_otsu_threshold(values.astype(np.float32)) + # Repeated intensities take the counted-rank weighted Otsu kernels. + self.weighted_otsu_threshold(np.repeat(values, 2).astype(np.float32)) + # Two separated modes: the histogram minimum exists. + self.minimum_threshold(np.repeat(np.array([0.0, 1.0]), 8)) self.li_threshold(values) self.triangle_threshold(values) self.isodata_threshold(values) diff --git a/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_diagnostics.py b/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_diagnostics.py index 8b3634eece..54d3bfedbc 100644 --- a/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_diagnostics.py +++ b/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_diagnostics.py @@ -97,57 +97,6 @@ def _threshold_weighted_variance_from_sums( return (fg_variance * fg_count + bg_variance * bg_count) / (fg_count + bg_count) -@njit(cache=True) -def _threshold_weighted_variance_unmasked_finite_numba( - image: np.ndarray, - binary_image: np.ndarray, -) -> float: - height, width = image.shape - if height == 0 or width == 0: - return 0.0 - - max_value = image[0, 0] - for y in range(height): - for x in range(width): - value = image[y, x] - if value > max_value: - max_value = value - - minval = max_value / 256.0 - if minval == 0.0: - return 0.0 - - fg_count = 0 - bg_count = 0 - fg_sum = 0.0 - bg_sum = 0.0 - fg_sumsq = 0.0 - bg_sumsq = 0.0 - for y in range(height): - for x in range(width): - value = image[y, x] - if value < minval: - value = minval - log_value = math.log2(value) - if binary_image[y, x]: - fg_count += 1 - fg_sum += log_value - fg_sumsq += log_value * log_value - else: - bg_count += 1 - bg_sum += log_value - bg_sumsq += log_value * log_value - - return _threshold_weighted_variance_from_sums( - fg_count, - bg_count, - fg_sum, - bg_sum, - fg_sumsq, - bg_sumsq, - ) - - @njit(cache=True) def _threshold_diagnostics_unmasked_finite_numba( image: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu.py b/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu.py index 73093e0ccd..a57c6d4b67 100644 --- a/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu.py +++ b/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu.py @@ -15,7 +15,6 @@ _histogram_range_numba, _isodata_threshold_numba, _mean_threshold_numba, - _minimum_cross_entropy_threshold_numba, _minimum_threshold_numba, _multiotsu_three_class_thresholds_numba, _otsu_threshold_numba, diff --git a/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu_histogram_multiotsu.py b/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu_histogram_multiotsu.py index ec386f6b50..787a933ec4 100644 --- a/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu_histogram_multiotsu.py +++ b/openhcs/processing/backends/cellprofiler/thresholding_threshold_numba_otsu_histogram_multiotsu.py @@ -479,64 +479,6 @@ def _multiotsu_interval_score_numba( return np.float32((weighted_index * weighted_index) / probability) -@njit(cache=True) -def _minimum_cross_entropy_threshold_numba( - values: np.ndarray, - bin_count: int, -) -> float: - if values.size == 0: - return 0.0 - if bin_count < 2: - bin_count = 2 - minimum, maximum = _histogram_range_numba(values) - if maximum == minimum: - return float(minimum) - - counts = _histogram_counts_numba(values, bin_count, minimum, maximum) - width = (maximum - minimum) / float(bin_count) - cumulative_count = np.zeros(bin_count, dtype=np.float64) - cumulative_weighted = np.zeros(bin_count, dtype=np.float64) - total_count = 0.0 - total_weighted = 0.0 - for index in range(bin_count): - center = minimum + (float(index) + 0.5) * width - count = float(counts[index]) - total_count += count - total_weighted += count * center - cumulative_count[index] = total_count - cumulative_weighted[index] = total_weighted - - if total_count == 0.0: - return 0.0 - - best_index = 0 - best_cross_entropy = np.inf - for index in range(bin_count - 1): - foreground_count = cumulative_count[index] - background_count = total_count - foreground_count - foreground_weighted = cumulative_weighted[index] - background_weighted = total_weighted - foreground_weighted - if ( - foreground_count <= 0.0 - or background_count <= 0.0 - or foreground_weighted <= 0.0 - or background_weighted <= 0.0 - ): - continue - foreground_mean = foreground_weighted / foreground_count - background_mean = background_weighted / background_count - if foreground_mean <= 0.0 or background_mean <= 0.0: - continue - cross_entropy = -( - foreground_weighted * math.log(foreground_mean) - + background_weighted * math.log(background_mean) - ) - if cross_entropy < best_cross_entropy: - best_cross_entropy = cross_entropy - best_index = index - return minimum + (float(best_index) + 0.5) * width - - @njit(cache=True) def _sauvola_threshold_image_numba( image: np.ndarray, diff --git a/openhcs/processing/backends/cellprofiler/worms.py b/openhcs/processing/backends/cellprofiler/worms.py index d01251bd52..dbc6378f9e 100644 --- a/openhcs/processing/backends/cellprofiler/worms.py +++ b/openhcs/processing/backends/cellprofiler/worms.py @@ -64,7 +64,7 @@ MeasurementsArtifactType, ObjectLabelsArtifactType, ) -from openhcs.core.callable_contract import KeywordRuntimeParameter +from openhcs.core.callable_contract import KeywordRuntimeParameter, processing_prepare from openhcs.core.runtime_array_values import RuntimeArrayData from openhcs.core.measurement_row_materialization import ( ConcatenatedColumnarRows, @@ -3264,3 +3264,11 @@ def rows(self) -> MeasurementSparseColumnarRows: parse_cellprofiler_float, ), ) + + +@processing_prepare(untangle_worms, untangle_worms_with_overlap, untangle_worms_both) +def prepare_untangle_worms() -> None: + """Compile the ordered-thinning kernel before a worm pipeline runs.""" + mask = np.zeros((16, 16), dtype=bool) + mask[4:12, 3:13] = True + skeletonize_worm_mask(mask) diff --git a/openhcs/processing/backends/cellprofiler/zernike.py b/openhcs/processing/backends/cellprofiler/zernike.py index a81eb11e36..2bcbf6c219 100644 --- a/openhcs/processing/backends/cellprofiler/zernike.py +++ b/openhcs/processing/backends/cellprofiler/zernike.py @@ -808,34 +808,6 @@ def intensity_zernike_moments( ) -> IntensityZernikeMoments: """Return intensity-weighted Zernike magnitudes and phases.""" - def intensity_zernike_moments_batch( - self, - images: Sequence[np.ndarray], - labels: np.ndarray, - measured_labels: np.ndarray, - *, - max_order: int, - ) -> tuple[ZernikeMomentIndexes, tuple[tuple[np.ndarray, np.ndarray], ...]]: - """Return intensity Zernikes for images sharing one label geometry.""" - rows: list[tuple[np.ndarray, np.ndarray]] = [] - zernike_numbers: ZernikeMomentIndexes | None = None - for image in images: - image_zernike_numbers, magnitudes, phases = self.intensity_zernike_moments( - image, - labels, - measured_labels, - image_mask=None, - max_order=max_order, - ) - if zernike_numbers is None: - zernike_numbers = image_zernike_numbers - elif zernike_numbers != image_zernike_numbers: - raise ValueError( - "Batched intensity Zernike indexes changed across images." - ) - rows.append((magnitudes, phases)) - return zernike_numbers or (), tuple(rows) - class LegacyFastNumpyShapeZernikeBackendStrategy(ShapeZernikeBackendStrategy): """Default shape-Zernike backend using shared label geometry.""" @@ -1131,106 +1103,6 @@ def intensity_zernike_moments( ).write_if_enabled() return zernike_numbers, magnitudes, phases - def intensity_zernike_moments_batch( - self, - images: Sequence[np.ndarray], - labels: np.ndarray, - measured_labels: np.ndarray, - *, - max_order: int, - ) -> tuple[ZernikeMomentIndexes, tuple[tuple[np.ndarray, np.ndarray], ...]]: - image_arrays = tuple(np.asarray(image, dtype=np.float64) for image in images) - if not image_arrays: - zernike_numbers_array = _zernike_indexes_array(int(max_order)) - return tuple((int(n), int(m)) for n, m in zernike_numbers_array), () - - labels_array = np.asarray(labels, dtype=np.int32) - measured_label_ids = np.asarray(measured_labels, dtype=np.int32) - zernike_numbers_array = _zernike_indexes_array(int(max_order)) - zernike_numbers = tuple((int(n), int(m)) for n, m in zernike_numbers_array) - if measured_label_ids.size == 0 or zernike_numbers_array.size == 0: - empty_rows = tuple( - ( - np.zeros( - (measured_label_ids.size, len(zernike_numbers)), dtype=float - ), - np.zeros( - (measured_label_ids.size, len(zernike_numbers)), dtype=float - ), - ) - for _image in image_arrays - ) - return zernike_numbers, empty_rows - - geometry = self.zernike_label_geometry(labels_array, measured_label_ids) - if geometry.y_coords.size == 0: - empty_rows = tuple( - ( - np.full((measured_label_ids.size, len(zernike_numbers)), np.nan), - np.full((measured_label_ids.size, len(zernike_numbers)), np.nan), - ) - for _image in image_arrays - ) - return zernike_numbers, empty_rows - - centers = geometry.centers - radii = geometry.radii - y_coords = geometry.y_coords.astype(np.int64, copy=False) - x_coords = geometry.x_coords.astype(np.int64, copy=False) - label_values = geometry.label_values - valid = ( - (label_values > 0) - & (label_values <= measured_label_ids.size) - & np.isfinite(radii[label_values - 1]) - & (radii[label_values - 1] > 0) - ) - y_coords = np.ascontiguousarray(y_coords[valid], dtype=np.int64) - x_coords = np.ascontiguousarray(x_coords[valid], dtype=np.int64) - label_values = np.ascontiguousarray(label_values[valid], dtype=np.int32) - if y_coords.size == 0: - empty_rows = tuple( - ( - np.full((measured_label_ids.size, len(zernike_numbers)), np.nan), - np.full((measured_label_ids.size, len(zernike_numbers)), np.nan), - ) - for _image in image_arrays - ) - return zernike_numbers, empty_rows - - coefficients, exponents, term_counts = self.zernike_radial_terms( - zernike_numbers_array - ) - score_context = ( - np.ascontiguousarray(centers, dtype=np.float64), - np.ascontiguousarray(radii, dtype=np.float64), - np.ascontiguousarray(zernike_numbers_array, dtype=np.int64), - coefficients, - exponents, - term_counts, - int(measured_label_ids.size), - ) - image_stack = np.ascontiguousarray(np.stack(image_arrays), dtype=np.float64) - score_started_at = time.perf_counter() - magnitudes, phases = _score_intensity_zernike_moments_batch_numba( - image_stack, - label_values, - y_coords, - x_coords, - score_context, - ) - runtime_profiler.log( - "zernike_intensity_batch_score", - time.perf_counter() - score_started_at, - images=len(image_arrays), - objects=int(measured_label_ids.size), - pixels=int(y_coords.size), - orders=zernike_numbers_array.shape[0], - ) - return zernike_numbers, tuple( - (magnitudes[image_index], phases[image_index]) - for image_index in range(len(image_arrays)) - ) - class NativeNumpyShapeZernikeBackendStrategy(ShapeZernikeBackendStrategy): """CellProfiler-source Zernike backend for parity-sensitive measurements.""" @@ -1239,31 +1111,6 @@ class NativeNumpyShapeZernikeBackendStrategy(ShapeZernikeBackendStrategy): backend_provider = NativeBackendProvider is_default_backend = True - def intensity_zernike_moments_batch( - self, - images: Sequence[np.ndarray], - labels: np.ndarray, - measured_labels: np.ndarray, - *, - max_order: int, - ) -> tuple[ZernikeMomentIndexes, tuple[tuple[np.ndarray, np.ndarray], ...]]: - rows = tuple( - self.intensity_zernike_moments( - image, - labels, - measured_labels, - image_mask=None, - max_order=max_order, - ) - for image in images - ) - if not rows: - indexes = _zernike_indexes_array(int(max_order)) - return tuple((int(n), int(m)) for n, m in indexes), () - return rows[0][0], tuple( - (magnitudes, phases) for _indexes, magnitudes, phases in rows - ) - def shape_zernike_moments( self, labels: np.ndarray, @@ -1668,26 +1515,6 @@ def intensity_zernike_moments( ) -def intensity_zernike_moments_batch( - images: Sequence[np.ndarray], - labels: np.ndarray, - measured_labels: np.ndarray, - *, - max_order: int, - backend_provider: BackendProviderInput = DEFAULT_CELLPROFILER_BACKEND_SELECTION, -) -> tuple[ZernikeMomentIndexes, tuple[tuple[np.ndarray, np.ndarray], ...]]: - """Return intensity Zernikes for images sharing one label geometry.""" - return ShapeZernikeBackendStrategy.for_memory_type( - MemoryType.NUMPY, - backend_provider=backend_provider, - ).intensity_zernike_moments_batch( - images, - labels, - measured_labels, - max_order=max_order, - ) - - def _array_content_key(array: np.ndarray) -> tuple[str, tuple[int, ...], bytes]: contiguous = np.ascontiguousarray(array) digest = hashlib.sha1(contiguous.view(np.uint8)).digest() @@ -1729,7 +1556,6 @@ def _zernike_indexes_array(max_order: int) -> np.ndarray: ZernikeIntensityDebugArrayTrace, ZernikeIntensityDebugTrace, intensity_zernike_moments, - intensity_zernike_moments_batch, shape_zernike_moments, ) @@ -1963,92 +1789,3 @@ def _score_intensity_zernike_moments_direct_numba( ) return magnitudes, phases - - -@njit(cache=True) -def _score_intensity_zernike_moments_batch_numba( - images: np.ndarray, - label_values: np.ndarray, - y_coords: np.ndarray, - x_coords: np.ndarray, - score_context: tuple, -) -> tuple[np.ndarray, np.ndarray]: - ( - centers, - radii, - zernike_numbers, - coefficients, - exponents, - term_counts, - object_count, - ) = score_context - image_count = images.shape[0] - zernike_count = zernike_numbers.shape[0] - real_sums = np.zeros((image_count, object_count, zernike_count), dtype=np.float64) - imag_sums = np.zeros((image_count, object_count, zernike_count), dtype=np.float64) - areas = np.zeros(object_count, dtype=np.float64) - max_order = _zernike_max_order_numba(zernike_numbers) - rho_powers = np.empty(max_order + 1, dtype=np.float64) - cos_by_m = np.empty(max_order + 1, dtype=np.float64) - sin_by_m = np.empty(max_order + 1, dtype=np.float64) - for pixel_index in range(label_values.size): - object_index = label_values[pixel_index] - 1 - if object_index < 0 or object_index >= object_count: - continue - radius = radii[object_index] - if not np.isfinite(radius) or radius <= 0.0: - continue - y = y_coords[pixel_index] - x = x_coords[pixel_index] - areas[object_index] += 1.0 - normalized_y = (y - centers[object_index, 0]) / radius - normalized_x = (x - centers[object_index, 1]) / radius - if not _prepare_zernike_pixel_basis_numba( - normalized_y, - normalized_x, - max_order, - rho_powers, - cos_by_m, - sin_by_m, - ): - continue - for zernike_index in range(zernike_count): - radial = 0.0 - for term_index in range(term_counts[zernike_index]): - radial += ( - coefficients[zernike_index, term_index] - * rho_powers[exponents[zernike_index, term_index]] - ) - m = abs(zernike_numbers[zernike_index, 1]) - real_basis = radial * cos_by_m[m] - imag_basis = radial * sin_by_m[m] - for image_index in range(image_count): - weight = images[image_index, y, x] - real_sums[image_index, object_index, zernike_index] += ( - weight * real_basis - ) - imag_sums[image_index, object_index, zernike_index] += ( - weight * imag_basis - ) - - magnitudes = np.empty((image_count, object_count, zernike_count), dtype=np.float64) - phases = np.empty((image_count, object_count, zernike_count), dtype=np.float64) - for image_index in range(image_count): - for object_index in range(object_count): - area = areas[object_index] - for zernike_index in range(zernike_count): - real_value = real_sums[image_index, object_index, zernike_index] - imag_value = imag_sums[image_index, object_index, zernike_index] - if area <= 0.0: - magnitudes[image_index, object_index, zernike_index] = np.nan - phases[image_index, object_index, zernike_index] = np.nan - else: - magnitudes[image_index, object_index, zernike_index] = ( - np.sqrt(real_value * real_value + imag_value * imag_value) - / area - ) - phases[image_index, object_index, zernike_index] = np.arctan2( - real_value, - imag_value, - ) - return magnitudes, phases diff --git a/openhcs/runtime/zmq_execution_server.py b/openhcs/runtime/zmq_execution_server.py index 1da59560d5..c30990575a 100644 --- a/openhcs/runtime/zmq_execution_server.py +++ b/openhcs/runtime/zmq_execution_server.py @@ -12,6 +12,7 @@ from types import ModuleType from polystore import cleanup_backend_connections +from polystore.base import ensure_storage_registry from zmqruntime.config import TransportMode from zmqruntime.execution import ExecutionServer from zmqruntime.messages import ( @@ -27,10 +28,12 @@ from objectstate.object_state import ObjectState from objectstate.object_state_registry import ObjectStateRegistry from openhcs.core.config_document import ConfigDocumentCodec +from openhcs.core.dataset_sources.discovery import discover_domain_registries from openhcs.core.orchestrator.cancellation import ExecutionCancelledError from openhcs.core.orchestrator.worker_execution import PreparedForkWorkerLaneRunner from openhcs.core.pipeline_document import PipelineDocumentCodec from openhcs.core.progress import ProgressEvent, ProgressExecutionContext +from openhcs.core.runtime_execution_validation import runtime_output_roots from openhcs.core.steps.function_step import FunctionStep from openhcs.runtime.environment_provenance import RuntimeEnvironmentSnapshot from openhcs.runtime.import_authority import ImportRootEndpoint @@ -47,6 +50,7 @@ from openhcs.runtime.zmq_execution_signature import ( OpenHCSExecutionConfigBundle, ZMQAuxiliaryExecutionParams, + ZMQAuxiliaryParamField, ZMQExecutionRequestPayload, ) from openhcs.runtime.zmq_orchestrator_environment import ( @@ -148,8 +152,15 @@ def prepare_runtime_capabilities( self, status_callback: EndpointStartupStatusCallback | None = None, ) -> None: - """Prepare catalog metadata before physical endpoint readiness.""" + """Prepare the catalog and what a first request would otherwise run once. + + The storage backend registry and the family registries that domains + extend are completed by their owners here, so workers forked after + this inherit them and no request pays for them. + """ self._function_catalog_preparation.prepare_before_serving(status_callback) + ensure_storage_registry() + discover_domain_registries() def start(self) -> None: self.prepare_runtime_capabilities() @@ -814,9 +825,6 @@ def _export_runtime_observation( if export_path is None: return - from openhcs.core.runtime_execution_validation import runtime_output_roots - from openhcs.runtime.zmq_execution_signature import ZMQAuxiliaryParamField - execution_bundle = compilation.execution_bundle output_roots = runtime_output_roots( execution_bundle.runtime_contexts, diff --git a/openhcs/runtime/zmq_execution_signature.py b/openhcs/runtime/zmq_execution_signature.py index 0120786f71..6d20cb2e16 100644 --- a/openhcs/runtime/zmq_execution_signature.py +++ b/openhcs/runtime/zmq_execution_signature.py @@ -17,6 +17,11 @@ from openhcs.core.config import GlobalPipelineConfig, PipelineConfig from openhcs.core.orchestrator.execution_result import RuntimeObservationMode +from openhcs.core.runtime_exports import RuntimeExportObservation +from openhcs.runtime.zmq_execution_observation import ( + ZMQRuntimeExecutionObservationExport, + ZMQRuntimeExecutionOutcomeExport, +) if TYPE_CHECKING: from openhcs.core.compiled_execution import CompiledExecutionBundle @@ -89,10 +94,6 @@ class ExportRuntimeValues(ZMQRuntimeObservationExportScope): def export( *, execution_bundle, execution_results, output_roots, server_environment, execution_id ): - from openhcs.runtime.zmq_execution_observation import ( - ZMQRuntimeExecutionObservationExport, - ) - return ZMQRuntimeExecutionObservationExport.from_execution( executed_contexts=execution_results.executed_contexts, execution_results=execution_results, @@ -118,11 +119,6 @@ def check_export_path(cls, export_path: object) -> None: def export( *, execution_bundle, execution_results, output_roots, server_environment, execution_id ): - from openhcs.core.runtime_exports import RuntimeExportObservation - from openhcs.runtime.zmq_execution_observation import ( - ZMQRuntimeExecutionOutcomeExport, - ) - return ZMQRuntimeExecutionOutcomeExport.from_execution( compiled_axis_ids=execution_bundle.runtime_contexts, execution_results=execution_results, diff --git a/openhcs/runtime/zmq_orchestrator_environment.py b/openhcs/runtime/zmq_orchestrator_environment.py index 4f2cb40033..c35f290960 100644 --- a/openhcs/runtime/zmq_orchestrator_environment.py +++ b/openhcs/runtime/zmq_orchestrator_environment.py @@ -5,6 +5,9 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +from polystore.base import reset_memory_backend, storage_registry + +from openhcs.core.dataset_sources.dataset_roots import DatasetRootRule from openhcs.runtime.zmq_execution_signature import ZMQExecutionIdentity if TYPE_CHECKING: @@ -28,8 +31,6 @@ class ZMQOrchestratorEnvironmentRequest(ZMQExecutionIdentity): debug_execution_config: DebugExecutionConfig | None def prepare(self) -> ZMQOrchestratorEnvironment: - from polystore.base import reset_memory_backend, storage_registry - from openhcs.core.debug import DebugExecutionPolicy reset_memory_backend() @@ -45,8 +46,6 @@ def prepare(self) -> ZMQOrchestratorEnvironment: ) def prepared_dataset_root(self, storage_registry) -> str: - from openhcs.core.dataset_sources.dataset_roots import DatasetRootRule - dataset_root = str( self.subject_id if self.execution_subject_id is None else self.execution_subject_id ) diff --git a/tests/integration/helpers/numba_trace/sitecustomize.py b/tests/integration/helpers/numba_trace/sitecustomize.py new file mode 100644 index 0000000000..172951d25f --- /dev/null +++ b/tests/integration/helpers/numba_trace/sitecustomize.py @@ -0,0 +1,130 @@ +"""Count Numba work per process, for ``test_no_jit_after_readiness``. + +Python imports this module at start-up in every process whose ``PYTHONPATH`` +holds this directory. It does nothing unless ``OPENHCS_NUMBA_TRACE_DIR`` is set. +Forked children inherit it. Each process appends JSON lines to ``/.jsonl``: + +- ``compiled``: Numba finished a JIT compilation (its ``numba:compile`` event). +- ``loaded``: Numba looked a kernel up in its disk cache (``hit`` says whether it + was there). A hit after readiness is a first-use cost too; a miss is followed + by a compilation. +- ``ready``: the execution server finished preparing, with the kernels that + have no compiled signature in memory or on disk (``unprepared``). +""" + +import os +import sys +import time + +_DIRECTORY = os.environ.get("OPENHCS_NUMBA_TRACE_DIR") +if _DIRECTORY: + import importlib.abc + import json + + def _record(kind, **data): + data.update(kind=kind, pid=os.getpid(), parent=os.getppid(), time=time.time()) + descriptor = os.open( + os.path.join(_DIRECTORY, f"{os.getpid()}.jsonl"), + os.O_WRONLY | os.O_APPEND | os.O_CREAT, + 0o644, + ) + try: + os.write(descriptor, (json.dumps(data, default=str) + "\n").encode()) + finally: + os.close(descriptor) + + def _kernel_name(py_func): + return f"{getattr(py_func, '__module__', '?')}.{getattr(py_func, '__qualname__', '?')}" + + def _count_cache_lookups(caching): + load = caching.Cache.load_overload + + def load_overload(self, sig, target_context): + found = load(self, sig, target_context) + _record( + "loaded", + kernel=_kernel_name(self._py_func), + signature=str(sig), + hit=found is not None, + ) + return found + + caching.Cache.load_overload = load_overload + + def _count_compilations(event): + class CompilationListener(event.Listener): + def on_start(self, event): + pass + + def on_end(self, event): + dispatcher = event.data.get("dispatcher") + _record( + "compiled", + kernel=_kernel_name(getattr(dispatcher, "py_func", None)), + signature=str(event.data.get("args")), + ) + + event.register("numba:compile", CompilationListener()) + + def _unprepared_kernels(): + """First-party kernels with no signature compiled here or cached on disk. + + A kernel inlined into its callers (``inline="always"``) never has one. + """ + from numba.core.dispatcher import Dispatcher + + kernels = {} + for module in tuple(sys.modules.values()): + for value in tuple(vars(module).values()) if module is not None else (): + if isinstance(value, Dispatcher): + kernels[id(value)] = value + return sorted( + _kernel_name(kernel.py_func) + for kernel in kernels.values() + if kernel.py_func.__module__.partition(".")[0] == "openhcs" + and kernel.targetoptions.get("inline") != "always" + and not kernel.overloads + and not kernel._cache._cache_file._load_index() + ) + + def _mark_readiness(server_module): + server = server_module.ZMQExecutionServer + prepare = server.prepare_runtime_capabilities + + def prepare_runtime_capabilities(self, *args, **kwargs): + result = prepare(self, *args, **kwargs) + _record("ready", unprepared=_unprepared_kernels()) + return result + + server.prepare_runtime_capabilities = prepare_runtime_capabilities + + _AFTER_IMPORT = { + "numba.core.caching": _count_cache_lookups, + "numba.core.event": _count_compilations, + "openhcs.runtime.zmq_execution_server": _mark_readiness, + } + + class _AfterImport(importlib.abc.MetaPathFinder): + """Run a function on a watched module right after it is first imported.""" + + def find_spec(self, fullname, path=None, target=None): + after = _AFTER_IMPORT.get(fullname) + if after is None: + return None + for finder in sys.meta_path: + if finder is self or not hasattr(finder, "find_spec"): + continue + spec = finder.find_spec(fullname, path, target) + if spec is None: + continue + execute = spec.loader.exec_module + + def exec_module(module, execute=execute, after=after): + execute(module) + after(module) + + spec.loader.exec_module = exec_module + return spec + return None + + sys.meta_path.insert(0, _AfterImport()) diff --git a/tests/integration/test_no_jit_after_readiness.py b/tests/integration/test_no_jit_after_readiness.py new file mode 100644 index 0000000000..72be1a02b9 --- /dev/null +++ b/tests/integration/test_no_jit_after_readiness.py @@ -0,0 +1,136 @@ +"""Every Numba kernel is compiled by catalog preparation, never after readiness. + +Benchmarks time pipeline compilation and execution, so neither may contain a +one-time cost. One cold run (empty Numba cache, empty catalog cache) of two +pipelines over a real execution server is counted with Numba's own compile +event, in every process: the server, the workers it forks and the client. + +What this cannot see: a kernel reached only with a signature (dtype, rank, +layout) or on a branch that neither pipeline nor any preparation exercises; the +second test narrows that to signatures, since it fails for any kernel that +preparation leaves with no compiled signature at all. GPU kernels (CuPy, +pyclesperanto program builds) are not counted: the run is CPU-only. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPOSITORY = Path(__file__).resolve().parents[2] +TRACER = Path(__file__).resolve().parent / "helpers" / "numba_trace" +PIPELINES = ( + "tests/integration/test_main.py", + "tests/integration/test_cellprofiler_generated_pipeline.py" + "::test_cppipe_import_executes_through_canonical_zmq_path", +) + + +@pytest.fixture(scope="module") +def cold_run_events(tmp_path_factory: pytest.TempPathFactory) -> list[dict]: + """Run the pipelines from empty caches and return what each process recorded.""" + root = tmp_path_factory.mktemp("cold_run") + trace = root / "trace" + trace.mkdir() + environment = { + **os.environ, + "PYTHONPATH": os.pathsep.join( + (str(TRACER), str(REPOSITORY), os.environ.get("PYTHONPATH", "")) + ), + "OPENHCS_NUMBA_TRACE_DIR": str(trace), + "NUMBA_CACHE_DIR": str(root / "numba"), + "XDG_DATA_HOME": str(root / "data"), + "OPENHCS_CPU_ONLY": "1", + "QT_QPA_PLATFORM": "offscreen", + "DISPLAY": "", + "WAYLAND_DISPLAY": "", + } + completed = subprocess.run( + [ + sys.executable, "-m", "pytest", *PIPELINES, + "--it-backends", "disk", "--it-microscopes", "ImageXpress", + "--it-dims", "3d", "--it-exec-mode", "multiprocessing", + "--it-zmq-mode", "zmq", "--it-visualizers", "none", + "-q", "--tb=short", "-p", "no:cacheprovider", "-p", "no:randomly", + f"--basetemp={root / 'pipelines'}", + ], + cwd=REPOSITORY, + env=environment, + capture_output=True, + text=True, + timeout=1800, + ) + assert completed.returncode == 0, completed.stdout[-4000:] + completed.stderr[-2000:] + return sorted( + ( + json.loads(line) + for path in trace.glob("*.jsonl") + for line in path.read_text().splitlines() + ), + key=lambda event: event["time"], + ) + + +def _readiness_by_server(events: list[dict]) -> dict[int, dict]: + """The first readiness mark of each execution server, by its process id.""" + ready: dict[int, dict] = {} + for event in events: + if event["kind"] == "ready": + ready.setdefault(event["pid"], event) + assert ready, "no execution server reported readiness; the tracer saw no server" + return ready + + +def _is_after_readiness(event: dict, ready: dict[int, dict]) -> bool: + """A server and the children it forks are judged against its own readiness. + + Any other process is a client: it executes nothing, so it has no reason to + compile or load a kernel at any time. + """ + server = ready.get(event["pid"]) or ready.get(event["parent"]) + return server is None or event["time"] > server["time"] + + +def test_no_kernel_is_compiled_or_loaded_after_readiness( + cold_run_events: list[dict], +) -> None: + ready = _readiness_by_server(cold_run_events) + numba_events = [ + event for event in cold_run_events if event["kind"] in ("compiled", "loaded") + ] + assert any( + event["kind"] == "compiled" and not _is_after_readiness(event, ready) + for event in numba_events + ), "preparation compiled nothing, so the counter proves nothing" + + late = [ + f"{event['kind']} pid {event['pid']} {event['kernel']} {event['signature']}" + for event in numba_events + if _is_after_readiness(event, ready) + ] + assert late == [], ( + "Numba worked after the server was ready, or in a client. Extend the " + "preparation of the kernel's owner so catalog preparation compiles it:\n" + + "\n".join(late) + ) + + +def test_preparation_compiles_every_kernel(cold_run_events: list[dict]) -> None: + """Derived from the code, not from a pipeline: every kernel has a signature.""" + unprepared = sorted( + { + kernel + for event in _readiness_by_server(cold_run_events).values() + for kernel in event["unprepared"] + } + ) + assert unprepared == [], ( + "No preparation reaches these kernels, so their first use would compile " + "them. Reach each from its owner's preparation, or delete it if nothing " + "calls it:\n" + "\n".join(unprepared) + ) diff --git a/tests/unit/test_callable_abi_preparation.py b/tests/unit/test_callable_abi_preparation.py index 32b05ae80b..6631bb2d46 100644 --- a/tests/unit/test_callable_abi_preparation.py +++ b/tests/unit/test_callable_abi_preparation.py @@ -21,7 +21,10 @@ attach_callable_contract_metadata, attach_processing_prepare, ) from openhcs.core.function_contract_metadata import FunctionContractAttribute -from openhcs.core.function_patterns import NormalizedFunctionGroup +from openhcs.core.function_patterns import ( + NormalizedFunctionGroup, + normalize_prepared_function_pattern, +) from openhcs.core.function_reference import ImportableFunctionReference from openhcs.core.processing_preparation import CallablePreparation, PreparationCacheBatch from openhcs.core.runtime_image_values import ImagePayloadMetadata @@ -177,8 +180,8 @@ def prepare(owner): raw.__annotations__["image"] = desired_annotations.pop(0) monkeypatch.setattr(CallablePreparation, "prepare", prepare) - first = NormalizedFunctionGroup.from_chain("group", (raw,)).items[0].contract - second = NormalizedFunctionGroup.from_chain("group", (raw,)).items[0].contract + first = next(normalize_prepared_function_pattern([raw]).iter_items()).contract + second = next(normalize_prepared_function_pattern([raw]).iter_items()).contract assert first.canonical_parameter_annotations["image"] is ArrayPayload assert second.canonical_parameter_annotations["image"] is np.ndarray assert FunctionContractAttribute.canonical_signature not in vars(raw) @@ -205,12 +208,24 @@ def test_unprepared_function_reference_compilation_reads_prepared_live_declarati import_identity=CallableImportIdentity.from_callable(raw_numeric), composite_key="test:raw_numeric", ) - contract = NormalizedFunctionGroup.from_chain("group", (reference,)).items[0].contract + contract = next(normalize_prepared_function_pattern([reference]).iter_items()).contract assert contract.func is reference assert contract.canonical_parameter_annotations["image"] is np.ndarray assert fake_preparation.call_count == 1 +def test_lowering_a_pattern_without_compiling_it_prepares_nothing(fake_preparation): + """Authoring and validation read the declaration; they never touch a kernel.""" + contract = NormalizedFunctionGroup.from_chain("group", (raw_numeric,)).items[0].contract + declared = CallableProjection.from_callable(raw_numeric).declared_contract() + assert fake_preparation.call_count == 0 + assert contract.canonical_signature == declared.canonical_signature + assert ( + declared.canonical_signature + == CallableContract.from_prepared_callable(raw_numeric).canonical_signature + ) + + @pytest.mark.parametrize("serializer", [pickle, cloudpickle]) def test_prepared_signature_metadata_namespace_and_reference_transport(serializer, fake_preparation): contract = CallableContract.from_prepared_callable(raw_numeric)