Merge branch 'main' into refactor-track-rename

Bring PR #260 (track --init/--cluster + rename-channels + cellpose) up to date
with main. main had already merged the deskew/flat-field/reconstruct (#255/#257/
#258), cytoland virtual-stain (#267/#271), and stitch-vendor (#238) refactors,
which the refactor-* stack implemented in parallel; those subsystems are
resolved in favour of main.

Conflict resolution:
- Took main's versions of deskew.py, flat_field_correction.py,
  apply_inverse_transfer_function.py, virtual_stain.py (cytoland rewrite,
  supersedes the stack's VisCy --init/--copy), cli/utils.py, and the
  deskew/flat_field/reconstruct/virtual_stain Nextflow modules + tests.
- Took main's nextflow.config (adds the gpu label/queue that tracking.nf's
  run_track needs) and common.nf (JSON resources contract), and mantis-v2.nf.
- pyproject.toml/uv.lock resolved to main (identical after taking iohub>=0.3.8).
- Adopted main's JSON resources contract in track.py: --init now calls
  echo_resources(num_cpus, mem_gb, time_minutes) instead of the old
  space-separated 'RESOURCES:' line, so parse_resources (common.nf) works.
- Preserved PR #260's unique work (tracking.nf, rename_channels.nf, the
  refactored/cellpose track.py, example_track_settings.yml).

Author: tayllatheodoro

biahub/apply_inverse_transfer_function.py

@@ -1,15 +1,10 @@
-import logging
-import math
-
 from pathlib import Path
 
 import click
 import submitit
-import yaml
 
 from iohub.ngff import open_ome_zarr
 from iohub.ngff.utils import create_empty_plate
-from waveorder import sampling
 from waveorder.cli.apply_inverse_transfer_function import (
     apply_inverse_transfer_function_single_position,
     get_reconstruction_output_metadata,
@@ -29,113 +24,39 @@
     sbatch_to_submitit,
 )
 from biahub.cli.utils import (
+    echo_resources,
     get_submitit_cluster,
     yaml_to_model,
 )
 
-logger = logging.getLogger(__name__)
-
-
-def _resolve_config(
-    config_filepath: Path,
-    reference_position_path: Path,
-    output_dirpath: Path,
-) -> Path:
-    """Resolve pixel sizes from zarr metadata into the reconstruction config.
-
-    Reads voxel sizes from the input zarr scale metadata and injects them
-    into the config's transfer_function sections when absent.  Writes the
-    resolved config next to the output zarr as ``reconstruct_resolved.yml``.
-
-    Returns the path to the resolved config.
-    """
-    with open(config_filepath) as f:
-        cfg = yaml.safe_load(f)
-
-    with open_ome_zarr(str(reference_position_path), mode="r") as ds:
-        yx_pixel_size = float(ds.scale[-1])
-        z_pixel_size = float(ds.scale[-3])
-
-    resolved_any = False
-    for section in ("phase", "birefringence", "fluorescence"):
-        tf = (cfg.get(section) or {}).get("transfer_function")
-        if tf is not None:
-            if "yx_pixel_size" not in tf or tf["yx_pixel_size"] is None:
-                tf["yx_pixel_size"] = yx_pixel_size
-                resolved_any = True
-                logger.warning(
-                    "%s.transfer_function.yx_pixel_size not set in config; "
-                    "resolved %.4f from input zarr metadata",
-                    section,
-                    yx_pixel_size,
-                )
-            if "z_pixel_size" not in tf or tf["z_pixel_size"] is None:
-                tf["z_pixel_size"] = z_pixel_size
-                resolved_any = True
-                logger.warning(
-                    "%s.transfer_function.z_pixel_size not set in config; "
-                    "resolved %.4f from input zarr metadata",
-                    section,
-                    z_pixel_size,
-                )
-
-    if not resolved_any:
-        click.echo("Pixel sizes already present in config; no resolution needed")
-
-    resolved_config = output_dirpath.parent / "reconstruct_resolved.yml"
-    resolved_config.parent.mkdir(parents=True, exist_ok=True)
-    with open(resolved_config, "w") as f:
-        yaml.safe_dump(cfg, f, default_flow_style=False, sort_keys=False)
-
-    click.echo(f"Resolved config written to {resolved_config}")
-    return resolved_config
-
-
-def _upsampled_zyx(
-    settings: ReconstructionSettings,
-    zyx_shape: tuple[int, int, int],
-) -> tuple[int, int, int]:
-    """Return the Nyquist-upsampled ZYX shape used during TF computation."""
-    Z, Y, X = zyx_shape
-    if settings.phase is not None:
-        tf = settings.phase.transfer_function
-        trans_nyq = sampling.transverse_nyquist(
-            tf.wavelength_illumination,
-            tf.numerical_aperture_illumination,
-            tf.numerical_aperture_detection,
-        )
-        axial_nyq = sampling.axial_nyquist(
-            tf.wavelength_illumination,
-            tf.numerical_aperture_detection,
-            tf.index_of_refraction_media,
-        )
-        yx_factor = math.ceil(tf.yx_pixel_size / trans_nyq)
-        z_factor = math.ceil(tf.z_pixel_size / axial_nyq)
-        Z, Y, X = Z * z_factor, Y * yx_factor, X * yx_factor
-    return Z, Y, X
-
 
 def _init_output_plate(
     input_position_dirpaths: list[Path],
     output_dirpath: Path,
     config_filepath: Path,
-) -> tuple[tuple[int, int, int, int, int], list[str], Path]:
+    settings: ReconstructionSettings,
+) -> tuple[tuple[int, int, int, int, int], list[str]]:
     """Create the empty reconstruction output plate.
 
-    Resolves pixel sizes from zarr metadata, creates the output store, and
-    copies per-position metadata from the input plate.
+    ``settings`` is the validated ReconstructionSettings loaded by the caller.
+    Output shape/scale/version/channel names are derived by waveorder from the
+    config, and per-position metadata is copied from the input plate.
 
-    Returns (input_shape, channel_names, resolved_config_path).
-    """
-    resolved_config = _resolve_config(
-        config_filepath, input_position_dirpaths[0], output_dirpath
-    )
+    create_empty_plate is idempotent: re-running with the same positions is a
+    no-op, and new positions get appended. Safe to call from both the init
+    orchestrator and from per-position runs.
 
+    Returns the input (T, C, Z, Y, X) shape and output channel names.
+    """
     with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as ds:
         input_shape = ds.data.shape
 
+    # Pixel sizes come from the config (the source of truth). waveorder's
+    # get_reconstruction_output_metadata already emits a PixelSizeMismatchWarning
+    # when the config pixel sizes disagree (>5%) with the input zarr scale, so we
+    # do not warn again here.
     output_metadata = get_reconstruction_output_metadata(
-        input_position_dirpaths[0], resolved_config
+        input_position_dirpaths[0], config_filepath
     )
     output_metadata.pop("plate_metadata", None)
     channel_names = output_metadata["channel_names"]
@@ -148,18 +69,15 @@ def _init_output_plate(
         metadata_sources=input_plate,
     )
 
-    click.echo(f"Created {output_dirpath} ({len(input_position_dirpaths)} positions)")
-    return input_shape, channel_names, resolved_config
+    return input_shape, channel_names
 
 
 def apply_inverse_transfer_function(
     input_position_dirpaths: list[Path],
     transfer_function_dirpath: Path,
     config_filepath: Path,
     output_dirpath: Path,
-    num_processes: int = 1,
     sbatch_filepath: Path | None = None,
-    local: bool = False,
     cluster: str = "slurm",
     monitor: bool = True,
     init_only: bool = False,
@@ -176,12 +94,8 @@ def apply_inverse_transfer_function(
         Path to YAML reconstruction config.
     output_dirpath : Path
         Path to output zarr.
-    num_processes : int
-        Number of parallel processes per position.
     sbatch_filepath : Path, optional
         SBATCH file with slurm parameter overrides.
-    local : bool
-        Legacy flag — use ``cluster='local'`` instead.
     cluster : str
         Execution cluster: 'slurm', 'local', or 'debug'.
     monitor : bool
@@ -192,69 +106,36 @@ def apply_inverse_transfer_function(
     output_dirpath = Path(output_dirpath)
     slurm_out_path = output_dirpath.parent / "slurm_output"
 
-    if init_only:
-        input_shape, _, resolved_config = _init_output_plate(
-            input_position_dirpaths, output_dirpath, config_filepath
-        )
-        T, C, Z, Y, X = input_shape
-        settings = yaml_to_model(resolved_config, ReconstructionSettings)
-
-        num_cpus, mem_per_cpu = wo_estimate_resources(list(input_shape), settings, 1)
-        click.echo(f"RESOURCES:{num_cpus} {num_cpus * mem_per_cpu}")
-
-        uZ, uY, uX = _upsampled_zyx(settings, (Z, Y, X))
-        tf_cpus, tf_mem = wo_estimate_resources([1, 1, uZ, uY, uX], settings, 1)
-        click.echo(f"TF_RESOURCES:{tf_cpus} {tf_cpus * tf_mem}")
-        return
-
-    output_metadata = get_reconstruction_output_metadata(
-        input_position_dirpaths[0], config_filepath
+    settings = yaml_to_model(config_filepath, ReconstructionSettings)
+    input_shape, channel_names = _init_output_plate(
+        input_position_dirpaths, output_dirpath, config_filepath, settings
     )
-    channel_names = output_metadata["channel_names"]
 
-    resolved_cluster = get_submitit_cluster(local=local, cluster=cluster)
+    max_num_cpus = 16
+    num_cpus, mem_per_cpu = wo_estimate_resources(list(input_shape), settings, max_num_cpus)
+    mem_gb = num_cpus * mem_per_cpu
+    time_minutes = 360
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
-    # --cluster debug: apply inverse TF per-position in-process.
-    # Nextflow already handles per-position fan-out and resource scheduling,
-    # so the CLI runs work in-process via submitit's DebugExecutor rather than
-    # submitting its own SLURM jobs.  See also:
-    # examples/submitit_debug_nextflow/2026-05-27-submitit-debug-nextflow-concerns.md
-    if resolved_cluster == "debug":
-        for pos_path in input_position_dirpaths:
-            output_position = output_dirpath / Path(*pos_path.parts[-3:])
-            apply_inverse_transfer_function_single_position(
-                pos_path,
-                transfer_function_dirpath,
-                config_filepath,
-                output_position,
-                num_processes,
-                channel_names,
-            )
-            click.echo(f"Apply-inv-tf done: {'/'.join(pos_path.parts[-3:])}")
+    if init_only:
+        click.echo(
+            f"Created {output_dirpath} ({len(input_position_dirpaths)} positions, "
+            f"{len(settings.output_channel_names)} output channels)"
+        )
         return
 
-    # SLURM / local: create output plate and fan out via submitit
-    output_metadata.pop("plate_metadata", None)
-    input_plate = Path(input_position_dirpaths[0]).parents[2]
-    create_empty_plate(
-        store_path=output_dirpath,
-        position_keys=[p.parts[-3:] for p in input_position_dirpaths],
-        **output_metadata,
-        metadata_sources=input_plate,
-    )
-
-    settings = yaml_to_model(config_filepath, ReconstructionSettings)
-    with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as ds:
-        input_shape = ds.data.shape
-    num_cpus, gb_ram_per_cpu = wo_estimate_resources(
-        list(input_shape), settings, num_processes
-    )
+    resolved_cluster = get_submitit_cluster(cluster=cluster)
 
+    # Fan out one job per position via submitit. With --cluster debug, submitit's
+    # DebugExecutor runs the work in-process (the slurm_* parameters are ignored):
+    # Nextflow already handles per-position fan-out and resource scheduling, so the
+    # CLI must NOT submit its own SLURM jobs. See also:
+    # examples/submitit_debug_nextflow/2026-05-27-submitit-debug-nextflow-concerns.md
     slurm_args = {
         "slurm_job_name": "apply-inverse-transfer-function",
-        "slurm_mem_per_cpu": f"{gb_ram_per_cpu}G",
+        "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
-        "slurm_time": 360,
+        "slurm_time": time_minutes,
         "slurm_partition": "cpu",
     }
 
@@ -276,7 +157,7 @@ def apply_inverse_transfer_function(
                     transfer_function_dirpath,
                     config_filepath,
                     output_dirpath / Path(*pos_path.parts[-3:]),
-                    num_processes,
+                    num_cpus,
                     channel_names,
                 )
             )
@@ -287,6 +168,16 @@ def apply_inverse_transfer_function(
     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 when .wait()/.done()/.result() is called. Run each
+    # one in the foreground and stream progress; monitor's async polling UI is
+    # pointless against synchronous in-process jobs.
+    if resolved_cluster == "debug":
+        for job, path in zip(jobs, input_position_dirpaths, strict=True):
+            job.wait()
+            click.echo(f"Apply-inv-tf complete: {path}")
+        return
+
     if monitor:
         monitor_jobs(jobs, input_position_dirpaths)
 

biahub/cli/utils.py

@@ -1,9 +1,11 @@
+import json
 import logging
 import os
 
 from pathlib import Path
 from typing import Literal
 
+import click
 import numpy as np
 import yaml
 
@@ -14,6 +16,35 @@
 logger = logging.getLogger(__name__)
 
 
+def echo_resources(num_cpus: int, mem_gb: int, time_minutes: int) -> None:
+    """Emit the per-position resource request consumed by the Nextflow pipeline.
+
+    Every step CLI calls this from its ``--init`` path so there is a single
+    source of truth for per-position CPU, memory, and wall-clock time. The
+    Nextflow ``init_*`` process captures this line on stdout and
+    ``parse_resources`` (``nextflow/modules/common.nf``) reads the JSON payload
+    to set the per-position task's ``cpus``/``memory``/``time`` directives. The
+    same values also feed the CLI's own ``slurm_*`` submission args, so the
+    SLURM fan-out and the Nextflow fan-out request identical resources.
+
+    A single JSON payload keeps the contract order-independent and extensible
+    (new fields can be added without breaking the positional parsing).
+
+    Parameters
+    ----------
+    num_cpus : int
+        CPUs per position.
+    mem_gb : int
+        TOTAL memory per position in GB (not per-CPU).
+    time_minutes : int
+        Wall-clock budget per position in minutes.
+    """
+    # Coerce to plain int: estimators may return numpy integers, which json
+    # cannot serialize.
+    payload = {"cpus": int(num_cpus), "mem_gb": int(mem_gb), "time_minutes": int(time_minutes)}
+    click.echo("RESOURCES:" + json.dumps(payload))
+
+
 def get_submitit_cluster(
     local: bool = False,
     cluster: str | None = None,