diff --git a/.claude/skills/reconstruct-with-nextflow/SKILL.md b/.claude/skills/reconstruct-with-nextflow/SKILL.md index 0a96ad72..6a1a4897 100644 --- a/.claude/skills/reconstruct-with-nextflow/SKILL.md +++ b/.claude/skills/reconstruct-with-nextflow/SKILL.md @@ -3,10 +3,10 @@ name: reconstruct-with-nextflow description: >- Run a mantis-v2 microscope dataset through the biahub Nextflow reconstruction pipeline (nextflow/mantis-v2.nf: flat-field → deskew → reconstruct → - virtual-stain → assemble → project → segment → track → QC) on the Bruno HPC - cluster. Assemble, project, segment, track and QC are optional and selected per + virtual-stain → assemble → project → segment → track → pyramid → QC) on the Bruno + HPC cluster. Assemble, project, segment, track and QC are optional and selected per dataset family: A549 runs all of them, neuromast runs assemble + QC and no - segmentation or tracking. Locates the raw + segmentation or tracking; pyramids are on whenever assemble runs. Locates the raw acquisition under /hpc/instruments/cm.mantis, picks the output project directory, scaffolds configs, launches the run in a tmux session in the foreground so it can be watched, monitors it unattended every 30 minutes @@ -358,8 +358,14 @@ Do not run anything yet. Show the user: | family | steps | |---|---| - | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → segment → track → QC** | - | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → QC** (no projection, segmentation or tracking) | + | A549 / cell-line / organelle | flat-field → deskew → reconstruct → virtual-stain → **assemble → project → segment → track → pyramid → QC** | + | neuromast / zebrafish / dynatrack | flat-field → deskew → reconstruct → virtual-stain → **assemble → pyramid → QC** (no projection, segmentation or tracking) | + + Pyramid is on by default whenever assemble runs: it adds two coarser levels + (3 in total, DCA v0.2) to every image and label image of the assembled and + projected plates, in place, once segmentation and tracking have finished + writing into them. QC is the very last step and waits for the pyramids. + `--pyramid_levels 1` skips them. Write the list out in the plan rather than naming the family, and say which optional steps are being skipped and why. If the user asks to skip anything diff --git a/.claude/skills/reconstruct-with-nextflow/references/caveats.md b/.claude/skills/reconstruct-with-nextflow/references/caveats.md index 16312670..73c6123a 100644 --- a/.claude/skills/reconstruct-with-nextflow/references/caveats.md +++ b/.claude/skills/reconstruct-with-nextflow/references/caveats.md @@ -265,22 +265,22 @@ Assemble is one large SLURM job — the single longest step. `biahub concatenate --resume` (passed by the Nextflow task) makes a late kill recoverable per (t, c) unit. -## 7. Shard geometry and `shards_ratio` - -`shards_ratio` is a ratio over the **chunk** shape -(`shards = chunks * shards_ratio`, elementwise TCZYX), not an absolute shard -shape — changing T without recomputing ZYX silently changes the spatial extent. -Larger T shards mean a larger RAM request and a worse blast radius when one -tears; leave the T ratio alone without a measured reason. Channels cannot be -sharded — keep C at 1. The same two keys exist in `project.yml`: the A549 -template uses one-frame `256 x 256` chunks (the assembled and tracking grid) -and `shards_ratio: [16, 1, 1, 7, 5]`, 16-timepoint full-frame shards; its YX -entries are `ceil(Y/256), ceil(X/256)` and must be recomputed for a different -frame size, exactly as for `concatenate.yml`. Auto-sizing is -[iohub#458](https://github.com/czbiohub-sf/iohub/issues/458); most of this -section goes away when it lands. - -**A T ratio above 1 needs [iohub#460](https://github.com/czbiohub-sf/iohub/pull/460), +## 7. Shard geometry and `shards` + +`shards` in `concatenate.yml` and `project.yml` is a **size target** +(`shards: 5GB`, `shards: 200MB`): iohub builds the largest chunk-aligned shard +under it, filling X, Y and Z first and then timepoints, and never spans +channels ([iohub#461](https://github.com/czbiohub-sf/iohub/pull/461)). Nothing +is recomputed per dataset: the same `5GB` gives 5 timepoints by the full volume +on the A549 plate and 9 on the dynatrack one. `"XY"` / `"XYZ"` (one shard per +plane / volume) and an explicit TCZYX shape in array elements are also accepted. +The old `shards_ratio` (a multiplier over the chunk shape) is refused with a +message naming the new key. A larger target means more timepoints per shard, +hence a larger RAM request (concatenate reads the T extent back from the plate it +initialised to size its workers) and a worse blast radius when one tears; leave +the templates' values alone without a measured reason. + +**Shards spanning several timepoints need [iohub#460](https://github.com/czbiohub-sf/iohub/pull/460), which shipped in iohub 0.3.11** — the floor `pyproject.toml` already requires, so a synced checkout has it. @@ -288,7 +288,7 @@ Before that fix, `concatenate` aborted with a message that points nowhere near t cause — `numpy ... Unable to allocate 651. TiB for an array with shape (1048576, 86, 1664, 1193)`, whose first axis is not T, not a shard, not anything. It is **not** evidence that the shard buffer is too large, and -unsetting `shards_ratio` "fixes" it for the wrong reason. What happens: a +unsetting `shards` "fixes" it for the wrong reason. What happens: a shard-aligned batch writes a whole shard's worth of timepoints in one call, iohub drops the blank ones (`Skipping t=1, c=0 due to all zeros or nans` in the step's `.command.out`), and the surviving indices are left with a *gap* — which @@ -440,3 +440,17 @@ Consequences worth knowing: differ" means nothing here; `corr` and max abs diff are the measures that do. A real virtual-stain regression moves the distribution, not the last bit. Deskew, reconstruct and assemble are fully deterministic. + +## 15. Pyramids follow tracking and precede QC, and are not a resume unit + +`biahub pyramid` rewrites each position's `zarr.json` (the multiscales list and +its `biahub-pyramid` record), as do tracking and imaging-qc, so the pipeline +runs it after tracking and holds QC, the very last step, until it has finished. A run that fails after the pyramid step +left the plates with 3 levels; a relaunch with `-resume` reuses the cached +pyramid tasks, and a relaunch that recomputes earlier steps rewrites level 0 +only — rerun `biahub pyramid -i /*/*/* --levels 3` by hand in that +case, since the coarser levels are not recomputed by anything else. A plate +that already carries a pyramid of another depth is refused at init +(`already has a pyramid of a different depth`); the fix is iohub's +`delete_pyramid` on each position or matching `--pyramid_levels`. + diff --git a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh index a4c1be97..4d16ff4f 100755 --- a/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh +++ b/.claude/skills/reconstruct-with-nextflow/templates/run_mantis_v2.sh @@ -204,6 +204,11 @@ unset CLAUDECODE # --qc_track_config was removed: qc_project.yaml counts the labels tracking writes # into the projected plate, so passing the old flag errors at launch. # +# Pyramids are not a flag here: the assembled and projected plates get 3 +# resolution levels (images mean-pooled, labels max-pooled; DCA v0.2, biahub#359) +# after tracking by default, and QC runs last, after them. Add +# `--pyramid_levels 1` to skip them. +# # TO SKIP A STEP, DELETE ITS LINE. Do not comment it out: `#` inside a # backslash-continued command does not start a comment line — the continuation # swallows it, every flag below it is dropped including `-resume`, and bash then diff --git a/README.md b/README.md index 058acc53..4593fb09 100644 --- a/README.md +++ b/README.md @@ -99,7 +99,7 @@ Some steps include an `estimate-*` command that writes a YAML of parameters you | | | | `flat-field` | Apply flat-field correction to selected channels | | `flip` | Flip images in a dataset | -| `pyramid` | Create multiscale pyramid levels | +| `pyramid` | Compute multiscale pyramids of images and label images, in place | | `process-with-config` | Apply arbitrary YAML-defined functions to a dataset | | | | | `estimate-deskew` | Estimate deskewing parameters | diff --git a/biahub/cli/main.py b/biahub/cli/main.py index d8d14104..a9ed017b 100644 --- a/biahub/cli/main.py +++ b/biahub/cli/main.py @@ -93,7 +93,7 @@ def format_options(self, ctx, formatter): { "name": "pyramid", "import_path": "biahub.pyramid.pyramid_cli", - "help": "Create pyramid levels for a dataset", + "help": "Compute multiscale pyramids of a plate's images and label images, in place", }, { "name": "apply-transform", diff --git a/biahub/concatenate.py b/biahub/concatenate.py index df36e3ed..f1930f5b 100644 --- a/biahub/concatenate.py +++ b/biahub/concatenate.py @@ -651,7 +651,7 @@ def _resolve_concatenate_inputs( output_metadata = { "shape": (len(input_time_indices), len(all_channel_names)) + tuple(cropped_shape_zyx), "chunks": chunk_size, - "shards_ratio": settings.shards_ratio, + "shards": settings.shards, "version": resolve_ome_zarr_version( all_data_paths[0], settings.output_ome_zarr_version ), @@ -793,7 +793,12 @@ def concatenate( # 16-worker tasks took 7-26 min; 0.15 min/volume budgets 60 min here. T_out, C_out, _, _, _ = prep["output_metadata"]["shape"] _, _, Z, Y, X = prep["shape"] - batch_size = settings.shards_ratio[0] if settings.shards_ratio else 1 + # iohub writes one shard's timepoints per unit, so the shard's T extent is the + # batch a worker holds. Read it back from the plate init just created rather than + # predicting it: the config may give the shards as a size target. + with open_ome_zarr(output_dirpath, layout="hcs", mode="r") as output_plate: + _, first_position = next(output_plate.positions()) + batch_size = (first_position.data.shards or first_position.data.chunks)[0] time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( shape=(max(T_out // batch_size, 1), C_out, Z, Y, X), ram_multiplier=8 * batch_size, diff --git a/biahub/project.py b/biahub/project.py index 977596fb..42aa3874 100644 --- a/biahub/project.py +++ b/biahub/project.py @@ -221,7 +221,9 @@ def _init_output_plate( ), dtype=dtype, chunks=(1, *settings.chunks_czyx) if settings.chunks_czyx else None, - shards_ratio=tuple(settings.shards_ratio) if settings.shards_ratio else None, + shards=tuple(settings.shards) + if isinstance(settings.shards, list) + else settings.shards, metadata_sources=input_plate, metadata_keys=PROVENANCE_METADATA_KEYS, extra_metadata={PROVENANCE_KEY: settings.model_dump(mode="json")}, diff --git a/biahub/pyramid.py b/biahub/pyramid.py index e4aa84b0..31b814eb 100644 --- a/biahub/pyramid.py +++ b/biahub/pyramid.py @@ -1,140 +1,275 @@ -import datetime +"""Compute multiscale pyramids in place, on a plate's images and its label images. + +Every position of the input plate gets ``levels`` resolution levels: level 0 is the +existing array and each further level is 2x coarser along the spatial axes longer than +one voxel, so a projected (Z=1) plate is downsampled in Y and X only. Images are pooled +with ``method`` and the label images under ``labels/`` with ``label_method``, one that +never invents a label id. The work is iohub's ``compute_pyramid``, run once per +position; this module fans it out and records what it did. +""" + +import math from pathlib import Path import click +import numpy as np import submitit from iohub.ngff import open_ome_zarr +from biahub.cli.monitor import monitor_jobs from biahub.cli.parsing import ( + cluster, + init_only, input_position_dirpaths, - local, + monitor, sbatch_filepath, sbatch_to_submitit, ) +from biahub.settings import LABEL_PYRAMID_METHODS, PYRAMID_METHODS, PyramidSettings from biahub.utils.cluster import ( apply_sbatch_overrides, + echo_resources, estimate_resources, get_submitit_cluster, + local_executor_parameters, ) +PROVENANCE_KEY = "biahub-pyramid" + + +def pyramid_dims(shape_zyx: tuple[int, int, int]) -> set[str]: + """Spatial axes a pyramid downsamples: those longer than one voxel. -def pyramid(fov_path: Path, levels: int, method: str) -> None: + A singleton axis cannot be halved, and scaling its metadata would misreport the + voxel size, so the Z=1 plate ``biahub project`` writes is downsampled in Y and X + only. Time and channel are never downsampled. """ - Create pyramid levels for a single field of view. + return {name for name, size in zip("zyx", shape_zyx, strict=True) if size > 1} - Delegates to iohub's Position.compute_pyramid() which uses cascade downsampling, - where each level is downsampled from the previous level rather than from level 0. - This avoids aliasing artifacts and chunk boundary issues. - Parameters - ---------- - fov_path : Path - Path to the FOV position directory - levels : int - Total number of resolution levels (including level 0). - E.g., levels=4 creates arrays "0", "1", "2", "3". - method : str - Downsampling method (e.g., 'mean', 'max', 'min') +def _check_levels(position_dirpath: Path, levels: int) -> None: + """Refuse a position whose image or labels already have a different pyramid. + + ``compute_pyramid`` recomputes an existing pyramid of the same depth in place, which + is what a retry wants, but a different depth needs ``delete_pyramid`` first; catch + that before any job is submitted, on the first position. """ - click.echo(f"Computing pyramid for FOV: {fov_path}") + with open_ome_zarr(position_dirpath, mode="r") as position: + found = {"image": len(position.array_keys())} + for name, label in position.labels(): + found[f"labels/{name}"] = len(label.array_keys()) + mismatched = {key: n for key, n in found.items() if n not in (1, levels)} + if mismatched: + raise ValueError( + f"{position_dirpath} already has a pyramid of a different depth than the " + f"{levels} levels requested: {mismatched}. Delete it (iohub's delete_pyramid) " + "or request that depth." + ) - with open_ome_zarr(fov_path, mode="r+") as dataset: - dataset.compute_pyramid(levels=levels, method=method) +def _pyramid_position(position_dirpath: Path, settings: PyramidSettings) -> None: + """Compute the pyramid of one position's image and of each of its label images.""" + click.echo(f"Computing pyramid for {position_dirpath}") + with open_ome_zarr(position_dirpath, mode="r+") as position: + dims = pyramid_dims(position.data.shape[-3:]) + position.compute_pyramid(levels=settings.levels, method=settings.method, dims=dims) + for name, label in position.labels(): + click.echo(f"Computing pyramid for labels/{name} of {position_dirpath}") + label.compute_pyramid( + levels=settings.levels, method=settings.label_method, dims=dims + ) + position.zattrs[PROVENANCE_KEY] = {**settings.model_dump(), "dims": sorted(dims)} -@click.command("pyramid") -@input_position_dirpaths() -@sbatch_filepath() -@local() -@click.option( - "--levels", - "-lv", - type=int, - default=4, - show_default=True, - help="Total number of resolution levels including level 0. E.g., levels=4 creates 0, 1, 2, 3.", -) -@click.option( - "--method", - "-m", - type=click.Choice( - [ - "stride", - "median", - "mode", - "mean", - "min", - "max", - ] - ), - default="mean", - show_default=True, - help="Downsampling method to use.", -) -def pyramid_cli( + +def pyramid( input_position_dirpaths: list[Path], - levels: int = 4, + levels: int = 3, method: str = "mean", - sbatch_filepath: Path | None = None, - local: bool = False, + label_method: str = "max", + sbatch_filepath: str | None = None, + cluster: str = "slurm", + monitor: bool = True, + init_only: bool = False, ) -> None: - """Create multi-scale pyramids for OME-Zarr datasets. + """Compute multiscale pyramids of the given positions, images and labels alike. - Uses cascade downsampling to generate progressively downscaled pyramid levels. - Each level is 2x downsampled from the previous (e.g., levels=4 creates the original - plus 2x, 4x, and 8x downsampled versions as arrays "0", "1", "2", "3"). - - >>> biahub pyramid -i ./data.zarr/*/*/* --levels 4 --local + Parameters + ---------- + input_position_dirpaths : list[Path] + Paths to the position directories; the pyramids are written into them. + levels : int, optional + Total number of resolution levels, counting level 0. + method : str, optional + Downsampling method for the images, one of ``PYRAMID_METHODS``. + label_method : str, optional + Downsampling method for the label images, one of ``LABEL_PYRAMID_METHODS``. + sbatch_filepath : str, optional + Path to the SLURM batch file. + cluster : str, optional + Execution cluster: 'slurm' submits to a Slurm cluster, 'local' runs jobs as + subprocesses on this machine, 'debug' runs jobs in-process in the foreground. + monitor : bool, optional + If True, monitor the submitted jobs. + init_only : bool, optional + Only validate the request against the first position and report the resources + a position needs; skip per-position processing. """ - if levels <= 1: - click.echo("No pyramid levels to create (levels must be > 1).") - return + slurm_out_path = input_position_dirpaths[0].parents[3] / "slurm_output" - # Estimate resources based on first FOV data shape - with open_ome_zarr(input_position_dirpaths[0], mode="r") as dataset: - T, C, Z, Y, X = dataset.data.shape + settings = PyramidSettings(levels=levels, method=method, label_method=label_method) + _check_levels(input_position_dirpaths[0], settings.levels) - _, num_cpus, gb_ram = estimate_resources(shape=(T, C, Z, Y, X), ram_multiplier=5) + with open_ome_zarr(input_position_dirpaths[0], mode="r") as position: + T, C, Z, Y, X = position.data.shape + dtype = position.data.dtype + shards = position.data.shards or position.data.chunks + + # RAM follows the SHARD, not the volume: iohub writes a pyramid level one target + # shard at a time and reads the source region feeding it, 8x the target in 3D. + # Measured on one A549 position (67 x 6 x 86 x 1664 x 1193 float32, level-0 + # shards 5 x 1 x 96 x 1792 x 1280 = 4.1 GiB, levels sized to the same bytes): + # 44 GB peak RSS and 41.5 min single-threaded, i.e. ~10x the shard and ~0.1 min + # per (t, c) volume; the projected plate's 0.14 GiB shards peaked at 1.7 GB in 27 s. + # All 54 positions of that plate in the v9 debug run (shards 6 x 1 x 88 x 1792 x + # 1280 = 4.51 GiB): 46.5 GB peak RSS on every position (10.3x the shard), and + # 45.6 / 54.5 / 67.3 min min / median / max, i.e. up to 0.17 min per volume, the + # slowest on the 120 GB cpu-c nodes. The old 12x / 0.15 asked 56 GB / 70 min, + # 4% above the slowest position. 14x the shard (~36% over the peak) and + # 0.3 min/volume (~1.8x the slowest) leave room for slower nodes. + itemsize = np.dtype(dtype).itemsize + shard_gb = math.prod(shards) * itemsize / 2**30 + volume_gb = Z * Y * X * itemsize / 2**30 + time_minutes, num_cpus, gb_ram_per_cpu = estimate_resources( + shape=(T, C, Z, Y, X), + dtype=dtype, + ram_multiplier=max(4.0, 14 * shard_gb / volume_gb / 4), # per CPU, 4 CPUs share it + time_multiplier=0.3, + max_num_cpus=4, + ) + mem_gb = num_cpus * gb_ram_per_cpu + echo_resources(num_cpus, mem_gb, time_minutes) - cluster = get_submitit_cluster(local) + if init_only: + click.echo( + f"Validated {settings.levels}-level pyramids for {len(input_position_dirpaths)} " + "positions" + ) + return slurm_args = { "slurm_job_name": "pyramid", - "slurm_partition": "cpu", + "slurm_mem": f"{mem_gb}G", "slurm_cpus_per_task": num_cpus, - "slurm_mem_per_cpu": f"{gb_ram}G", - "slurm_time": 30, - "slurm_array_parallelism": 100, + "slurm_array_parallelism": 100, # process up to 100 positions at a time + "slurm_time": time_minutes, + "slurm_partition": "preempted", } - # Override with sbatch file parameters if provided if sbatch_filepath: slurm_args = apply_sbatch_overrides(slurm_args, sbatch_to_submitit(sbatch_filepath)) - slurm_out_path = Path("slurm_output") - slurm_out_path.mkdir(exist_ok=True) - - executor = submitit.AutoExecutor(folder=slurm_out_path, cluster=cluster) + resolved_cluster = get_submitit_cluster(cluster=cluster) + click.echo(f"Preparing jobs on cluster='{resolved_cluster}': {slurm_args}") + executor = submitit.AutoExecutor(folder=slurm_out_path, cluster=resolved_cluster) + if resolved_cluster == "local": + # Explicit --sbatch-filepath values (e.g. '#LOCAL --timeout-min') win over + # the computed local executor parameters. + slurm_args = {**local_executor_parameters(time_minutes), **slurm_args} executor.update_parameters(**slurm_args) - click.echo( - f"Submitting {len(input_position_dirpaths)} pyramid jobs with resources: {slurm_args}" - ) - + click.echo("Submitting jobs...") jobs = [] with submitit.helpers.clean_env(), executor.batch(): - for fov_path in input_position_dirpaths: - job = executor.submit(pyramid, fov_path=fov_path, levels=levels, method=method) - jobs.append(job) + for input_position_path in input_position_dirpaths: + jobs.append(executor.submit(_pyramid_position, input_position_path, settings)) job_ids = [job.job_id for job in jobs] - timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") - log_path = slurm_out_path / f"pyramid-jobs_{timestamp}.log" + slurm_out_path.mkdir(exist_ok=True) + log_path = slurm_out_path / "submitit_jobs_ids.log" with log_path.open("w") as log_file: log_file.write("\n".join(job_ids)) + # submitit's DebugExecutor is lazy: .submit() wraps the callable in a DebugJob but + # execution only happens on .wait()/.done()/.result(). On the Nextflow path + # (--cluster debug) run each position in the foreground. + if resolved_cluster == "debug": + for job, path in zip(jobs, input_position_dirpaths, strict=True): + job.wait() + click.echo(f"Pyramid complete: {path}") + return + + if monitor: + monitor_jobs(jobs, input_position_dirpaths) + + +@click.command("pyramid") +@input_position_dirpaths() +@sbatch_filepath() +@cluster() +@monitor() +@init_only() +@click.option( + "--levels", + type=int, + default=3, + show_default=True, + help="Total number of resolution levels, counting level 0: 3 adds 2x and 4x levels.", +) +@click.option( + "--method", + type=click.Choice(PYRAMID_METHODS), + default="mean", + show_default=True, + help="Downsampling method for the images.", +) +@click.option( + "--label-method", + type=click.Choice(LABEL_PYRAMID_METHODS), + default="max", + show_default=True, + help="Downsampling method for the label images under labels/; none invents a label id.", +) +def pyramid_cli( + input_position_dirpaths: list[Path], + sbatch_filepath: str | None = None, + cluster: str = "slurm", + monitor: bool = False, + init_only: bool = False, + levels: int = 3, + method: str = "mean", + label_method: str = "max", +) -> None: + """Compute multiscale pyramids of a plate's images and label images, in place. + + Each level is 2x coarser than the previous along the spatial axes longer than one + voxel, so a Z=1 plate is downsampled in Y and X only. Images use --method, the label + images under labels/ use --label-method. Rerunning recomputes the same levels. + + \b + SLURM fan-out of positions across a whole plate: + >>> biahub pyramid -i ./plate.zarr/*/*/* --levels 3 + + \b + Validate the request and report the per-position resources only: + >>> biahub pyramid --init -i ./plate.zarr/*/*/* --levels 3 + + \b + In-process run of a single position (e.g. from a Nextflow worker): + >>> biahub pyramid --cluster debug -i ./plate.zarr/A/1/0 --levels 3 + """ # noqa: D301 + pyramid( + input_position_dirpaths=input_position_dirpaths, + levels=levels, + method=method, + label_method=label_method, + sbatch_filepath=sbatch_filepath, + cluster=cluster, + monitor=monitor, + init_only=init_only, + ) + if __name__ == "__main__": pyramid_cli() diff --git a/biahub/settings.py b/biahub/settings.py index 9387b48d..08db3298 100644 --- a/biahub/settings.py +++ b/biahub/settings.py @@ -8,6 +8,7 @@ import torch import yaml +from iohub.core.sharding import parse_shard_size from pydantic import ( BaseModel, BeforeValidator, @@ -27,6 +28,44 @@ class MyBaseModel(BaseModel): model_config = ConfigDict(extra="forbid") +# Shard geometry as iohub's `shards=` takes it: a size target such as "5GB" (the largest +# chunk-aligned shard under it, filled X, Y, Z then T; C is never grown), "XY"/"XYZ" (one +# shard per plane / per volume), or an explicit TCZYX shape in array elements. +ShardsSpec = str | list[int] | None + + +def _validate_shards(v: ShardsSpec) -> ShardsSpec: + if v is None: + return v + if isinstance(v, str): + if v.strip().upper() in ("XY", "XYZ"): + return v + try: + parse_shard_size(v) + except ValueError as err: + raise ValueError( + f"shards: {err}. Use a size such as '5GB', 'XY' / 'XYZ', or a TCZYX shape." + ) from err + return v + if len(v) != 5 or any(i < 1 for i in v): + raise ValueError( + "shards must be 5 positive integers (T, C, Z, Y, X), in array elements" + ) + if v[1] != 1: + raise ValueError("shards: iohub does not shard along C; the C entry must be 1") + return v + + +def _reject_shards_ratio(data): + if isinstance(data, dict) and "shards_ratio" in data: + raise ValueError( + "shards_ratio (a multiplier over the chunk shape) was replaced by shards, which " + "needs no chunk arithmetic: a size such as '5GB', 'XY' / 'XYZ', or a TCZYX shape " + "in array elements." + ) + return data + + class DetectPeaksSettings(MyBaseModel): threshold_abs: float = 110 nms_distance: int = 16 @@ -488,11 +527,11 @@ class ProjectionSettings(MyBaseModel): channels: dict[str, ChannelProjection] = Field(min_length=1) focus: ProjectionFocusSettings | None = None - # Output chunk shape (C, Z, Y, X) and shards ratio (T, C, Z, Y, X), as in concatenate: - # a shard is chunks * ratio, elementwise. None takes iohub's defaults (one-frame - # 256 x 256 chunks, one (t, c) frame per shard). Sharding needs OME-Zarr 0.5. + # Output chunk shape (C, Z, Y, X) and shard geometry (see ShardsSpec), as in + # concatenate. None takes iohub's defaults (one-frame 256 x 256 chunks, one (t, c) + # frame per shard). Sharding needs OME-Zarr 0.5. chunks_czyx: list[int] | None = None - shards_ratio: list[int] | None = None + shards: ShardsSpec = None # When None, preserve the OME-Zarr version of the input store. output_ome_zarr_version: Literal["0.4", "0.5"] | None = None @@ -503,16 +542,15 @@ def check_chunk_size(cls, v): raise ValueError("chunks_czyx must be 4 positive integers (C, Z, Y, X)") return v - @field_validator("shards_ratio") + @field_validator("shards") @classmethod - def check_shards_ratio(cls, v): - if v is not None and (len(v) != 5 or any(i < 1 for i in v)): - raise ValueError("shards_ratio must be 5 positive integers (T, C, Z, Y, X)") - if v is not None and v[1] != 1: - raise ValueError( - "shards_ratio: iohub does not shard along C; the C entry must be 1" - ) - return v + def check_shards(cls, v): + return _validate_shards(v) + + @model_validator(mode="before") + @classmethod + def reject_shards_ratio(cls, data): + return _reject_shards_ratio(data) @model_validator(mode="after") def check_focus(self): @@ -534,6 +572,27 @@ def uses_focus(self) -> bool: return any(c.uses_focus for c in self.channels.values()) +PYRAMID_METHODS = ("mean", "median", "mode", "min", "max", "stride") +# Methods that only ever pick a value present in the window, so no level of a +# label pyramid holds an id that level 0 does not (DCA v0.2: "maximum filter or +# decimation filter for label arrays"). +LABEL_PYRAMID_METHODS = ("max", "stride", "mode") + + +class PyramidSettings(MyBaseModel): + """`biahub pyramid`: the pyramid every image and label image of a position gets. + + ``levels`` counts level 0, so ``3`` adds two downsampled levels, each 2x coarser + than the previous along the spatial axes longer than one voxel (Z is left alone + on a Z=1 plate). ``method`` pools the images and ``label_method`` the label + images under ``labels/``. + """ + + levels: Annotated[int, Field(ge=2)] = 3 + method: Literal[PYRAMID_METHODS] = "mean" + label_method: Literal[LABEL_PYRAMID_METHODS] = "max" + + class AffineTransformSettings(MyBaseModel): t_reference: Literal["first", "previous"] = "first" transform_type: Literal["euclidean", "similarity", "affine"] = "euclidean" @@ -725,7 +784,9 @@ class ConcatenateSettings(MyBaseModel): Y_slice: list | list[list | Literal["all"]] | Literal["all"] = "all" Z_slice: list | list[list | Literal["all"]] | Literal["all"] = "all" chunks_czyx: Literal[None] | list[int] = None - shards_ratio: list[int] | None = None + # Shard geometry of the output (see ShardsSpec); None takes iohub's default, one + # shard per (t, c) volume on OME-Zarr 0.5. + shards: ShardsSpec = None ensure_unique_positions: bool | None = False # Concatenate is the migration path into v0.5 stores, so it defaults to # "0.5". Set to None to preserve the input store's OME-Zarr version, or @@ -849,6 +910,16 @@ def check_chunk_size(cls, v): raise ValueError("chunks_czyx must be a list of 4 integers (C, Z, Y, X)") return v + @field_validator("shards") + @classmethod + def check_shards(cls, v): + return _validate_shards(v) + + @model_validator(mode="before") + @classmethod + def reject_shards_ratio(cls, data): + return _reject_shards_ratio(data) + @model_validator(mode="after") def validate_slice_lengths(self): # Per-source lists (channel_names, X/Y/Z_slice) must be one entry per diff --git a/biahub/track.py b/biahub/track.py index 8a3de596..08a742d7 100644 --- a/biahub/track.py +++ b/biahub/track.py @@ -639,9 +639,18 @@ def write_tracking_labels(label_array, labels: ArrayLike, output_mode: str) -> N The label image has the image's axes with one channel and the plate's Z: one plane for a projected plate, where 2D tracking writes its (T, Y, X) labels; the - full stack for 3D tracking, whose (T, Z, Y, X) labels fill it. + full stack for 3D tracking, whose (T, Z, Y, X) labels fill it. The labels are + stored in the label image's dtype (the source label's, see ``label_layout``); + track ids beyond it are refused rather than wrapped. """ - labels = np.asarray(labels, dtype=np.uint32) + labels = np.asarray(labels) + dtype = np.dtype(label_array.dtype) + if labels.size and labels.max() > np.iinfo(dtype).max: + raise ValueError( + f"Track ids reach {int(labels.max())}, more than the label image's {dtype} holds " + f"({np.iinfo(dtype).max}); segment with a wider label dtype." + ) + labels = labels.astype(dtype, copy=False) if output_mode == "2D": # A single-plane series that kept its Z axis. if labels.ndim == 4 and labels.shape[1] == 1: @@ -837,8 +846,9 @@ def _init_labels( The input's dimensionality must match ``settings.output_mode`` (see :func:`check_input_dimensionality`), and with ``settings.labels`` every position must already have that label image (``biahub segment --init`` creates it, so a - pipeline missing its segment step fails here). The tracked labels are (T, Z, Y, X) - with the plate's Z: one plane for a projected plate, the stack for 3D tracking. + pipeline missing its segment step fails here); the tracked labels then take its + dtype, chunks and shards. The tracked label image is (T, 1, Z, Y, X) with the + plate's Z: one plane for a projected plate, the stack for 3D tracking. Positions that already have the output label are left untouched, and the track settings are recorded on each label image. @@ -859,7 +869,9 @@ def _init_labels( "`biahub segment` on the plate first (its model name is the label name)." ) - shape, layout = label_layout(input_position_dirpaths[0]) + # The tracked labels take the chunks and shards of the label image they are + # linked from (segment's), so the two read alike; without one, the defaults. + shape, layout = label_layout(input_position_dirpaths[0], source_label=settings.labels) create_empty_labels( Path(input_position_dirpaths[0]).parents[2], [Path(p).parts[-3:] for p in input_position_dirpaths], diff --git a/biahub/utils/ngff.py b/biahub/utils/ngff.py index 038f11bf..fba70c91 100644 --- a/biahub/utils/ngff.py +++ b/biahub/utils/ngff.py @@ -160,20 +160,27 @@ def labels_exist(position_dirpaths: list[Path], label_name: str) -> bool: def label_layout( position_dirpath: Path, + source_label: str | None = None, ) -> tuple[tuple[int, int, int, int, int], dict[str, Any]]: """Return the image's (T, C, Z, Y, X) shape and the label-image layout for its plate. The layout is the ``label_dtype`` / ``label_chunks`` / ``label_shards`` keyword - arguments of iohub's ``create_empty_labels``: uint32, :data:`LABEL_CHUNKS` and the - :data:`LABEL_SHARDS` size target (no shards on an OME-Zarr 0.4 plate), for a label - of shape (T, 1, Z, Y, X): the image's axes with one channel. + arguments of iohub's ``create_empty_labels``, for a label of shape (T, 1, Z, Y, X): + the image's axes with one channel. With ``source_label``, the dtype, chunks and + shards are those of the position's ``labels/`` level 0, so a label + derived from another (track's from segment's) is stored like it; otherwise uint32, + :data:`LABEL_CHUNKS` and the :data:`LABEL_SHARDS` size target (no shards on an + OME-Zarr 0.4 plate). """ with open_ome_zarr(str(position_dirpath), mode="r") as dataset: T, C, Z, Y, X = dataset.data.shape version = dataset.version - layout = { - "label_dtype": np.uint32, - "label_chunks": label_chunks((T, 1, Z, Y, X)), - "label_shards": LABEL_SHARDS if version == "0.5" else None, - } + if source_label is not None: + source = dataset.get_label(source_label).data + dtype, chunks, shards = np.dtype(source.dtype), tuple(source.chunks), source.shards + shards = tuple(shards) if shards is not None else None + else: + dtype, chunks = np.dtype(np.uint32), label_chunks((T, 1, Z, Y, X)) + shards = LABEL_SHARDS if version == "0.5" else None + layout = {"label_dtype": dtype, "label_chunks": chunks, "label_shards": shards} return (T, C, Z, Y, X), layout diff --git a/nextflow/configs/a549/concatenate.yml b/nextflow/configs/a549/concatenate.yml index 5ab9438a..8377a6ab 100644 --- a/nextflow/configs/a549/concatenate.yml +++ b/nextflow/configs/a549/concatenate.yml @@ -14,11 +14,14 @@ # A range past the end of the shorter positions pads them the same way. time_indices: all channel_names: all -chunks_czyx: [1, 16, 256, 256] +chunks_czyx: [1, 8, 256, 256] # 2 MiB float32 chunks; the pyramid levels keep this shape output_ome_zarr_version: "0.5" -# shards_ratio is a ratio over the CHUNK shape (shards = chunks * ratio, -# elementwise TCZYX), not an absolute shard shape. Leave unset unless you have a -# measured reason: raising the T entry multiplies the shard buffer and the RAM -# request, and one torn shard then loses every timepoint in it. -shards_ratio: [5, 1, 6, 7, 5] # ZYX entries recomputed for THIS output shape: ceil(Z/16), ceil(Y/256), ceil(X/256) +# shards is a size target: iohub builds the largest chunk-aligned shard under it, +# filling X, Y, Z first and then timepoints (C is never sharded). 5GB on this +# plate's 86 x 1664 x 1193 float32 volumes gives shards of 6 timepoints by the +# full volume (88 x 1792 x 1280 in chunks, 4.5 GiB). Do +# not raise it without a measured reason: the T extent multiplies the shard +# buffer and the RAM request, and one torn shard loses every timepoint in it. +# Also accepts "XY" / "XYZ" (one shard per plane / volume) or a TCZYX shape. +shards: 5GB diff --git a/nextflow/configs/a549/project.yml b/nextflow/configs/a549/project.yml index 20c5af85..ad3f4f09 100644 --- a/nextflow/configs/a549/project.yml +++ b/nextflow/configs/a549/project.yml @@ -60,9 +60,10 @@ focus: # plate can follow it. Chunks are one frame, 256 x 256: the same YX grid as the # assembled and tracking stores, so a 256-aligned crop reads one chunk per channel in # every store (the standard's 16-timepoint chunk suits uint8 frames; for float32 it -# would decode 16 frames per single-frame patch). Shards span 16 timepoints and the -# full frame (the standard's floor for the T shard; ~150 MB each here, 30 files per -# position instead of 400). iohub writes one shard's timepoints per unit, so --resume -# is 16 frames coarse and a torn shard loses 16 frames of one channel. +# would decode 16 frames per single-frame patch). Shards are a size target: iohub +# fills the full frame and then as many timepoints as fit, 21 frames of 1792 x 1280 +# float32 (193 MB) here, above the standard's 16-timepoint floor and ~24 files per +# position instead of 400. iohub writes one shard's timepoints per unit, so --resume +# is one shard of frames coarse and a torn shard loses that many frames of one channel. chunks_czyx: [1, 1, 256, 256] -shards_ratio: [16, 1, 1, 7, 5] # YX entries recomputed for THIS output shape: ceil(Y/256), ceil(X/256) +shards: 200MB diff --git a/nextflow/configs/zebrafish/concatenate.yml b/nextflow/configs/zebrafish/concatenate.yml index 3124fe18..76d32dda 100644 --- a/nextflow/configs/zebrafish/concatenate.yml +++ b/nextflow/configs/zebrafish/concatenate.yml @@ -14,11 +14,14 @@ # A range past the end of the shorter positions pads them the same way. time_indices: all channel_names: all -chunks_czyx: [1, 16, 256, 256] +chunks_czyx: [1, 8, 256, 256] # 2 MiB float32 chunks; the pyramid levels keep this shape output_ome_zarr_version: "0.5" -# shards_ratio is a ratio over the CHUNK shape (shards = chunks * ratio, -# elementwise TCZYX), not an absolute shard shape. Leave unset unless you have a -# measured reason: raising the T entry multiplies the shard buffer and the RAM -# request, and one torn shard then loses every timepoint in it. -shards_ratio: [10, 1, 11, 3, 4] +# shards is a size target: iohub builds the largest chunk-aligned shard under it, +# filling X, Y, Z first and then timepoints (C is never sharded). 5GB on a +# 171 x 768 x 794 float32 volume gives shards of 9 timepoints by the full volume +# (4.98 GB; the old 10-timepoint shards were 5.2 GiB, above the DCA 5 GB SHOULD). +# Do not raise it without a measured reason: the T extent multiplies the shard +# buffer and the RAM request, and one torn shard loses every timepoint in it. +# Also accepts "XY" / "XYZ" (one shard per plane / volume) or a TCZYX shape. +shards: 5GB diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf index 7645007c..e7fa6c27 100644 --- a/nextflow/mantis-v2.nf +++ b/nextflow/mantis-v2.nf @@ -23,17 +23,21 @@ nextflow.enable.dsl = 2 // workflow body for why that is possible and what it buys. // // Flat-field → deskew → reconstruct → virtual-stain → assemble → project → -// segment → track → QC is the full chain: assemble concatenates the +// segment → track → pyramid → QC is the full chain: assemble concatenates the // deskew/reconstruct/virtual-stain channels into one plate, project reduces // that 3D plate to the 2D (Z=1) plate 2D segmentation needs, segment writes // cellpose labels into it, track links those labels and writes the tracked -// labels beside them, and QC reads the finished stores. Follow the chaining -// below for the pattern. +// labels beside them, pyramid adds the coarser resolution levels to the +// assembled and projected plates once nothing else writes pixels or labels +// into them, and QC, the very last step, reads the finished stores. Follow the +// chaining below for the pattern. // -// The last five are OPTIONAL and selected by whether their config is given, so -// a run performs the prefix it asks for: A549 wants assemble + project + -// segment + track + QC, a neuromast run wants assemble + QC and no tracking -// (issue #306). A skipped step never renumbers the directories of the ones +// Assemble, project, segment, track and QC are OPTIONAL and selected by +// whether their config is given, so a run performs the prefix it asks for: +// A549 wants assemble + project + segment + track + QC, a neuromast run wants +// assemble + QC and no segmentation or tracking (issue #306). Pyramids are on +// by default for whichever of the two plates exist (--pyramid_levels 1 turns +// them off). A skipped step never renumbers the directories of the ones // around it. // --------------------------------------------------------------------------- @@ -58,6 +62,11 @@ params.max_positions = 0 params.qc_config = null params.qc_project_config = null params.qc_track_config = null // removed; refused below +// Resolution levels of the assembled and projected plates, counting level 0: +// 3 is the DCA v0.2 target (issue #359) and adds a 2x and a 4x level, in Y and X +// (and Z when the plate has more than one plane), to every image and label +// image. 1 turns the step off. Images are mean-pooled, labels max-pooled. +params.pyramid_levels = 3 include { collect_positions; dataset_name; check_environment } from './modules/common' include { flat_field_init_wf; flat_field_run_wf } from './modules/flat_field' @@ -69,6 +78,12 @@ include { project_init_wf; project_run_wf } from './modules/project' include { segment_init_wf; segment_run_wf } from './modules/segment' include { track_init_wf; track_run_wf } from './modules/tracking' include { qc_plan_wf; qc_compute_wf; qc_report_wf; qc_report_spec } from './modules/qc' +// One module, two plates: a process can only be invoked once per workflow +// context, so the assembled and projected plates get their own aliases. +include { pyramid_init_wf as pyramid_assembled_init_wf; + pyramid_run_wf as pyramid_assembled_run_wf } from './modules/pyramid' +include { pyramid_init_wf as pyramid_projected_init_wf; + pyramid_run_wf as pyramid_projected_run_wf } from './modules/pyramid' include { notify_step; notify_run_start; notify_run_end } from './modules/notify' include { cleanup_decision; cleanup_targets; cleanup_intermediates_wf } from './modules/cleanup' @@ -106,8 +121,8 @@ def step_directories(performed) { reconstruct : 'reconstruct', virtual_stain : 'virtual-stain', assemble : 'assemble', - // project writes beside the assembled plate (/_proj.zarr) - // and segment and track write into the plate they read; none has a + // project writes beside the assembled plate (/_proj.zarr), + // and segment, track and pyramid write into the plates they read; none has a // directory of its own or takes a number. ] def layout = [:] @@ -155,6 +170,13 @@ workflow { def qc_image_on = params.qc_config as boolean def qc_project_on = params.qc_project_config as boolean def qc_on = qc_image_on || qc_project_on + // Pyramids go on the deliverable plates, so without assemble there is + // nothing to pyramid and the default is silently moot rather than an error. + def pyramid_levels = params.pyramid_levels as int + if (pyramid_levels < 1) { + error "--pyramid_levels must be at least 1 (1 = no pyramid), got ${params.pyramid_levels}." + } + def pyramid_on = pyramid_levels > 1 && assemble_on // A step cannot outlive the step whose output it reads. Refuse the // combination at launch, naming the config to add or the one to drop, rather @@ -330,6 +352,19 @@ workflow { init_signals << tk_init.done } + if (pyramid_on) { + // Pyramid's init scaffolds nothing: it checks the plate's first position + // against the requested depth and sizes the task. Gated on the init that + // scaffolded each plate; the label images tracking adds later get their + // levels at compute time, from the same depth. + pya_init = pyramid_assembled_init_wf(assemble_output, pyramid_levels, as_init.done) + init_signals << pya_init.done + if (project_on) { + pyp_init = pyramid_projected_init_wf(project_output, pyramid_levels, pj_init.done) + init_signals << pyp_init.done + } + } + // QC's planning verbs belong in this phase for the same reason the step // inits do — `plan-stage` and `estimate-resources` read the store's // structure and validate the config, and read no pixels. Each store's plan @@ -455,12 +490,39 @@ workflow { segment_on ? segment_done.done : plate_written) } + // ----- Pyramids --------------------------------------------------------- + // After the LAST step that writes pixels or labels into each plate: the + // assembled plate's tracking when it writes there (track without project), + // else assemble; the projected plate's tracking, else project. `biahub + // pyramid` rewrites a position's multiscales metadata and adds its own + // record to the position's zarr.json, so it must not overlap another writer + // of that file, and QC (which stamps the tables it wrote into the same + // attrs) is held until it has finished. Level 0 is never touched. + // The last step that writes into each plate: segment and track write into + // the projected plate when project runs, else into the assembled one. + def labels_written = track_on ? track_done.done : (segment_on ? segment_done.done : null) + def assembled_written = (labels_written != null && !project_on) ? labels_written : assemble_done.done + def projected_written = project_on ? (labels_written ?: project_done.done) : null + if (pyramid_on) { + pyramid_assembled_done = pyramid_assembled_run_wf(all_positions, assemble_output, + pyramid_levels, pya_init.resources, + assembled_written) + if (project_on) { + pyramid_projected_done = pyramid_projected_run_wf(all_positions, project_output, + pyramid_levels, pyp_init.resources, + projected_written) + } + } + // ----- QC compute ------------------------------------------------------- - // Planned in the init phase; the compute it planned is gated on the last - // step that writes each store: the ASSEMBLED plate on assemble_done (so its - // pixel QC runs concurrently with projection, segmentation and tracking), - // and the PROJECTED plate on the last of project, segment and track that - // runs (the tracked labels are counted there). + // THE VERY LAST STEP. Planned in the init phase; the compute it planned is + // gated on the last writer of each store: its pyramid when pyramids are on, + // else the last of segment and track where they write into that store (the + // projected plate, or the assembled one without project), else the step + // that created it. QC used to run concurrently with projection, segmentation + // and tracking; it now waits, so that imaging-qc's stamps on the position + // attrs never race another writer of the same zarr.json and the report + // describes the finished plates. // // A QC verdict cannot fail the pipeline: `imaging-qc gate` exits 0 whether // positions pass or fail, recording the verdict in each store's own @@ -479,13 +541,14 @@ workflow { // combined, for the reason spelled out at qc_plan_inputs above. qc_compute_ready = channel.empty() if (qc_image_on) { + def assembled_final = pyramid_on ? pyramid_assembled_done.done : assembled_written qc_compute_ready = qc_compute_ready.mix( - assemble_done.done.map { tuple(assemble_output, params.qc_config) } ) + assembled_final.map { tuple(assemble_output, params.qc_config) } ) } if (qc_project_on) { - def proj_written = track_on ? track_done.done : (segment_on ? segment_done.done : project_done.done) + def projected_final = pyramid_on ? pyramid_projected_done.done : projected_written qc_compute_ready = qc_compute_ready.mix( - proj_written.map { tuple(project_output, params.qc_project_config) } ) + projected_final.map { tuple(project_output, params.qc_project_config) } ) } qc = qc_compute_wf(qc_plan.items, qc_plan.stores, qc_compute_ready) @@ -529,6 +592,10 @@ workflow { if (project_on) final_signals << project_done.done if (segment_on) final_signals << segment_done.done if (track_on) final_signals << track_done.done + if (pyramid_on) { + final_signals << pyramid_assembled_done.done + if (project_on) final_signals << pyramid_projected_done.done + } if (qc_on) final_signals << qc_report.done final_gate = channel.empty() final_signals.each { signal -> final_gate = final_gate.mix(signal) } @@ -576,6 +643,12 @@ workflow { if (project_on) step_events << [label: 'project', done: project_done.done, output: project_output] if (segment_on) step_events << [label: 'segment', done: segment_done.done, output: track_input] if (track_on) step_events << [label: 'track', done: track_done.done, output: track_input] + if (pyramid_on) { + step_events << [label: 'pyramid (assembled)', done: pyramid_assembled_done.done, output: assemble_output] + if (project_on) { + step_events << [label: 'pyramid (projected)', done: pyramid_projected_done.done, output: project_output] + } + } if (qc_on) step_events << [label: 'QC', done: qc_report.done, output: qc_report_dir] if (cleanup_on) step_events << [label: 'cleanup intermediates', done: cleanup_run.done, output: cleanup_record] diff --git a/nextflow/modules/pyramid.nf b/nextflow/modules/pyramid.nf new file mode 100644 index 00000000..dea8c962 --- /dev/null +++ b/nextflow/modules/pyramid.nf @@ -0,0 +1,135 @@ +// Pyramid subworkflows: init (validate + resources) and run (fan-out × N). +// +// This module is PATH-AGNOSTIC. Callers pass the plate and the number of levels +// explicitly; the module has no idea where the plate sits in the pipeline +// directory layout or which step wrote it. mantis-v2.nf includes it twice, +// under aliases, for the assembled and the projected plates. +// +// `biahub pyramid` writes IN PLACE: every position of the plate gets `levels` +// resolution levels (level 0 is the existing array), on its image and on every +// label image under `labels/`, so it must run after the LAST step that writes +// pixels or labels into that plate, and before QC, which stamps the same +// position attrs it rewrites. Level 0 is never touched. Rerunning recomputes +// the same levels, which is what a retried task wants. +// +// INIT AND RUN ARE SEPARATE SUBWORKFLOWS, like every step's (see the INIT PHASE +// comment in mantis-v2.nf). The init here scaffolds nothing: it checks the first +// position against the requested depth (a plate already carrying a pyramid of +// another depth is refused now rather than after tracking) and prints the +// RESOURCES line one position's task is sized from. +// +// run_pyramid uses `--cluster debug` so that submitit's DebugExecutor runs the +// work in-process; Nextflow owns the per-position fan-out. See modules/project.nf, +// the template this follows. + +include { parse_resources; slurm_logs; slurm_log_dir; retry_time; retry_memory } from './common' + + +process init_pyramid { + label 'cpu_local' + + input: + val input_zarr + val levels + val trigger + + output: + stdout + + script: + """ + mkdir -p "${slurm_log_dir('pyramid')}" + biahub pyramid --init \ + -i "${input_zarr}"/*/*/* \ + --levels ${levels} + """ +} + +process run_pyramid { + // Two plates share this process (under the aliased workflows), so the tag + // names the store as well as the position. + tag "${new File(input_zarr as String).name}/${position}" + label 'cpu' + clusterOptions { slurm_logs('pyramid') } + cpus { meta.cpus } + // retry_memory, not a fixed request: an OOM kill (137) doubles the next + // attempt's memory. With the bare init estimate every retry asked for the same + // amount, so one low estimate failed all five attempts on every position and + // aborted the run (v9 of 2026_08_11_A549_SEC61B_DENV: 16 GB, 137 each time). + memory { retry_memory(meta.mem_gb, task) } + time { retry_time(meta.time_minutes, task) } + + input: + tuple val(position), val(meta) + val input_zarr + val levels + + output: + val position + + script: + """ + biahub pyramid --cluster debug \ + -i "${input_zarr}/${position}" \ + --levels ${levels} + """ +} + + +// Validate the request against the plate and size one position's task. +// Metadata-only and cheap, so it belongs in the pipeline's up-front init phase. +// +// take: +// input_zarr path to the plate.zarr the pyramids are written into +// levels total number of resolution levels, counting level 0 +// trigger gating channel — init starts once this emits (the plate exists) +// emit: +// resources the RESOURCES payload sizing one position's task +// done fires once the request has been validated +workflow pyramid_init_wf { + take: + input_zarr + levels + trigger + + main: + init_out = init_pyramid(input_zarr, levels, trigger.collect().map { 'done' }) + + emit: + // Value channels by the contract every step module emits on; see + // flat_field_init_wf for why the guarantee is made at the input. + resources = init_out.map { stdout_text -> parse_resources(stdout_text) } + done = init_out.map { 'done' } +} + + +// Fan out one pyramid task per position. +// +// take: +// positions collected channel of position keys (e.g. ['A/1/0', 'B/1/0']) +// input_zarr path to the plate.zarr the pyramids are written into +// levels total number of resolution levels, counting level 0 +// resources RESOURCES payload from pyramid_init_wf +// prev_done gating channel — the last step writing into the plate has finished +workflow pyramid_run_wf { + take: + positions + input_zarr + levels + resources + prev_done + + main: + // Gate channels are mapped to a token before combining, never combined raw; + // see flat_field_run_wf. + pos_meta = positions + .flatMap { items -> items } + .combine(resources) + .combine(prev_done.map { 'done' }) + .map { pos, meta, _gate -> [pos, meta] } + + py_done = run_pyramid(pos_meta, input_zarr, levels) | collect + + emit: + done = py_done +} diff --git a/pyproject.toml b/pyproject.toml index a3f260ba..39c0ea8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,9 +98,9 @@ biahub = "biahub.cli.main:cli" # uv-dynamic-versioning chain, which pulls in packaging>=24.2. index-strategy = "unsafe-best-match" # waveorder's own [tool.uv.sources] pins iohub to an older main commit -# (b7bc648, the czbiohub-sf/iohub#476 merge) than the one biahub needs below -# (d1f9e89, which adds #479). uv refuses two git URLs for one package, so -# force the whole graph onto biahub's rev. Drop together with the iohub pin. +# (b7bc648, the czbiohub-sf/iohub#476 merge) than the one biahub needs below. +# uv refuses two git URLs for one package, so force the whole graph onto +# biahub's rev. Drop together with the iohub pin. override-dependencies = [ "iohub @ git+https://github.com/czbiohub-sf/iohub@0cb3e35b860b6f774e7125e41d457192911a178e", ] diff --git a/settings/example_concatenate_settings.yml b/settings/example_concatenate_settings.yml index 0fbb61c7..a7cf3c31 100644 --- a/settings/example_concatenate_settings.yml +++ b/settings/example_concatenate_settings.yml @@ -66,11 +66,12 @@ Z_slice: "all" # - [1, 10, 100, 100]: Specify custom chunk sizes chunks_czyx: null -# Number of chunks in a shard for each dimension [T, C, Z, Y, X] -# Options: -# - null: No sharding -# - [1, 1, 4, 8, 8]: Specify custom shards ratio -shards_ratio: null +# Shard geometry (OME-Zarr 0.5 only). Options: +# - null: iohub's default, one shard per (t, c) volume +# - "5GB": the largest chunk-aligned shard under that size, filling X, Y, Z then T +# - "XY" / "XYZ": one shard per plane / per volume +# - [5, 1, 86, 1664, 1193]: an explicit TCZYX shape in array elements (C must be 1) +shards: null # Version of the OME-Zarr format to use for the output # Options: diff --git a/settings/example_project_settings.yml b/settings/example_project_settings.yml index 1d353a49..434eb8b8 100644 --- a/settings/example_project_settings.yml +++ b/settings/example_project_settings.yml @@ -57,14 +57,15 @@ focus: na_det: 1.35 lambda_ill: 0.5 -# Output chunk shape [C, Z, Y, X] and shards ratio [T, C, Z, Y, X], as in concatenate: a -# shard is chunks * ratio, elementwise (the C entry must be 1). null takes iohub's -# defaults: one-frame 256 x 256 chunks, one (t, c) frame per shard. For a timelapse -# consider shards spanning 16 timepoints and the full frame, e.g. [16, 1, 1, 7, 5] for -# 1664 x 1193 frames (ceil(Y/256), ceil(X/256)); see nextflow/configs/a549/project.yml. +# Output chunk shape [C, Z, Y, X] and shard geometry, as in concatenate. null takes +# iohub's defaults: one-frame 256 x 256 chunks, one (t, c) frame per shard. shards is a +# size target such as "200MB" (the largest chunk-aligned shard under it, filling the +# frame and then timepoints; C is never sharded), "XY" / "XYZ" (one shard per plane / +# volume), or an explicit TCZYX shape in array elements. For a timelapse a size target +# gives shards of many timepoints by the full frame; see nextflow/configs/a549/project.yml. # Sharding needs OME-Zarr 0.5. chunks_czyx: null -shards_ratio: null +shards: null # When null, keep the OME-Zarr version of the input store. output_ome_zarr_version: null diff --git a/tests/test_cli/test_project_cli.py b/tests/test_cli/test_project_cli.py index 2ce68724..001c2ed9 100644 --- a/tests/test_cli/test_project_cli.py +++ b/tests/test_cli/test_project_cli.py @@ -255,7 +255,7 @@ def test_project_writes_the_configured_chunks_and_shards(plate, tmp_path): {"GFP": {"method": "max"}, "Phase3D": {"method": "slice", "z": {"center": 1}}}, output_ome_zarr_version="0.5", chunks_czyx=[1, 1, 4, 4], - shards_ratio=[2, 1, 1, 2, 2], # (2, 1, 1, 8, 8): two timepoints per shard + shards=[2, 1, 1, 8, 8], # two timepoints per shard ) position = plate_path / "A" / "1" / "0" diff --git a/tests/test_cli/test_pyramid_cli.py b/tests/test_cli/test_pyramid_cli.py new file mode 100644 index 00000000..02c518bb --- /dev/null +++ b/tests/test_cli/test_pyramid_cli.py @@ -0,0 +1,144 @@ +"""biahub pyramid: multiscale levels on a plate's images and label images, in place.""" + +import json + +import numpy as np +import pytest + +from click.testing import CliRunner +from iohub.ngff import open_ome_zarr + +from biahub.cli.main import cli +from biahub.pyramid import PROVENANCE_KEY, pyramid_dims + +POSITIONS = (("A", "1", "0"), ("B", "1", "0")) + + +@pytest.fixture +def plate(create_custom_plate, tmp_path): + """(T=3, C=3, Z=4, Y=16, X=12) uint16 with a uint32 label image in each position.""" + plate_path, _ = create_custom_plate( + tmp_path, position_list=POSITIONS, z_size=4, y_size=16, x_size=12 + ) + for key in POSITIONS: + with open_ome_zarr(plate_path.joinpath(*key), mode="r+") as position: + labels = np.zeros((3, 1, 4, 16, 12), dtype=np.uint32) # (T, 1, Z, Y, X) + labels[..., 2:10, 1:7] = 5 + labels[..., 12:16, 8:12] = 9 + position.create_label("nuclei", labels) + return plate_path + + +def _run(*args): + return CliRunner().invoke(cli, ["pyramid", *map(str, args)]) + + +def _levels(position_path): + with open_ome_zarr(position_path, mode="r") as position: + label = position.get_label("nuclei") + return ( + {level: position[level].shape for level in position.array_keys()}, + {level: label[level].shape for level in label.array_keys()}, + dict(position.zattrs).get(PROVENANCE_KEY), + ) + + +def test_pyramid_dims_skips_singleton_axes(): + assert pyramid_dims((4, 16, 12)) == {"z", "y", "x"} + assert pyramid_dims((1, 16, 12)) == {"y", "x"} + + +def test_pyramid_writes_image_and_label_levels(plate): + positions = [plate.joinpath(*key) for key in POSITIONS] + + result = _run("--cluster", "debug", "-i", *positions, "--levels", "3") + + assert result.exit_code == 0, result.output + for position_path in positions: + image_levels, label_levels, record = _levels(position_path) + assert image_levels == { + "0": (3, 3, 4, 16, 12), + "1": (3, 3, 2, 8, 6), + "2": (3, 3, 1, 4, 3), + } + assert label_levels == { + "0": (3, 1, 4, 16, 12), + "1": (3, 1, 2, 8, 6), + "2": (3, 1, 1, 4, 3), + } + assert record == { + "levels": 3, + "method": "mean", + "label_method": "max", + "dims": ["x", "y", "z"], + } + with open_ome_zarr(position_path, mode="r") as position: + label = position.get_label("nuclei") + for level in ("1", "2"): + # Max pooling keeps only ids that exist at level 0, and keeps both blobs. + assert set(np.unique(label[level][:])) == {0, 5, 9} + assert position.get_effective_scale(level)[2:] == [2.0 ** int(level)] * 3 + assert label.get_effective_scale(level)[2:] == [2.0 ** int(level)] * 3 + assert position[1][:].mean() > 0 + + +def test_pyramid_z1_plate_is_downsampled_in_yx_only(create_custom_plate, tmp_path): + plate_path, _ = create_custom_plate( + tmp_path, position_list=POSITIONS[:1], z_size=1, y_size=16, x_size=12 + ) + position_path = plate_path.joinpath(*POSITIONS[0]) + with open_ome_zarr(position_path, mode="r+") as position: + position.create_label("nuclei", np.ones((3, 1, 1, 16, 12), dtype=np.uint32)) + + result = _run("--cluster", "debug", "-i", position_path, "--label-method", "stride") + + assert result.exit_code == 0, result.output + image_levels, label_levels, record = _levels(position_path) + assert image_levels["2"] == (3, 3, 1, 4, 3) + assert label_levels["2"] == (3, 1, 1, 4, 3) + assert record["dims"] == ["x", "y"] + assert record["label_method"] == "stride" + with open_ome_zarr(position_path, mode="r") as position: + # Z keeps its size and its scale at every level. + assert position.get_effective_scale("2")[2] == 1.0 + assert position.get_label("nuclei").get_effective_scale("2")[2] == 1.0 + + +def test_pyramid_rerun_recomputes_and_different_depth_is_refused(plate): + position_path = plate.joinpath(*POSITIONS[0]) + assert _run("--cluster", "debug", "-i", position_path).exit_code == 0 + + rerun = _run("--cluster", "debug", "-i", position_path, "--method", "max") + assert rerun.exit_code == 0, rerun.output + assert _levels(position_path)[2]["method"] == "max" + + deeper = _run("--cluster", "debug", "-i", position_path, "--levels", "4") + assert deeper.exit_code != 0 + assert "different depth" in str(deeper.exception) + assert "'image': 3" in str(deeper.exception) + + +def test_pyramid_init_validates_and_reports_resources(plate): + positions = [plate.joinpath(*key) for key in POSITIONS] + + result = _run("--init", "-i", *positions) + + assert result.exit_code == 0, result.output + line = next(ln for ln in result.output.splitlines() if ln.startswith("RESOURCES:")) + assert set(json.loads(line.removeprefix("RESOURCES:"))) == { + "cpus", + "mem_gb", + "time_minutes", + } + assert "Validated 3-level pyramids for 2 positions" in result.output + # Nothing was written. + image_levels, label_levels, record = _levels(positions[0]) + assert list(image_levels) == ["0"] and list(label_levels) == ["0"] and record is None + + +def test_pyramid_refuses_mean_for_labels_and_one_level(plate): + position_path = plate.joinpath(*POSITIONS[0]) + assert _run("--init", "-i", position_path, "--label-method", "mean").exit_code == 2 + one_level = _run("--init", "-i", position_path, "--levels", "1") + assert one_level.exit_code != 0 + assert "levels" in str(one_level.exception) diff --git a/tests/test_cli/test_track_cli.py b/tests/test_cli/test_track_cli.py index 7d67ca84..a6ca4310 100644 --- a/tests/test_cli/test_track_cli.py +++ b/tests/test_cli/test_track_cli.py @@ -622,6 +622,34 @@ def test_label_layout(tmp_path, example_tracking_plate): ) == (16, 1, 1, 1664, 1280) +def test_label_layout_copies_the_source_label(tmp_path, create_custom_plate): + """A derived label (track's from segment's) takes the source's dtype, chunks and shards.""" + plate_path, _ = create_custom_plate(tmp_path, position_list=(("A", "1", "0"),), z_size=1) + position = plate_path / "A" / "1" / "0" + with open_ome_zarr(str(position), mode="r+") as dataset: + dataset.create_label_zeros( + "cells", + shape=(3, 1, 1, 5, 6), + dtype=np.uint16, + chunks=(1, 1, 1, 5, 3), + shards=(3, 1, 1, 5, 6), + ) + _, layout = label_layout(position, source_label="cells") + assert layout == { + "label_dtype": np.uint16, + "label_chunks": (1, 1, 1, 5, 3), + "label_shards": (3, 1, 1, 5, 6), + } + + +def test_write_tracking_labels_keeps_the_label_dtype_and_refuses_overflow(): + label_array = np.zeros((2, 1, 1, 4, 4), dtype=np.uint16) + write_tracking_labels(label_array, np.full((2, 4, 4), 65535, dtype=np.int64), "2D") + assert label_array.dtype == np.uint16 and label_array.max() == 65535 + with pytest.raises(ValueError, match="more than the label image's uint16"): + write_tracking_labels(label_array, np.full((2, 4, 4), 65536, dtype=np.int64), "2D") + + def test_write_tracking_labels_2d_fills_the_single_plane(): label_array = np.zeros((2, 1, 1, 4, 4), dtype=np.uint32) diff --git a/tests/test_concatenate.py b/tests/test_concatenate.py index d1769006..980d9f75 100644 --- a/tests/test_concatenate.py +++ b/tests/test_concatenate.py @@ -490,11 +490,19 @@ def test_concatenate_with_cropping(create_custom_plate, tmp_path, sbatch_file): @pytest.mark.parametrize( - ["version", "shards_ratio_time"], - [["0.4", 1], ["0.5", None], ["0.5", 1], ["0.5", 2], ["0.5", 5]], + ["version", "shards"], + [ + ["0.4", None], + ["0.5", None], + ["0.5", [1, 1, 2, 8, 6]], + ["0.5", [2, 1, 2, 8, 6]], + ["0.5", [5, 1, 2, 8, 6]], # an explicit shape is kept as given, even past T = 3 + ["0.5", "XYZ"], + ["0.5", "1KiB"], # 48 B chunks: the whole volume fits, then timepoints + ], ) def test_concatenate_with_custom_chunks( - create_custom_plate, tmp_path, sbatch_file, version, shards_ratio_time + create_custom_plate, tmp_path, sbatch_file, version, shards ): """ Test concatenating with custom chunk sizes @@ -521,20 +529,20 @@ def test_concatenate_with_custom_chunks( # Define custom chunk sizes chunks = [1, 1, 2, 4, 3] # [C, Z, Y, X] - if version == "0.5": - if shards_ratio_time is None: - shards_ratio = None - else: - shards_ratio = [shards_ratio_time, 1, 1, 2, 2] - elif version == "0.4": - shards_ratio = None + expected_shards = { + None: None, + "XYZ": (1, 1, 4, 8, 6), + "1KiB": (2, 1, 4, 8, 6), # a 384 B volume: two timepoints fit under 1 KiB + }.get(shards if not isinstance(shards, list) else "list") + if isinstance(shards, list): + expected_shards = tuple(shards) # elements, already whole chunks, taken as given settings = ConcatenateSettings( concat_data_paths=[str(plate_1_path) + "/*/*/*", str(plate_2_path) + "/*/*/*"], channel_names=["all", "all"], time_indices="all", chunks_czyx=chunks[1:], - shards_ratio=shards_ratio, + shards=shards, output_ome_zarr_version=version, ) @@ -553,10 +561,8 @@ def test_concatenate_with_custom_chunks( output_plate = open_ome_zarr(output_path) for pos_name, pos in output_plate.positions(): assert pos.data.chunks == tuple(chunks) - if version == "0.5" and shards_ratio is not None: - assert pos.data.shards == tuple( - c * s for c, s in zip(chunks, shards_ratio, strict=True) - ) + if version == "0.5" and shards is not None: + assert pos.data.shards == expected_shards np.testing.assert_array_equal( pos.data.numpy(), np.concatenate( diff --git a/tests/test_example_settings.py b/tests/test_example_settings.py index b0004544..ed99da68 100644 --- a/tests/test_example_settings.py +++ b/tests/test_example_settings.py @@ -245,9 +245,9 @@ def test_nextflow_a549_project_config_validates(): "membrane_prediction", ] assert settings.channels["Phase3D"].method == "slice" - # One-frame 256 x 256 chunks (the assembled/track grid), 16-timepoint full-frame shards. + # One-frame 256 x 256 chunks (the assembled/track grid), ~200 MB full-frame shards. assert settings.chunks_czyx == [1, 1, 256, 256] - assert settings.shards_ratio == [16, 1, 1, 7, 5] + assert settings.shards == "200MB" assert settings.channels["GFP EX488 EM525-45"].method == "max" assert settings.focus.channel == "Phase3D" and settings.focus.per == "timepoint" diff --git a/tests/test_project_settings.py b/tests/test_project_settings.py index 0409c672..5fc8d46a 100644 --- a/tests/test_project_settings.py +++ b/tests/test_project_settings.py @@ -129,15 +129,23 @@ def test_chunks_and_shards_are_validated(): s = ProjectionSettings( channels={"a": {"method": "max"}}, chunks_czyx=[1, 1, 256, 256], - shards_ratio=[16, 1, 1, 7, 5], + shards=[16, 1, 1, 1664, 1193], ) - assert s.chunks_czyx == [1, 1, 256, 256] and s.shards_ratio == [16, 1, 1, 7, 5] + assert s.chunks_czyx == [1, 1, 256, 256] and s.shards == [16, 1, 1, 1664, 1193] + for form in ("200MB", "1.5GiB", "XY", "xyz"): + assert ( + ProjectionSettings(channels={"a": {"method": "max"}}, shards=form).shards == form + ) with pytest.raises(ValueError, match="4 positive integers"): ProjectionSettings(channels={"a": {"method": "max"}}, chunks_czyx=[1, 256, 256]) with pytest.raises(ValueError, match="5 positive integers"): - ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 1, 7, 5]) + ProjectionSettings(channels={"a": {"method": "max"}}, shards=[16, 1, 7, 5]) with pytest.raises(ValueError, match="C entry must be 1"): - ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 2, 1, 7, 5]) + ProjectionSettings(channels={"a": {"method": "max"}}, shards=[16, 2, 1, 7, 5]) + with pytest.raises(ValueError, match="Use a size such as"): + ProjectionSettings(channels={"a": {"method": "max"}}, shards="big") + with pytest.raises(ValueError, match="replaced by shards"): + ProjectionSettings(channels={"a": {"method": "max"}}, shards_ratio=[16, 1, 1, 7, 5]) def test_unknown_keys_are_refused(): diff --git a/tests/test_virtual_stain.py b/tests/test_virtual_stain.py index af6bb427..de112966 100644 --- a/tests/test_virtual_stain.py +++ b/tests/test_virtual_stain.py @@ -20,14 +20,14 @@ def implementation(request): return request.param -def _output_position(tmp_path, shards_ratio=None): +def _output_position(tmp_path, shards=None): store = tmp_path / "output.zarr" create_empty_plate( store_path=store, position_keys=[("A", "1", "0")], channel_names=["Nuclei", "Membrane"], shape=(T, C, Z, Y, X), - shards_ratio=shards_ratio, + shards=shards, version="0.5", ) return store / "A" / "1" / "0" @@ -91,7 +91,7 @@ def test_write_timepoints_recomputes(tmp_path, implementation, resume, token): def test_write_timepoints_batches_by_time_shard(tmp_path, implementation): """With two timepoints per shard, an interruption inside a shard recomputes it whole.""" - position = _output_position(tmp_path, shards_ratio=(2, 1, 1, 1, 1)) + position = _output_position(tmp_path, shards=(2, 1, Z, Y, X)) # two timepoints per shard with pytest.raises(RuntimeError): _write_timepoints( _predictor([], fail_at=1), position, T, resume=True, resume_token="a"