WaveOrder 4.0: five-block config schema, migrate-config, calibrate-tf

benchmarks/runner.py

@@ -42,6 +42,7 @@
 )
 from benchmarks.utils import TimingTree
 from waveorder import phantoms
+from waveorder.config.migrate import MigrationError, detect_version, migrate_3x_dict
 
 _PHANTOM_FUNCTIONS = {
     "single_bead": phantoms.single_bead,
@@ -99,6 +100,47 @@ def _wo_rec_argv(position_path: Path, config_path: Path, recon_path: Path, cli:
     return ["wo", cli, "-i", str(position_path), "-c", str(config_path), "-o", str(recon_path)]
 
 
+MIGRATED_CONFIG_FILENAME = "config_v4.yml"
+
+
+def _migrate_config_for_execution(config_path: Path) -> Path | None:
+    """Migrate a written 3.x ``config.yml`` to 4.0 via ``wo migrate-config``.
+
+    Parameters
+    ----------
+    config_path : Path
+        The case's ``config.yml``, already written to disk.
+
+    Returns
+    -------
+    Path or None
+        Path to the migrated 4.0 configuration, or None when the case
+        should run unmigrated: the configuration is already 4.0, or it is
+        one that intentionally stays 3.x this release (birefringence).
+
+    Raises
+    ------
+    RuntimeError
+        If the configuration is migratable but ``wo migrate-config`` fails.
+    """
+    raw = yaml.safe_load(config_path.read_text())
+    if detect_version(raw) >= 4:
+        return None
+    try:
+        migrate_3x_dict(raw)
+    except MigrationError:
+        return None
+    migrated_path = config_path.with_name(MIGRATED_CONFIG_FILENAME)
+    result = subprocess.run(
+        ["wo", "migrate-config", str(config_path), "-o", str(migrated_path)],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"wo migrate-config failed:\n{result.stderr}")
+    return migrated_path
+
+
 def _run_wo_rec(argv: list[str]) -> None:
     """Invoke ``wo rec`` via ``argv``, raising on non-zero exit."""
     env = {**os.environ, "PYTHONWARNINGS": "ignore::UserWarning,ignore::DeprecationWarning"}
@@ -226,21 +268,33 @@ def _run_and_measure(
     case_dir: Path,
     tt: TimingTree,
     ground_truth: torch.Tensor | None = None,
+    migrate_config: bool = False,
     cli: str = "rec",
 ) -> dict:
     """Write config, shell out to ``wo rec``/``wo tile-stitch``, read output, compute metrics.
 
     Must be called from inside a ``with tt.time("total"):`` block — it
-    adds ``reconstruct`` and ``metrics`` sub-timings.
+    adds ``reconstruct`` and ``metrics`` sub-timings. With
+    ``migrate_config=True`` the 3.x configuration is first converted with
+    ``wo migrate-config`` and the reconstruction runs from the 4.0 file;
+    metrics gain a ``config_migrated`` flag.
     """
     modality = infer_modality(recon_config)
     tf_settings = recon_block(recon_config).get(modality, {}).get("transfer_function", {})
 
     config_path = case_dir / "config.yml"
     config_path.write_text(yaml.dump(recon_config, default_flow_style=False))
 
+    exec_config_path = config_path
+    migrated = False
+    if migrate_config:
+        migrated_path = _migrate_config_for_execution(config_path)
+        if migrated_path is not None:
+            exec_config_path = migrated_path
+            migrated = True
+
     recon_path = case_dir / "reconstruction.zarr"
-    argv = _wo_rec_argv(position_path, config_path, recon_path, cli=cli)
+    argv = _wo_rec_argv(position_path, exec_config_path, recon_path, cli=cli)
     (case_dir / "cli_command.sh").write_text(f"#!/bin/bash\n{shlex.join(argv)}\n")
 
     with tt.time("reconstruct"):
@@ -251,13 +305,16 @@ def _run_and_measure(
         while recon_data.ndim > 3:
             recon_data = recon_data[0]
         NA_det, wavelength, pixel_size = _extract_optics(tf_settings)
-        return compute_metrics(
+        metrics = compute_metrics(
             recon_data,
             NA_det=NA_det,
             wavelength=wavelength,
             pixel_size=pixel_size,
             phantom=ground_truth,
         )
+    if migrate_config:
+        metrics["config_migrated"] = migrated
+    return metrics
 
 
 def _finalize_case(case_dir: Path, tt: TimingTree, metrics: dict, save_all: bool) -> None:
@@ -319,9 +376,16 @@ def _get_by_dotted_path(d: dict, path: str):
 def _load_optimized_parameters(case_dir: Path) -> dict[str, float]:
     """Return optimized ``{param_name: value}`` from config_optimized.yml.
 
-    Raises ValueError if the file is missing or has no transfer_function.
+    In migrated-config mode the reconstruction runs from ``config_v4.yml``,
+    so the optimizer writes ``config_v4_optimized.yml`` instead; both names
+    are accepted. Raises ValueError if the file is missing or has no
+    transfer_function.
     """
     optimized_path = case_dir / "config_optimized.yml"
+    if not optimized_path.exists():
+        migrated_optimized = case_dir / "config_v4_optimized.yml"
+        if migrated_optimized.exists():
+            optimized_path = migrated_optimized
     if not optimized_path.exists():
         raise ValueError(
             f"reference includes parameter.* keys for case '{case_dir.name}' "
@@ -396,6 +460,7 @@ def run_synthetic_case(
     reference: dict[str, ReferenceBound] | None = None,
     simulation: SimulationConfig | None = None,
     recovery: RecoveryConfig | None = None,
+    migrate_config: bool = False,
 ) -> dict:
     """Run a single synthetic benchmark case via ``wo rec``.
 
@@ -414,6 +479,11 @@ def run_synthetic_case(
         function zarrs are deleted and ``simulated.zarr`` /
         ``reconstruction.zarr`` larger than
         :data:`SIZE_LIMIT_BYTES` are deleted after metrics are computed.
+    migrate_config : bool
+        Migrate a 3.x configuration with ``wo migrate-config`` and
+        reconstruct from the 4.0 file (design document section 2
+        regression mode). Configurations that intentionally stay 3.x
+        (birefringence) run unmigrated.
 
     Returns
     -------
@@ -483,6 +553,7 @@ def run_synthetic_case(
                 case_dir=case_dir,
                 tt=tt,
                 ground_truth=ground_truth,
+                migrate_config=migrate_config,
             )
         ref_check = check_reference(case_dir, metrics, reference)
         if ref_check is not None:
@@ -647,6 +718,7 @@ def run_hpc_case(
     save_all: bool = False,
     reference: dict[str, ReferenceBound] | None = None,
     crop: CropConfig | None = None,
+    migrate_config: bool = False,
     cli: str = "rec",
 ) -> dict:
     """Run a single HPC benchmark case on existing data via ``wo rec``.
@@ -666,11 +738,16 @@ def run_hpc_case(
         transfer function zarr is deleted and ``reconstruction.zarr``
         is deleted if it exceeds :data:`SIZE_LIMIT_BYTES` after metrics
         are computed.
-    reference_parameters : dict, optional
-        Per-parameter expected values + tolerances for the optimizer.
+    reference : dict, optional
+        Min/max bounds on metrics and optimizer parameters.
     crop : CropConfig, optional
         If given, crop the input to this bbox before reconstructing.
         Speeds up iteration on large FOVs.
+    migrate_config : bool
+        Migrate a 3.x configuration with ``wo migrate-config`` and
+        reconstruct from the 4.0 file (design document section 2
+        regression mode). Configurations that intentionally stay 3.x
+        (birefringence) run unmigrated.
 
     Returns
     -------
@@ -692,6 +769,7 @@ def run_hpc_case(
             recon_config=recon_config,
             case_dir=case_dir,
             tt=tt,
+            migrate_config=migrate_config,
             cli=cli,
         )
         ref_check = check_reference(case_dir, metrics, reference)

tests/calibrate/test_beads.py

@@ -0,0 +1,146 @@
+"""Bead calibration (prior: beads) recovers signed aberrations."""
+
+import math
+
+import numpy as np
+import pytest
+import torch
+import yaml
+from iohub.ngff import open_ome_zarr
+from iohub.ngff.models import TransformationMeta
+from typer.testing import CliRunner
+
+from waveorder.calibrate.beads import detect_beads, run_bead_calibration, sphere_form_factor
+from waveorder.calibrate.synthesis import fluorescence_otf
+from waveorder.calibrate.tfzarr import read_tf_zarr
+from waveorder.cli.main import app
+from waveorder.config.schema import V4Config
+
+ZYX_SHAPE = (16, 48, 48)
+OPTICS = dict(
+    yx_pixel_size=0.2,
+    z_pixel_size=0.4,
+    wavelength_emission=0.532,
+    index_of_refraction_media=1.3,
+    numerical_aperture_detection=1.2,
+)
+DIAMETER = 0.3
+TRUTH = {"astigmatism_0": 0.15}
+POSITIONS = [(8, 12, 12), (8, 12, 36), (8, 36, 24)]
+
+
+def _bead_config(max_iter: int = 300) -> V4Config:
+    return V4Config(
+        **{
+            "version": 4,
+            "data": {"input_channel_names": ["GFP"], "pixel_size": {"yx": 0.2, "z": 0.4}},
+            "imaging_system": {
+                "contrast": "fluorescence",
+                "index_of_refraction_media": 1.3,
+                "detection": {
+                    "wavelength": 0.532,
+                    "numerical_aperture": 1.2,
+                    "pupil_aberrations": {"astigmatism_0": {"nominal": 0.0, "tolerance": "free"}},
+                },
+            },
+            "object": {"prior": {"beads": {"diameter": DIAMETER}}},
+            "noise": {"gaussian": {"sigma": 1.0}},
+            "solver": {"calibration": {"score": "evidence", "search": [{"gradient": {"max_iter": max_iter}}]}},
+        }
+    )
+
+
+@pytest.fixture(scope="module")
+def bead_volume():
+    """Spheres imaged through an astigmatic PSF with Gaussian noise."""
+    otf = fluorescence_otf(ZYX_SHAPE, zernike_coefficients=TRUTH, **OPTICS)
+    form = sphere_form_factor(ZYX_SHAPE, OPTICS["yx_pixel_size"], OPTICS["z_pixel_size"], DIAMETER)
+
+    freqs = [torch.fft.fftfreq(n) for n in ZYX_SHAPE]
+    grids = torch.meshgrid(*freqs, indexing="ij")
+    object_fft = torch.zeros(ZYX_SHAPE, dtype=torch.complex64)
+    for z, y, x in POSITIONS:
+        phase = grids[0] * z + grids[1] * y + grids[2] * x
+        object_fft = object_fft + 3000.0 * torch.exp(-2j * math.pi * phase)
+    object_fft = object_fft * form
+
+    clean = torch.real(torch.fft.ifftn(object_fft * otf)) + 10.0
+    generator = torch.Generator().manual_seed(3)
+    return (clean + torch.randn(ZYX_SHAPE, generator=generator)).numpy()
+
+
+def test_sphere_form_factor_dc_and_decay():
+    form = sphere_form_factor(ZYX_SHAPE, 0.2, 0.4, DIAMETER)
+    assert float(form[0, 0, 0]) == 1.0
+    assert float(form.abs().max()) <= 1.0
+
+
+def test_detect_beads_finds_the_peaks(bead_volume):
+    positions = detect_beads(torch.as_tensor(bead_volume))
+    found = {tuple(p.tolist()) for p in positions}
+    for truth_position in POSITIONS:
+        assert any(max(abs(a - b) for a, b in zip(p, truth_position)) <= 2 for p in found), (truth_position, found)
+
+
+def test_bead_calibration_recovers_signed_aberration(bead_volume):
+    result = run_bead_calibration(_bead_config(), bead_volume)
+    recovered = dict(zip(result.parameter_names, result.parameters[0]))
+    # The known object breaks the global sign-flip twin: the sign is
+    # identifiable here, unlike the blind evidence path.
+    assert abs(recovered["astigmatism_0"] - TRUTH["astigmatism_0"]) < 0.04, recovered
+    assert result.positions_zyx.shape[0] >= len(POSITIONS)
+    variances = np.diag(result.parameters_covariance[0])
+    assert np.all(np.isfinite(variances)) and np.all(variances > 0)
+
+
+def test_bead_calibration_cli(tmp_path, bead_volume):
+    input_path = tmp_path / "beads.zarr"
+    dataset = open_ome_zarr(input_path, layout="hcs", mode="w", channel_names=["GFP"])
+    position = dataset.create_position("0", "0", "0")
+    position.create_image(
+        "0",
+        np.asarray(bead_volume)[None, None].astype(np.float32),
+        transform=[TransformationMeta(type="scale", scale=[1, 1, 0.4, 0.2, 0.2])],
+    )
+    dataset.close()
+
+    config_path = tmp_path / "beads.yml"
+    config_path.write_text(yaml.safe_dump(_bead_config(max_iter=150).model_dump()))
+
+    output_path = tmp_path / "tf_beads.zarr"
+    result = CliRunner().invoke(
+        app,
+        ["calibrate-tf", "-i", str(input_path / "0" / "0" / "0"), "-c", str(config_path), "-o", str(output_path)],
+    )
+    assert result.exit_code == 0, result.output
+
+    store = read_tf_zarr(output_path)
+    assert store.provenance["kind"] == "beads"
+    assert store.parameter_names == ["astigmatism_0"]
+    assert abs(store.parameters[0, 0] - TRUTH["astigmatism_0"]) < 0.06
+    assert store.parameters_covariance[0, 0, 0] > 0
+
+
+def test_detect_beads_adaptive_relaxation(bead_volume):
+    # A deliberately strict threshold finds nothing at first; the adaptive
+    # fallback relaxes it (with a warning) instead of crashing.
+    scaled = torch.as_tensor(bead_volume)
+    strict = 1e6
+    with pytest.raises(ValueError, match="snr_threshold"):
+        detect_beads(scaled, snr_threshold=strict, adaptive=False)
+    # Moderate SNR: pick a threshold ~1.5x the actual peak SNR so only the
+    # relaxed pass finds the beads.
+    median = scaled.median()
+    mad = (scaled - median).abs().median()
+    peak_snr = float((scaled.max() - median) / (1.4826 * mad))
+    with pytest.warns(UserWarning, match="relaxed"):
+        positions = detect_beads(scaled, snr_threshold=1.5 * peak_snr)
+    assert positions.shape[0] >= 1
+
+
+def test_bead_detection_knobs_from_config(bead_volume):
+    config = _bead_config(max_iter=1)
+    config.object.prior.beads.snr_threshold = 3.0
+    config.object.prior.beads.max_beads = 2
+    result = run_bead_calibration(config, bead_volume)
+    assert result.positions_zyx.shape[0] == 2