diff --git a/benchmarks/runner.py b/benchmarks/runner.py index 6ffdec90..0e3ee1f8 100644 --- a/benchmarks/runner.py +++ b/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,12 +268,16 @@ 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", {}) @@ -239,8 +285,16 @@ def _run_and_measure( 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) diff --git a/tests/calibrate/__init__.py b/tests/calibrate/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/calibrate/test_beads.py b/tests/calibrate/test_beads.py new file mode 100644 index 00000000..c54c8d88 --- /dev/null +++ b/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 diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py new file mode 100644 index 00000000..3b599614 --- /dev/null +++ b/tests/calibrate/test_evidence_search.py @@ -0,0 +1,223 @@ +"""Evidence/GCV-scored calibration recovers synthetic aberrations.""" + +import numpy as np +import pytest +import torch + +from waveorder.calibrate.evidence_search import ( + build_parameter_specs, + gcv_score, + run_evidence_calibration, +) +from waveorder.calibrate.synthesis import fluorescence_otf +from waveorder.calibrate.tfzarr import TransferFunctionStore +from waveorder.config.schema import V4Config + +ZYX_SHAPE = (16, 32, 32) +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, +) +TRUTH = {"astigmatism_0": 0.12, "coma_x": 0.06} +SIGMA, LAM = 0.05, 1e-2 + + +def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA) -> 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": 0.2}, + "coma_x": {"nominal": 0.0, "tolerance": 0.2}, + }, + }, + }, + "object": {"prior": "tikhonov", "weight": weight}, + "noise": {"gaussian": {"sigma": sigma}}, + "solver": { + "calibration": { + "score": score, + "search": search or [{"gradient": {"max_iter": 150}}], + } + }, + } + ) + + +@pytest.fixture(scope="module") +def synthetic_data(): + otf_truth = fluorescence_otf(ZYX_SHAPE, zernike_coefficients=TRUTH, **OPTICS) + generator = torch.Generator().manual_seed(0) + prior_std = SIGMA / LAM**0.5 + f = torch.randn(ZYX_SHAPE, generator=generator) * prior_std + g = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf_truth)) + g = g + torch.randn(ZYX_SHAPE, generator=generator) * SIGMA + return g.numpy() + + +def _twin_error(parameters: np.ndarray, names: list[str]) -> float: + """Error vs truth, minimized over the global sign flip (exact twin).""" + truth = np.array([TRUTH[name] for name in names]) + direct = np.abs(parameters - truth).max() + flipped = np.abs(-parameters - truth).max() + return min(direct, flipped) + + +def test_build_parameter_specs_orders_by_noll(): + specs = build_parameter_specs(_config()) + assert [s.name for s in specs] == ["astigmatism_0", "coma_x"] # Noll 6, 8 + assert specs[0].tolerance == 0.2 + + +def test_evidence_gradient_recovers_truth(synthetic_data): + result = run_evidence_calibration(_config(), synthetic_data) + assert _twin_error(result.parameters[0], result.parameter_names) < 0.05, result.parameters + assert len(result.score_trace) == 150 + variances = np.diag(result.parameters_covariance[0]) + assert np.all(np.isfinite(variances)) and np.all(variances > 0) + assert np.sqrt(variances).max() < 0.1 # tight posterior at this SNR + + +def test_evidence_grid_then_gradient(synthetic_data): + config = _config(search=[{"grid": {"points": 5}}, {"gradient": {"max_iter": 60}}]) + result = run_evidence_calibration(config, synthetic_data) + assert _twin_error(result.parameters[0], result.parameter_names) < 0.05 + assert result.score_grid is not None + assert result.score_grid.shape == (25, 3) # 5 x 5 points, 2 coords + score + + +def test_random_then_gradient(synthetic_data): + config = _config(search=[{"random": {"samples": 8, "seed": 1}}, {"gradient": {"max_iter": 60}}]) + result = run_evidence_calibration(config, synthetic_data) + assert _twin_error(result.parameters[0], result.parameter_names) < 0.06 + + +def test_gcv_recovers_truth(synthetic_data): + result = run_evidence_calibration(_config(score="gcv"), synthetic_data) + assert _twin_error(result.parameters[0], result.parameter_names) < 0.06, result.parameters + + +def test_sigma_estimate_and_weight_auto(synthetic_data): + # The score is exactly symmetric under the global sign flip of the + # coefficients, so the zero nominal is a stationary point; a grid first + # stage (as the design prescribes for multi-modal searches) moves the + # start off it before sigma and lambda adapt. + config = _config( + sigma="estimate", + weight="auto", + search=[{"grid": {"points": 5}}, {"gradient": {"max_iter": 100}}], + ) + result = run_evidence_calibration(config, synthetic_data) + assert _twin_error(result.parameters[0], result.parameter_names) < 0.06 + # The jointly-estimated noise level should land near the truth. + assert 0.3 * SIGMA < result.noise_sigma < 3.0 * SIGMA, result.noise_sigma + assert np.isfinite(result.weight) and result.weight > 0 + + +def _store(values: dict[str, float], variances: dict[str, float] | None = None): + names = list(values) + covariance = np.zeros((1, len(names), len(names))) + if variances is not None: + covariance[0] = np.diag([variances[name] for name in names]) + return TransferFunctionStore( + parameters=np.array([[values[name] for name in names]]), + parameter_names=names, + tile_grid={}, + parameters_covariance=covariance, + ) + + +def test_init_store_sets_start_without_prior(synthetic_data): + config = _config(search=[{"gradient": {"max_iter": 1}}]) + init = _store({"astigmatism_0": 0.3, "coma_x": -0.2}) + result = run_evidence_calibration(config, synthetic_data, init_store=init) + # One Adam step (lr 0.01) cannot move far from the --init start. + np.testing.assert_allclose(result.parameters[0], [0.3, -0.2], atol=0.02) + + +def test_prior_covariance_overrides_tolerances(synthetic_data): + # A tight -t covariance pins the estimate near the prior center even + # though the configuration's own tolerance (0.2) would let it move. + prior = _store( + {"astigmatism_0": 0.3, "coma_x": -0.1}, + variances={"astigmatism_0": 1e-8, "coma_x": 1e-8}, + ) + result = run_evidence_calibration(_config(), synthetic_data, prior_store=prior) + np.testing.assert_allclose(result.parameters[0], [0.3, -0.1], atol=0.02) + + +def test_zero_covariance_moves_center_only(synthetic_data): + # An all-zero covariance (e.g. from simulate) carries no uncertainty: + # the config tolerances stay in force, only the start/center moves. + prior = _store({"astigmatism_0": 0.3, "coma_x": -0.1}) + result = run_evidence_calibration(_config(), synthetic_data, prior_store=prior) + assert _twin_error(result.parameters[0], result.parameter_names) < 0.06 + + +def test_sequential_calibration(synthetic_data): + """Design case 12: -t yesterday's tf_hat.zarr as today's prior.""" + first = run_evidence_calibration(_config(), synthetic_data) + store = TransferFunctionStore( + parameters=first.parameters, + parameter_names=first.parameter_names, + tile_grid={}, + parameters_covariance=first.parameters_covariance, + ) + + otf_truth = fluorescence_otf(ZYX_SHAPE, zernike_coefficients=TRUTH, **OPTICS) + generator = torch.Generator().manual_seed(7) + f = torch.randn(ZYX_SHAPE, generator=generator) * (SIGMA / LAM**0.5) + second_data = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf_truth)) + second_data = (second_data + torch.randn(ZYX_SHAPE, generator=generator) * SIGMA).numpy() + + second = run_evidence_calibration(_config(), second_data, prior_store=store) + assert _twin_error(second.parameters[0], second.parameter_names) < 0.05 + # Two volumes carry more information than one: the posterior tightens. + assert np.trace(second.parameters_covariance[0]) < np.trace(first.parameters_covariance[0]) * 1.5 + + +def test_gcv_score_matches_definition(): + otf = fluorescence_otf(ZYX_SHAPE, **OPTICS) + g = torch.randn(ZYX_SHAPE, generator=torch.Generator().manual_seed(2)) + g_fft = torch.fft.fftn(g, norm="ortho") + lam = 1e-2 + h2 = torch.abs(otf) ** 2 + shrink = lam / (h2 + lam) + n = g_fft.numel() + expected = ((torch.abs(g_fft) ** 2) * shrink**2).sum() / ((shrink.sum() / n) ** 2) + assert torch.allclose(gcv_score(g_fft, otf, lam), expected) + + +def test_singular_prior_covariance_rejected(synthetic_data): + singular = _store( + {"astigmatism_0": 0.3, "coma_x": -0.1}, + variances={"astigmatism_0": 1e-8, "coma_x": 0.0}, # zero-variance row + ) + # Force the non-zero branch: one finite variance, one zero -> singular. + singular.parameters_covariance[0, 0, 0] = 1e-8 + with pytest.raises(ValueError, match="singular"): + run_evidence_calibration(_config(), synthetic_data, prior_store=singular) + + +def test_grid_range_centers_on_prior_center(synthetic_data): + # A -t prior (zero covariance -> tolerances kept) moves the grid's + # center: with the config tolerance 0.2, the grid spans the prior + # center +/- 0.4, so the best grid point tracks the prior center. + prior = _store({"astigmatism_0": 0.5, "coma_x": 0.5}) + config = _config(search=[{"grid": {"points": 3}}]) + result = run_evidence_calibration(_config(search=[{"grid": {"points": 3}}]), synthetic_data) + moved = run_evidence_calibration(config, synthetic_data, prior_store=prior) + # Without the prior the 3-point grid centered at 0 cannot reach 0.5. + assert abs(result.parameters[0, 0]) <= 0.4 + 1e-9 + # With the prior the grid includes points around 0.5. + assert moved.parameters[0, 0] >= 0.1 - 1e-9 diff --git a/tests/calibrate/test_run_tfzarr.py b/tests/calibrate/test_run_tfzarr.py new file mode 100644 index 00000000..cb0aa46e --- /dev/null +++ b/tests/calibrate/test_run_tfzarr.py @@ -0,0 +1,186 @@ +"""4.0 configs drive the existing Zernike recovery; tf.zarr round-trips.""" + +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.run import build_recovery_settings, zernike_calibration_requested +from waveorder.calibrate.tfzarr import read_tf_zarr, write_tf_zarr +from waveorder.cli.main import app +from waveorder.config.schema import V4Config +from waveorder.models.shift_variant_fluorescent_3d import apply_shift_variant_forward +from waveorder.zernike import SpatialPolynomialPupil + + +def _base_config(**imaging_system_updates) -> dict: + cfg = { + "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": {"zernike_up_to_order": {"order": 4, "tolerance": "free"}}, + }, + }, + "solver": { + "calibration": { + "score": "sharpness", + "search": [{"gradient": {"max_iter": 30}}], + } + }, + } + cfg["imaging_system"].update(imaging_system_updates) + return cfg + + +def test_zernike_calibration_requested(): + cfg = V4Config(**_base_config()) + assert zernike_calibration_requested(cfg) + plain = _base_config() + plain["imaging_system"]["detection"].pop("pupil_aberrations") + plain.pop("solver") + assert not zernike_calibration_requested(V4Config(**plain)) + + +def test_case_5_shift_invariant_settings(): + """Design case 5 (AMBSF): order-4 free Zernikes, one full-FOV tile.""" + cfg = V4Config(**_base_config()) + settings, noll_indices = build_recovery_settings(cfg, (128, 96)) + assert noll_indices == [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + assert settings.tile_size_yx == {"y": 128, "x": 96} + assert settings.loss == "midband" + assert settings.n_iter == 30 + + +def test_case_17_tiled_settings(): + """Design case 17 (patchwise recovery): per-tile independent estimates.""" + cfg = V4Config(**_base_config(field_dependence={"tiled": {"size_yx": [32, 32], "overlap_yx": [8, 8]}})) + settings, _ = build_recovery_settings(cfg, (128, 128)) + assert settings.tile_size_yx == {"y": 32, "x": 32} + assert settings.tile_overlap_yx == {"y": 8.0, "x": 8.0} + + +def test_named_term_selects_single_mode(): + raw = _base_config() + raw["imaging_system"]["detection"]["pupil_aberrations"] = {"astigmatism_0": {"nominal": 0.0, "tolerance": "free"}} + _, noll_indices = build_recovery_settings(V4Config(**raw), (64, 64)) + assert noll_indices == [6] + + +def test_unsupported_paths_raise(): + evidence = _base_config() + evidence["solver"]["calibration"]["score"] = "evidence" + with pytest.raises(NotImplementedError, match="score"): + build_recovery_settings(V4Config(**evidence), (64, 64)) + + numeric = _base_config() + numeric["imaging_system"]["detection"]["pupil_aberrations"] = { + "zernike_up_to_order": {"order": 4, "tolerance": 0.05} + } + with pytest.raises(NotImplementedError, match="tolerance"): + build_recovery_settings(V4Config(**numeric), (64, 64)) + + +def test_tfzarr_round_trip(tmp_path): + parameters = np.arange(6, dtype=np.float64).reshape(2, 3) + tile_grid = {"size_yx": [32, 32], "overlap_yx": [0, 0], "centers_yx": [[-8.0, 0.0], [8.0, 0.0]]} + tf = (np.random.default_rng(0).random((2, 4, 8, 8)) + 0j).astype(np.complex64) + write_tf_zarr( + tmp_path / "tf.zarr", + parameters=parameters, + parameter_names=["defocus", "astigmatism_45", "astigmatism_0"], + tile_grid=tile_grid, + provenance={"kind": "refined", "sources": ["x.zarr/0/0/0"]}, + transfer_function=tf, + ) + store = read_tf_zarr(tmp_path / "tf.zarr") + np.testing.assert_array_equal(store.parameters, parameters) + np.testing.assert_array_equal(store.parameters_covariance, np.zeros((2, 3, 3))) + np.testing.assert_array_equal(store.transfer_function, tf) + assert store.parameter_names == ["defocus", "astigmatism_45", "astigmatism_0"] + assert store.tile_grid == tile_grid + assert store.provenance["kind"] == "refined" + + +def test_tfzarr_shape_validation(tmp_path): + with pytest.raises(ValueError, match="N_tiles"): + write_tf_zarr(tmp_path / "bad.zarr", np.zeros(3), ["a", "b", "c"], {}) + with pytest.raises(ValueError, match="parameter_names"): + write_tf_zarr(tmp_path / "bad.zarr", np.zeros((1, 3)), ["a"], {}) + + +@pytest.fixture +def aberrated_beads_zarr(tmp_path): + """A small bead grid imaged through a constant astigmatic pupil.""" + torch.manual_seed(0) + Z, Y, X = 24, 64, 64 + phantom = torch.zeros(Z, Y, X) + for y in (16, 32, 48): + for x in (16, 32, 48): + phantom[Z // 2, y, x] = 1.0 + + truth = 0.15 + pupil = SpatialPolynomialPupil({(6, 0, 0): truth}, field_extent_um=(6.4, 6.4)) + data = apply_shift_variant_forward( + phantom, + 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, + spatial_pupil=pupil, + n_tiles_yx=(2, 2), + ) + + 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", + data.numpy()[None, None].astype(np.float32), + transform=[TransformationMeta(type="scale", scale=[1, 1, 0.4, 0.2, 0.2])], + ) + dataset.close() + return input_path, truth + + +def test_calibrate_tf_cli_end_to_end(tmp_path, aberrated_beads_zarr): + input_path, truth = aberrated_beads_zarr + raw = _base_config() + raw["imaging_system"]["detection"]["pupil_aberrations"] = { + "astigmatism_0": {"nominal": 0.0, "tolerance": "free"}, + "astigmatism_45": {"nominal": 0.0, "tolerance": "free"}, + } + raw["solver"]["calibration"]["search"] = [{"gradient": {"max_iter": 120}}] + config_path = tmp_path / "calibrate.yml" + config_path.write_text(yaml.safe_dump(raw)) + + output_path = tmp_path / "tf_hat.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.parameter_names == ["astigmatism_45", "astigmatism_0"] + assert store.parameters.shape == (1, 2) + assert store.transfer_function is not None # shift-invariant stores carry H + assert store.provenance["kind"] == "refined" + assert store.tile_grid["size_yx"] == [64, 64] + + recovered = dict(zip(store.parameter_names, store.parameters[0])) + # The sharpness score sees |H| only, so the sign of a lone astigmatism + # pair is ambiguous; check the recovered magnitude. + magnitude = np.hypot(recovered["astigmatism_0"], recovered["astigmatism_45"]) + assert abs(magnitude - truth) < 0.08, recovered diff --git a/tests/calibrate/test_synthesis_scores.py b/tests/calibrate/test_synthesis_scores.py new file mode 100644 index 00000000..5f874876 --- /dev/null +++ b/tests/calibrate/test_synthesis_scores.py @@ -0,0 +1,143 @@ +"""PSF/OTF synthesis matches existing numerics; evidence recovers the truth.""" + +import torch + +from waveorder.calibrate.scores import calibration_penalty, whittle_evidence +from waveorder.calibrate.synthesis import fluorescence_otf, fluorescence_psf +from waveorder.models.shift_variant_fluorescent_3d import calculate_psf_at_field_position +from waveorder.zernike import SpatialPolynomialPupil + +ZYX_SHAPE = (24, 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, +) + + +def _reference_psf(coefficients: dict[tuple[int, int, int], float]) -> torch.Tensor: + """PSF from the existing shift-variant model at the field center.""" + pupil = SpatialPolynomialPupil(coefficients, field_extent_um=(10.0, 10.0)) + return calculate_psf_at_field_position( + ZYX_SHAPE, + OPTICS["yx_pixel_size"], + OPTICS["z_pixel_size"], + OPTICS["wavelength_emission"], + OPTICS["index_of_refraction_media"], + OPTICS["numerical_aperture_detection"], + pupil, + x_o_um=0.0, + y_o_um=0.0, + ) + + +def test_psf_matches_shift_variant_model_unaberrated(): + psf = fluorescence_psf(ZYX_SHAPE, **OPTICS) + reference = _reference_psf({}) + assert torch.allclose(psf, reference, atol=1e-6) + + +def test_psf_matches_shift_variant_model_aberrated(): + coefficients = {"spherical": 0.1, "coma_x": 0.05, "astigmatism_45": -0.03} + psf = fluorescence_psf(ZYX_SHAPE, zernike_coefficients=coefficients, **OPTICS) + reference = _reference_psf({(11, 0, 0): 0.1, (8, 0, 0): 0.05, (5, 0, 0): -0.03}) + assert torch.allclose(psf, reference, atol=1e-6) + + +def test_noll_and_named_keys_agree(): + psf_named = fluorescence_psf(ZYX_SHAPE, zernike_coefficients={"defocus": 0.2}, **OPTICS) + psf_noll = fluorescence_psf(ZYX_SHAPE, zernike_coefficients={4: 0.2}, **OPTICS) + assert torch.equal(psf_named, psf_noll) + + +def test_psf_normalized_and_centered(): + psf = fluorescence_psf(ZYX_SHAPE, **OPTICS) + assert abs(float(psf.sum()) - 1.0) < 1e-5 + # In-focus (z index 0 in FFT order) and origin-centered in y, x. + peak = torch.argmax(psf) + z, y, x = torch.unravel_index(peak, psf.shape) + assert int(z) == 0 and int(y) == 0 and int(x) == 0 + + +def test_otf_peak_normalized(): + otf = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"spherical": 0.1}, **OPTICS) + assert abs(float(torch.abs(otf).max()) - 1.0) < 1e-5 + + +def test_gradients_flow_through_coefficients(): + coefficient = torch.tensor(0.05, requires_grad=True) + psf = fluorescence_psf(ZYX_SHAPE, zernike_coefficients={"spherical": coefficient}, **OPTICS) + psf.sum().backward() # sum is constant 1, but graph must exist + assert coefficient.grad is not None + + +def _simulate(otf_truth: torch.Tensor, noise_sigma: float, regularization_strength: float, seed: int = 0): + """Draw g = H f + n with f from the Tikhonov prior S_k = sigma^2 / lambda.""" + generator = torch.Generator().manual_seed(seed) + shape = otf_truth.shape + prior_std = noise_sigma / regularization_strength**0.5 + f = torch.randn(shape, generator=generator) * prior_std + g = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf_truth)) + g = g + torch.randn(shape, generator=generator) * noise_sigma + return g + + +def test_evidence_prefers_true_astigmatism_magnitude(): + # Astigmatism changes |H|, which the Whittle evidence sees. Two known + # invariances temper what can be asserted: + # * pure defocus is close to an axial shift, which leaves |H| nearly + # unchanged, so it is only weakly identifiable from the evidence; + # * flipping the sign of the whole aberration vector reflects the + # intensity PSF through the origin, leaving |H| exactly unchanged + # (the classic blind-deconvolution PSF-flip ambiguity), so only the + # magnitude of a lone astigmatism term is identifiable. + sigma, lam = 0.1, 1e-2 + truth = 0.15 + otf_truth = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"astigmatism_0": truth}, **OPTICS) + g = _simulate(otf_truth, sigma, lam) + g_fft = torch.fft.fftn(g, norm="ortho") + + candidates = torch.linspace(-0.3, 0.6, 10).tolist() + [truth, -truth] + scores = {} + for c in candidates: + otf = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"astigmatism_0": c}, **OPTICS) + scores[c] = float(whittle_evidence(data_fft=g_fft, otf=otf, noise_sigma=sigma, regularization_strength=lam)) + best = min(scores, key=scores.get) + assert abs(abs(best) - truth) < 0.11, scores + + +def test_evidence_prefers_true_spherical(): + sigma, lam = 0.05, 1e-2 + truth = 0.2 + otf_truth = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"spherical": truth}, **OPTICS) + g = _simulate(otf_truth, sigma, lam, seed=1) + g_fft = torch.fft.fftn(g, norm="ortho") + + candidates = [0.0, 0.1, truth, 0.3, 0.4] + scores = [] + for c in candidates: + otf = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"spherical": c}, **OPTICS) + scores.append(float(whittle_evidence(data_fft=g_fft, otf=otf, noise_sigma=sigma, regularization_strength=lam))) + assert candidates[scores.index(min(scores))] == truth, dict(zip(candidates, scores)) + + +def test_evidence_is_differentiable(): + sigma, lam = 0.1, 1e-2 + otf_truth = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"defocus": 0.1}, **OPTICS) + g = _simulate(otf_truth, sigma, lam) + coefficient = torch.tensor(0.0, requires_grad=True) + otf = fluorescence_otf(ZYX_SHAPE, zernike_coefficients={"defocus": coefficient}, **OPTICS) + score = whittle_evidence(data=g, otf=otf, noise_sigma=sigma, regularization_strength=lam) + score.backward() + assert coefficient.grad is not None + assert torch.isfinite(coefficient.grad) + + +def test_calibration_penalty(): + values = {"na": 1.25, "spherical": 0.1, "defocus": 0.5} + nominals = {"na": 1.2, "spherical": 0.0, "defocus": 0.0} + tolerances = {"na": 0.05, "spherical": 0.1, "defocus": "free"} + penalty = float(calibration_penalty(values, nominals, tolerances)) + assert abs(penalty - (0.5 + 0.5)) < 1e-9 diff --git a/tests/calibrate/test_zernike_names.py b/tests/calibrate/test_zernike_names.py new file mode 100644 index 00000000..f2ae83b1 --- /dev/null +++ b/tests/calibrate/test_zernike_names.py @@ -0,0 +1,78 @@ +"""Named Zernike terms map to the standard Noll indices and orientations.""" + +import math + +import pytest +import torch + +from waveorder.calibrate.zernike_names import ( + NAMED_TERMS, + name_to_noll, + noll_indices_up_to_order, + noll_to_name, + validate_zernike_name, +) +from waveorder.zernike import zernike_noll + + +def test_named_terms(): + assert name_to_noll("defocus") == 4 + assert name_to_noll("astigmatism_45") == 5 + assert name_to_noll("astigmatism_0") == 6 + assert name_to_noll("coma_y") == 7 + assert name_to_noll("coma_x") == 8 + assert name_to_noll("trefoil_y") == 9 + assert name_to_noll("trefoil_x") == 10 + assert name_to_noll("spherical") == 11 + + +def test_zn_names(): + assert name_to_noll("z12") == 12 + assert name_to_noll("z4") == 4 + assert noll_to_name(12) == "z12" + assert noll_to_name(4) == "defocus" + + +def test_unknown_names_rejected(): + for bad in ("comatose", "z0", "zx", "Z4", ""): + with pytest.raises(ValueError): + validate_zernike_name(bad) + + +def test_orientation_conventions(): + """x-named terms are cosine terms (peak at theta = 0), y-named are sine.""" + rho = torch.tensor([1.0]) + theta_0 = torch.tensor([0.0]) + theta_90 = torch.tensor([math.pi / 2]) + + def at(name, theta): + return float(zernike_noll(name_to_noll(name), rho, theta)) + + assert at("coma_x", theta_0) > 1.0 + assert abs(at("coma_x", theta_90)) < 1e-5 + assert at("coma_y", theta_90) > 1.0 + assert abs(at("coma_y", theta_0)) < 1e-5 + assert at("astigmatism_0", theta_0) > 1.0 + assert at("astigmatism_0", theta_90) < -1.0 + assert abs(at("astigmatism_45", theta_0)) < 1e-5 + + +def test_orthonormal_normalization(): + """Coefficients are RMS waves: mean of Z^2 over the unit disk is 1.""" + n = 512 + y, x = torch.meshgrid(torch.linspace(-1, 1, n), torch.linspace(-1, 1, n), indexing="ij") + rho = torch.sqrt(x**2 + y**2) + theta = torch.atan2(y, x) + inside = rho <= 1.0 + for name in NAMED_TERMS: + z = zernike_noll(name_to_noll(name), rho, theta) + rms2 = (z[inside] ** 2).mean() + assert abs(float(rms2) - 1.0) < 0.02, name + + +def test_up_to_order(): + assert noll_indices_up_to_order(2) == [4, 5, 6] + assert noll_indices_up_to_order(3) == [4, 5, 6, 7, 8, 9, 10] + assert noll_indices_up_to_order(4) == [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + with pytest.raises(ValueError): + noll_indices_up_to_order(1) diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/config/repo_configs.py b/tests/config/repo_configs.py new file mode 100644 index 00000000..1ed8920c --- /dev/null +++ b/tests/config/repo_configs.py @@ -0,0 +1,62 @@ +"""Every reconstruction configuration YAML committed to the repository. + +The design document (section 2) requires the regression tests to cover +every 3.x configuration in the repository and in the published examples. +This module is the single source of that list; +``test_migrate.test_every_repo_config_is_listed`` guards it against new +configuration directories appearing without coverage. +""" + +from pathlib import Path + +import yaml + +REPO_ROOT = Path(__file__).parent.parent.parent + +#: Directories holding reconstruction configuration files. +CONFIG_DIRS = ( + REPO_ROOT / "docs/examples/cli/configs", + REPO_ROOT / "docs/examples/optimization", + REPO_ROOT / "benchmarks/configs", +) + +ALL_CONFIG_PATHS = sorted(path for config_dir in CONFIG_DIRS for path in config_dir.glob("*.yml")) + +#: Birefringence stays 3.x this release (design document section 2). +BIREFRINGENCE_CONFIG_PATHS = [p for p in ALL_CONFIG_PATHS if "birefringence" in p.name] + + +def _is_tile_stitch_config(path: Path) -> bool: + """Whether a configuration is a ``wo tile-stitch`` configuration. + + Tile-stitch configurations wrap a reconstruction block under a top-level + ``recon`` key alongside ``tile`` / ``blend`` stitching parameters, rather + than being a plain ``wo rec`` configuration. Detection is by structure, not + filename, so new tile-stitch cases are classified correctly regardless of + how they are named. + + Parameters + ---------- + path : Path + Path to a configuration YAML file. + + Returns + ------- + bool + ``True`` if the configuration has the tile-stitch structure. + """ + try: + doc = yaml.safe_load(path.read_text()) + except (OSError, yaml.YAMLError): + return False + return isinstance(doc, dict) and "recon" in doc and ("tile" in doc or "blend" in doc) + + +#: Tile-stitch configurations wrap a reconstruction block in stitching +#: parameters; they are not plain ``wo rec`` configurations. +TILE_STITCH_CONFIG_PATHS = [p for p in ALL_CONFIG_PATHS if _is_tile_stitch_config(p)] + +#: Every 3.x phase / fluorescence configuration in the repository. +MIGRATABLE_CONFIG_PATHS = [ + p for p in ALL_CONFIG_PATHS if p not in BIREFRINGENCE_CONFIG_PATHS and p not in TILE_STITCH_CONFIG_PATHS +] diff --git a/tests/config/test_fresh_imports.py b/tests/config/test_fresh_imports.py new file mode 100644 index 00000000..0e76ffcb --- /dev/null +++ b/tests/config/test_fresh_imports.py @@ -0,0 +1,32 @@ +"""Fresh-process imports, in the order the CLI triggers them. + +pytest imports modules in whatever order the collected tests dictate, +which can hide circular imports that a cold ``wo`` process hits. These +tests spawn a clean interpreter per import path. +""" + +import subprocess +import sys + +import pytest + +IMPORT_STATEMENTS = [ + # The first import `wo rec` performs when it loads a config. + "from waveorder.config.loader import load_reconstruction_settings", + # The reverse entry: calibration numerics before any config import. + "from waveorder.calibrate import run_evidence_calibration, run_zernike_calibration", + # The schema alone (imports the Zernike names from waveorder.calibrate). + "from waveorder.config.schema import V4Config", + # The CLI app with every registered verb. + "import waveorder.cli.main", +] + + +@pytest.mark.parametrize("statement", IMPORT_STATEMENTS) +def test_fresh_process_import(statement): + result = subprocess.run( + [sys.executable, "-c", statement], + capture_output=True, + text=True, + ) + assert result.returncode == 0, f"{statement!r} failed:\n{result.stderr}" diff --git a/tests/config/test_migrate.py b/tests/config/test_migrate.py new file mode 100644 index 00000000..55a56a8c --- /dev/null +++ b/tests/config/test_migrate.py @@ -0,0 +1,272 @@ +"""3.x -> 4.0 migration: deterministic, lossless, and round-trippable.""" + +from pathlib import Path + +import pytest +import yaml + +from tests.config.repo_configs import ( + ALL_CONFIG_PATHS, + BIREFRINGENCE_CONFIG_PATHS, + CONFIG_DIRS, + MIGRATABLE_CONFIG_PATHS, + REPO_ROOT, + TILE_STITCH_CONFIG_PATHS, +) +from waveorder.cli.settings import ReconstructionSettings +from waveorder.config import ( + MigrationError, + NotRunnableInLegacyError, + detect_version, + migrate_3x_dict, + v4_to_legacy, + v4_to_yaml_text, +) +from waveorder.config.schema import V4Config + + +def _load(path: Path) -> dict: + with open(path) as f: + return yaml.safe_load(f) + + +def test_every_repo_config_is_listed(): + """Guard: every configuration YAML in docs/ and benchmarks/ is covered. + + A new configuration directory (or a renamed one) must be added to + ``tests.config.repo_configs`` so the migration and end-to-end + equivalence tests pick it up. Experiment YAMLs describe benchmark + cases, not reconstructions, and are exempt. + """ + for config_dir in CONFIG_DIRS: + assert list(config_dir.glob("*.yml")), f"{config_dir} is empty or missing" + found = { + path + for root in (REPO_ROOT / "docs", REPO_ROOT / "benchmarks") + for path in root.rglob("*.yml") + if "experiments" not in path.parts + } + assert found == set(ALL_CONFIG_PATHS) + covered = set(MIGRATABLE_CONFIG_PATHS) | set(BIREFRINGENCE_CONFIG_PATHS) | set(TILE_STITCH_CONFIG_PATHS) + assert covered == set(ALL_CONFIG_PATHS) + + +def test_version_detection(): + assert detect_version({}) == 3 + assert detect_version({"input_channel_names": ["BF"]}) == 3 + assert detect_version({"version": 4}) == 4 + + +@pytest.mark.parametrize("config_path", MIGRATABLE_CONFIG_PATHS, ids=lambda p: p.name) +def test_migration_round_trip(config_path): + """migrate(3.x) -> legacy reproduces the directly-parsed settings field for field.""" + raw = _load(config_path) + direct = ReconstructionSettings(**raw) + migrated = migrate_3x_dict(raw) + round_tripped = v4_to_legacy(migrated) + assert round_tripped.model_dump() == direct.model_dump() + + +@pytest.mark.parametrize("config_path", MIGRATABLE_CONFIG_PATHS, ids=lambda p: p.name) +def test_migration_is_deterministic(config_path): + raw = _load(config_path) + assert migrate_3x_dict(raw).model_dump() == migrate_3x_dict(raw).model_dump() + + +@pytest.mark.parametrize("config_path", MIGRATABLE_CONFIG_PATHS, ids=lambda p: p.name) +def test_migrated_yaml_reparses(config_path): + """The written 4.0 YAML parses back to the same configuration.""" + migrated = migrate_3x_dict(_load(config_path)) + reparsed = V4Config(**yaml.safe_load(v4_to_yaml_text(migrated))) + assert reparsed.model_dump() == migrated.model_dump() + + +@pytest.mark.parametrize("config_path", BIREFRINGENCE_CONFIG_PATHS, ids=lambda p: p.name) +def test_birefringence_not_migratable(config_path): + with pytest.raises(MigrationError): + migrate_3x_dict(_load(config_path)) + + +@pytest.mark.parametrize("config_path", TILE_STITCH_CONFIG_PATHS, ids=lambda p: p.name) +def test_tile_stitch_not_migratable(config_path): + """Tile-stitch configs (tile/blend/recon) raise a clear MigrationError. + + ``wo tile-stitch`` reads the 3.x TileStitchSettings; there is no 4.0 + tile-stitch execution path, so the whole-config migration must refuse + with MigrationError (which the benchmark runner catches to run the case + unmigrated) instead of a raw pydantic ValidationError. + """ + with pytest.raises(MigrationError): + migrate_3x_dict(_load(config_path)) + + +def test_block_regrouping(): + """Spot-check the field map of design document section 2.""" + raw = { + "input_channel_names": ["BF"], + "time_indices": [0, 5], + "reconstruction_dimension": 2, + "phase": { + "transfer_function": { + "yx_pixel_size": 0.325, + "z_pixel_size": 2.0, + "z_padding": 3, + "wavelength_illumination": 0.45, + "index_of_refraction_media": 1.0, + "numerical_aperture_illumination": 0.4, + "numerical_aperture_detection": 0.55, + "invert_phase_contrast": True, + }, + "apply_inverse": { + "reconstruction_algorithm": "Tikhonov", + "regularization_strength": 0.001, + }, + }, + } + cfg = migrate_3x_dict(raw) + assert cfg.version == 4 + assert cfg.data.input_channel_names == ["BF"] + assert cfg.data.time_indices == [0, 5] + assert cfg.data.pixel_size.yx.y == 0.325 + assert cfg.data.pixel_size.z == 2.0 + assert cfg.object.dimension == 2 + assert cfg.solver.padding.z == 3 + assert cfg.imaging_system.contrast == "phase" + assert cfg.imaging_system.index_of_refraction_media == 1.0 + assert cfg.imaging_system.detection.numerical_aperture == 0.55 + assert cfg.imaging_system.illumination.wavelength == 0.45 + assert cfg.imaging_system.illumination.numerical_aperture == 0.4 + assert cfg.imaging_system.illumination.invert_phase_contrast is True + assert cfg.object.prior == "tikhonov" + assert cfg.object.weight == 0.001 + assert cfg.solver.reconstruction == "closed_form" + assert cfg.noise.kind == "gaussian" + assert cfg.noise.gaussian.sigma == 1.0 + + +def test_tv_maps_to_admm(): + raw = { + "input_channel_names": ["BF"], + "phase": { + "transfer_function": {}, + "apply_inverse": { + "reconstruction_algorithm": "TV", + "regularization_strength": 0.01, + "TV_rho_strength": 0.005, + "TV_iterations": 7, + }, + }, + } + cfg = migrate_3x_dict(raw) + assert cfg.object.prior == "tv" + assert cfg.object.weight == 0.01 + assert cfg.solver.reconstruction.admm.rho == 0.005 + assert cfg.solver.reconstruction.admm.iterations == 7 + round_tripped = v4_to_legacy(cfg) + assert round_tripped.model_dump() == ReconstructionSettings(**raw).model_dump() + + +def test_rl_maps_to_richardson_lucy_with_poisson(): + raw = { + "input_channel_names": ["GFP"], + "fluorescence": { + "transfer_function": {}, + "apply_inverse": { + "reconstruction_algorithm": "RLGC", + "rl": {"iterations": 13, "background": 2.0}, + }, + }, + } + cfg = migrate_3x_dict(raw) + assert cfg.object.prior == "none" + assert cfg.noise.kind == "poisson" + rl = cfg.solver.reconstruction.richardson_lucy + assert rl.iterations == 13 + assert rl.background == 2.0 + assert rl.gradient_consensus is True + round_tripped = v4_to_legacy(cfg) + assert round_tripped.model_dump() == ReconstructionSettings(**raw).model_dump() + + +def test_confocal_pinhole_maps_to_detection(): + raw = { + "input_channel_names": ["GFP"], + "fluorescence": { + "transfer_function": {"confocal_pinhole_diameter": 0.5, "wavelength_emission": 0.52}, + "apply_inverse": {}, + }, + } + cfg = migrate_3x_dict(raw) + assert cfg.imaging_system.detection.wavelength == 0.52 + assert cfg.imaging_system.detection.pinhole.diameter_um == 0.5 + round_tripped = v4_to_legacy(cfg) + assert round_tripped.model_dump() == ReconstructionSettings(**raw).model_dump() + + +def test_optimizable_floats_map_to_free_parameters(): + raw = { + "input_channel_names": ["BF"], + "reconstruction_dimension": 2, + "phase": { + "transfer_function": { + "z_focus_offset": {"init": 1.5, "lr": 0.05}, + "tilt_angle_zenith": {"init": 0.0, "lr": 0.005}, + }, + "apply_inverse": {"regularization_strength": 0.01}, + }, + "optimization": {"max_iterations": 50, "method": "adam"}, + } + cfg = migrate_3x_dict(raw) + z_offset = cfg.imaging_system.detection.z_focus_offset + assert z_offset.nominal == 1.5 + assert z_offset.tolerance == "free" + calibration = cfg.solver.calibration + assert calibration.score == "sharpness" + assert calibration.search[0].gradient.max_iter == 50 + assert calibration.advanced.optimizer == "adam" + assert calibration.advanced.learning_rates == { + "detection.z_focus_offset": 0.05, + "illumination.tilt.zenith": 0.005, + } + round_tripped = v4_to_legacy(cfg) + assert round_tripped.model_dump() == ReconstructionSettings(**raw).model_dump() + + +def test_auto_regularization_maps_to_weight_auto(): + raw = { + "input_channel_names": ["GFP"], + "fluorescence": { + "transfer_function": {}, + "apply_inverse": {"auto_regularization": {"rule": "l_curve", "num_samples": 11}}, + }, + } + cfg = migrate_3x_dict(raw) + assert cfg.object.weight.auto.rule == "l_curve" + assert cfg.object.weight.auto.num_samples == 11 + round_tripped = v4_to_legacy(cfg) + assert round_tripped.model_dump() == ReconstructionSettings(**raw).model_dump() + + +def test_detect_version_rejects_empty_file(): + with pytest.raises(ValueError, match="mapping"): + detect_version(None) + + +def test_free_parameters_without_calibration_fail_loudly(): + """A 4.0 config with tolerances but no calibration block must not silently + run the 3.x sharpness optimizer (the design default score is evidence).""" + raw = { + "version": 4, + "data": {"input_channel_names": ["BF"], "pixel_size": {"yx": 0.1, "z": 0.25}}, + "imaging_system": { + "contrast": "phase", + "detection": {"numerical_aperture": {"nominal": 1.2, "tolerance": "free"}}, + "illumination": {"wavelength": 0.532, "numerical_aperture": 0.9}, + }, + } + with pytest.raises(NotRunnableInLegacyError, match="solver.calibration"): + v4_to_legacy(V4Config(**raw)) + # An explicit sharpness block selects the 3.x optimizer path. + raw["solver"] = {"calibration": {"score": "sharpness", "search": [{"gradient": {"max_iter": 10}}]}} + legacy = v4_to_legacy(V4Config(**raw)) + assert legacy.phase.transfer_function.numerical_aperture_detection.is_optimizable diff --git a/tests/config/test_migrate_cli.py b/tests/config/test_migrate_cli.py new file mode 100644 index 00000000..0dd6e6fb --- /dev/null +++ b/tests/config/test_migrate_cli.py @@ -0,0 +1,152 @@ +"""migrate-config CLI, and 3.x vs migrated-4.0 end-to-end equivalence.""" + +import warnings +from pathlib import Path + +import numpy as np +import pytest +import yaml +from iohub.ngff import open_ome_zarr +from iohub.ngff.models import TransformationMeta +from typer.testing import CliRunner + +from tests.config.repo_configs import MIGRATABLE_CONFIG_PATHS +from waveorder.cli.main import app +from waveorder.config import loader +from waveorder.config.schema import V4Config + +FLUOR_3X_CONFIG = { + "input_channel_names": ["GFP"], + "reconstruction_dimension": 3, + "fluorescence": { + "transfer_function": { + "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, + }, + "apply_inverse": { + "reconstruction_algorithm": "Tikhonov", + "regularization_strength": 1e-3, + }, + }, +} + + +def _write_input_zarr(input_path: Path, channel_names: list[str], tzyx_shape: tuple[int, int, int, int]) -> Path: + """Write a tiny random single-position dataset with the given channels. + + Parameters + ---------- + input_path : Path + Destination zarr path. + channel_names : list[str] + Channel names; the channel axis has ``len(channel_names)`` planes. + tzyx_shape : tuple[int, int, int, int] + ``(T, Z, Y, X)`` shape of each channel. + + Returns + ------- + Path + ``input_path``, for chaining. + """ + t, z, y, x = tzyx_shape + dataset = open_ome_zarr(input_path, layout="hcs", mode="w", channel_names=channel_names) + position = dataset.create_position("0", "0", "0") + rng = np.random.default_rng(0) + position.create_image( + "0", + rng.random((t, len(channel_names), z, y, x)).astype(np.float32), + transform=[TransformationMeta(type="scale", scale=[1, 1, 0.4, 0.2, 0.2])], + ) + dataset.close() + return input_path + + +@pytest.fixture +def input_zarr(tmp_path): + """A tiny single-position fluorescence dataset.""" + return _write_input_zarr(tmp_path / "input.zarr", ["GFP"], (2, 8, 16, 16)) + + +def test_migrate_config_cli_writes_v4(tmp_path): + config_path = tmp_path / "config.yml" + config_path.write_text(yaml.safe_dump(FLUOR_3X_CONFIG)) + result = CliRunner().invoke(app, ["migrate-config", str(config_path)]) + assert result.exit_code == 0, result.output + + output_path = tmp_path / "config_v4.yml" + assert output_path.exists() + migrated = yaml.safe_load(output_path.read_text()) + assert migrated["version"] == 4 + V4Config(**migrated) # validates + + +def test_migrate_config_cli_rejects_v4_input(tmp_path): + config_path = tmp_path / "config.yml" + config_path.write_text(yaml.safe_dump({"version": 4})) + result = CliRunner().invoke(app, ["migrate-config", str(config_path)]) + assert result.exit_code == 0 + assert "already a 4.0 configuration" in result.output + + +def _reconstruct(input_zarr: Path, config_path: Path, output_path: Path) -> np.ndarray: + result = CliRunner().invoke( + app, + ["reconstruct", "-i", str(input_zarr / "0" / "0" / "0"), "-c", str(config_path), "-o", str(output_path)], + ) + assert result.exit_code == 0, result.output + with open_ome_zarr(output_path / "0" / "0" / "0", mode="r") as ds: + return ds.data[:] + + +def test_3x_and_migrated_v4_reconstructions_match(tmp_path, input_zarr): + """A migrated configuration produces the same output array, bit for bit.""" + config_3x = tmp_path / "config.yml" + config_3x.write_text(yaml.safe_dump(FLUOR_3X_CONFIG)) + + result = CliRunner().invoke(app, ["migrate-config", str(config_3x), "-o", str(tmp_path / "config_v4.yml")]) + assert result.exit_code == 0, result.output + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + recon_3x = _reconstruct(input_zarr, config_3x, tmp_path / "out_3x.zarr") + recon_v4 = _reconstruct(input_zarr, tmp_path / "config_v4.yml", tmp_path / "out_v4.zarr") + + np.testing.assert_array_equal(recon_3x, recon_v4) + + +@pytest.mark.parametrize("config_path", MIGRATABLE_CONFIG_PATHS, ids=lambda p: p.name) +def test_every_repo_config_migrates_and_matches(tmp_path, monkeypatch, config_path): + """Design section 2: migrate every 3.x configuration in the repository + and the published examples, run it, and assert the output matches the + 3.x output array by array.""" + monkeypatch.chdir(tmp_path) # optimizer log_dir entries are relative paths + raw = yaml.safe_load(config_path.read_text()) + input_path = _write_input_zarr(tmp_path / "input.zarr", list(raw["input_channel_names"]), (1, 8, 16, 16)) + + config_3x = tmp_path / "config.yml" + config_3x.write_text(config_path.read_text()) + result = CliRunner().invoke(app, ["migrate-config", str(config_3x), "-o", str(tmp_path / "config_v4.yml")]) + assert result.exit_code == 0, result.output + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + recon_3x = _reconstruct(input_path, config_3x, tmp_path / "out_3x.zarr") + recon_v4 = _reconstruct(input_path, tmp_path / "config_v4.yml", tmp_path / "out_v4.zarr") + + np.testing.assert_array_equal(recon_3x, recon_v4) + + +def test_3x_config_warns_and_prints_v4(tmp_path, input_zarr): + config_3x = tmp_path / "deprecation_check.yml" + config_3x.write_text(yaml.safe_dump(FLUOR_3X_CONFIG)) + + loader._already_warned_paths.clear() + with pytest.warns(DeprecationWarning, match="migrate-config"): + loader.load_reconstruction_settings(config_3x) + # Second load of the same path stays quiet. + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + loader.load_reconstruction_settings(config_3x) diff --git a/tests/config/test_simulate_v4.py b/tests/config/test_simulate_v4.py new file mode 100644 index 00000000..5ff4b8a6 --- /dev/null +++ b/tests/config/test_simulate_v4.py @@ -0,0 +1,114 @@ +"""The 4.0 simulate verb: standard outputs, determinism, calibration closure.""" + +import numpy as np +import pytest +import yaml +from iohub.ngff import open_ome_zarr +from typer.testing import CliRunner + +from waveorder.calibrate.evidence_search import run_evidence_calibration +from waveorder.calibrate.tfzarr import read_tf_zarr +from waveorder.cli.main import app +from waveorder.config.schema import V4Config +from waveorder.config.simulate import run_simulation + +SIM_RAW = { + "version": 4, + "simulate": {"seed": 0, "shape": [16, 32, 32]}, + "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": { + "spherical": 0.08, + "astigmatism_0": {"nominal": 0.0, "tolerance": 0.1}, + }, + }, + }, + "object": {"source": {"gaussian_spectrum": {"exponent": 2.0}}}, + "noise": {"gaussian": {"snr_db": 20}}, +} + + +def _config(**updates) -> V4Config: + raw = yaml.safe_load(yaml.safe_dump(SIM_RAW)) + raw.update(updates) + return V4Config(**raw) + + +def test_simulation_writes_standard_outputs(tmp_path): + truth = run_simulation(_config(), tmp_path / "sim") + + with open_ome_zarr(tmp_path / "sim" / "data.zarr" / "0" / "0" / "0") as data: + assert data.data.shape == (1, 1, 16, 32, 32) + assert data.channel_names == ["GFP"] + assert data.scale[-3:] == [0.4, 0.2, 0.2] + with open_ome_zarr(tmp_path / "sim" / "object.zarr" / "0" / "0" / "0") as obj: + assert obj.data.shape == (1, 1, 16, 32, 32) + + store = read_tf_zarr(tmp_path / "sim" / "tf.zarr") + # P = parameters with a tolerance: only astigmatism_0 is free. + assert store.parameter_names == ["astigmatism_0"] + assert store.parameters[0, 0] == truth["astigmatism_0"] + assert abs(truth["astigmatism_0"]) < 0.4 # drawn from N(0, 0.1) + assert truth["spherical"] == 0.08 # fixed term, in provenance not parameters + np.testing.assert_array_equal(store.parameters_covariance, np.zeros((1, 1, 1))) + assert store.transfer_function.shape == (16, 32, 32) + assert store.provenance["kind"] == "simulated" + assert store.provenance["true_coefficients_waves"]["spherical"] == 0.08 + + +def test_simulation_is_deterministic(tmp_path): + truth_a = run_simulation(_config(), tmp_path / "a") + truth_b = run_simulation(_config(), tmp_path / "b") + assert truth_a == truth_b + with open_ome_zarr(tmp_path / "a" / "data.zarr" / "0" / "0" / "0") as a: + with open_ome_zarr(tmp_path / "b" / "data.zarr" / "0" / "0" / "0") as b: + np.testing.assert_array_equal(a.data[:], b.data[:]) + truth_c = run_simulation(_config(simulate={"seed": 1, "shape": [16, 32, 32]}), tmp_path / "c") + assert truth_c["astigmatism_0"] != truth_a["astigmatism_0"] + + +def test_poisson_noise(tmp_path): + cfg = _config(noise={"poisson": {"photons": 50}}) + run_simulation(cfg, tmp_path / "sim") + with open_ome_zarr(tmp_path / "sim" / "data.zarr" / "0" / "0" / "0") as data: + counts = data.data[:] + assert counts.min() >= 0 + assert counts == pytest.approx(np.round(counts)) # integer counts + assert 25 < counts.mean() < 100 + + +def test_simulate_cli(tmp_path): + config_path = tmp_path / "truth.yml" + config_path.write_text(yaml.safe_dump(SIM_RAW)) + result = CliRunner().invoke(app, ["simulate", "-c", str(config_path), "-o", str(tmp_path / "sim")]) + assert result.exit_code == 0, result.output + assert (tmp_path / "sim" / "data.zarr").exists() + assert (tmp_path / "sim" / "object.zarr").exists() + assert (tmp_path / "sim" / "tf.zarr").exists() + + +def test_calibration_closure_on_simulated_data(tmp_path): + """Case 28-style closure: calibrate-tf on simulate's output recovers the truth.""" + truth = run_simulation(_config(), tmp_path / "sim") + with open_ome_zarr(tmp_path / "sim" / "data.zarr" / "0" / "0" / "0") as data: + zyx = np.asarray(data.data[0, 0]) + + recon_raw = yaml.safe_load(yaml.safe_dump(SIM_RAW)) + recon_raw.pop("simulate") + recon_raw["object"] = {"prior": "tikhonov", "weight": 1e-3} + recon_raw["noise"] = {"gaussian": {"sigma": "estimate"}} + recon_raw["solver"] = { + "calibration": {"score": "evidence", "search": [{"grid": {"points": 7}}, {"gradient": {"max_iter": 80}}]} + } + result = run_evidence_calibration(V4Config(**recon_raw), zyx) + recovered = result.parameters[0, 0] + # |H| only: the estimate is defined up to the global sign flip of the + # free terms... but the fixed spherical term breaks the exact twin here, + # so compare against the better of the two signings with a loose gate. + error = min(abs(recovered - truth["astigmatism_0"]), abs(-recovered - truth["astigmatism_0"])) + assert error < 0.05, (recovered, truth) diff --git a/tests/config/test_v4_schema.py b/tests/config/test_v4_schema.py new file mode 100644 index 00000000..db2f7e13 --- /dev/null +++ b/tests/config/test_v4_schema.py @@ -0,0 +1,230 @@ +"""The 4.0 schema parses the design document's examples and enforces its rules.""" + +import pytest +import yaml + +from waveorder.config.schema import ( + FreeParameter, + SearchGradient, + SearchGrid, + V4Config, + free_parameters, +) + +# Design document example 4.1: blind, tiled, evidence-scored phase reconstruction. +EXAMPLE_4_1 = """ +version: 4 + +data: + input_channel_names: [BF] + time_indices: all + pixel_size: {yx: 0.1, z: 0.25} + +imaging_system: + contrast: phase + field_dependence: + tiled: {size_yx: [256, 256], overlap_yx: [32, 32], tile_blend: linear} + index_of_refraction_media: 1.3 + detection: + numerical_aperture: {nominal: 1.2, tolerance: 0.05, bounds: [1.0, 1.3]} + pupil_aberrations: + zernike_up_to_order: {order: 4, tolerance: 0.02} + spherical: {nominal: 0.05, tolerance: 0.02} + illumination: + wavelength: 0.532 + numerical_aperture: 0.9 + tilt: {zenith: {nominal: 0.0, tolerance: 0.02}, azimuth: 0.0} + invert_phase_contrast: false + +object: + dimension: 3 + prior: tikhonov + weight: auto + constraints: [] + +noise: + gaussian: {sigma: estimate} + +solver: + padding: {z: 5, yx: 0} + reconstruction: closed_form + calibration: + score: evidence + search: + - grid: {points: 7} + - gradient: {max_iter: 200} + final_estimate: best_H + +outputs: [reconstruction, transfer_function, psf, parameters, parameters_covariance, + score_trace, score_grid, weight_trace, residual, tile_map] +""" + +# Design document example 4.2: the same system, non-blind, as 3.x ran it. +EXAMPLE_4_2 = """ +version: 4 +data: + input_channel_names: [BF] + pixel_size: {yx: 0.1, z: 0.25} +imaging_system: + contrast: phase + index_of_refraction_media: 1.3 + detection: {numerical_aperture: 1.2} + illumination: {wavelength: 0.532, numerical_aperture: 0.9} +object: {prior: tikhonov, weight: 1e-3} +""" + +# Design document example 4.3: bead calibration (fluorescence). +EXAMPLE_4_3 = """ +version: 4 +data: + input_channel_names: [GFP] + pixel_size: {yx: 0.1, z: 0.25} +imaging_system: + contrast: fluorescence + field_dependence: shift_invariant + index_of_refraction_media: 1.3 + detection: + wavelength: 0.520 + numerical_aperture: {nominal: 1.2, tolerance: free, bounds: [1.0, 1.3]} + pinhole: none + pupil_aberrations: {zernike_up_to_order: {order: 4, tolerance: free}} +object: + prior: {beads: {diameter: 0.2}} +noise: {poisson_gaussian: {gain: 2.1, offset: 100, sigma: estimate}} +solver: + calibration: {score: evidence, search: [{gradient: {max_iter: 100}}]} +outputs: [parameters, parameters_covariance, transfer_function, psf] +""" + + +def _parse(text: str) -> V4Config: + return V4Config(**yaml.safe_load(text)) + + +def test_example_4_1_parses(): + cfg = _parse(EXAMPLE_4_1) + assert cfg.imaging_system.field_dependence.tiled.size_yx == (256, 256) + na = cfg.imaging_system.detection.numerical_aperture + assert isinstance(na, FreeParameter) + assert na.tolerance == 0.05 + assert na.bounds == (1.0, 1.3) + aberrations = cfg.imaging_system.detection.pupil_aberrations + assert aberrations.zernike_up_to_order.order == 4 + assert aberrations.terms["spherical"].nominal == 0.05 + assert isinstance(cfg.solver.calibration.search[0], SearchGrid) + assert isinstance(cfg.solver.calibration.search[1], SearchGradient) + assert cfg.noise.gaussian.sigma == "estimate" + assert cfg.object.weight == "auto" + + +def test_example_4_2_parses_and_is_fixed(): + cfg = _parse(EXAMPLE_4_2) + assert free_parameters(cfg.imaging_system) == {} + assert cfg.solver.reconstruction == "closed_form" + assert cfg.noise.kind == "gaussian" + assert cfg.noise.gaussian.sigma == 1.0 # default reproduces 3.x + assert cfg.solver.padding.z == 0 # 3.x z_padding default + assert cfg.solver.calibration is None + + +def test_example_4_3_parses(): + cfg = _parse(EXAMPLE_4_3) + assert cfg.object.prior.beads.diameter == 0.2 + assert cfg.noise.kind == "poisson_gaussian" + params = free_parameters(cfg.imaging_system) + assert "detection.numerical_aperture" in params + assert "detection.pupil_aberrations.zernike_up_to_order" in params + + +def test_free_parameters_collects_named_terms(): + cfg = _parse(EXAMPLE_4_1) + params = free_parameters(cfg.imaging_system) + assert params["detection.pupil_aberrations.spherical"].tolerance == 0.02 + assert params["illumination.tilt.zenith"].tolerance == 0.02 + assert "illumination.tilt.azimuth" not in params + + +def test_nominal_outside_bounds_rejected(): + with pytest.raises(ValueError, match="outside bounds"): + FreeParameter(nominal=2.0, tolerance=0.1, bounds=(0.0, 1.0)) + + +def test_na_above_index_rejected(): + bad = yaml.safe_load(EXAMPLE_4_2) + bad["imaging_system"]["detection"]["numerical_aperture"] = 1.4 + with pytest.raises(ValueError, match="index_of_refraction_media"): + V4Config(**bad) + + +def test_unknown_zernike_name_rejected(): + bad = yaml.safe_load(EXAMPLE_4_3) + bad["imaging_system"]["detection"]["pupil_aberrations"] = {"comatose": 0.1} + with pytest.raises(ValueError, match="Zernike"): + V4Config(**bad) + + +def test_grid_stage_with_unbounded_free_parameter_rejected(): + bad = yaml.safe_load(EXAMPLE_4_3) + bad["imaging_system"]["detection"]["numerical_aperture"] = {"nominal": 1.2, "tolerance": "free"} + bad["imaging_system"]["detection"].pop("pupil_aberrations") + bad["solver"]["calibration"]["search"] = [{"grid": {"points": 5}}] + with pytest.raises(ValueError, match="bounds"): + V4Config(**bad) + + +def test_rl_requires_poisson_and_no_prior(): + raw = yaml.safe_load(EXAMPLE_4_2) + raw["solver"] = {"reconstruction": {"richardson_lucy": {"iterations": 10}}} + with pytest.raises(ValueError, match="richardson_lucy"): + V4Config(**raw) + raw["object"] = {"prior": "none"} + raw["noise"] = {"poisson": {}} + cfg = V4Config(**raw) + assert cfg.reconstruction_kind == "richardson_lucy" + + +def test_prior_none_with_closed_form_rejected(): + raw = yaml.safe_load(EXAMPLE_4_2) + raw["object"] = {"prior": "none"} + with pytest.raises(ValueError, match="unregularized"): + V4Config(**raw) + + +def test_weight_auto_requires_evidence_family_score(): + raw = yaml.safe_load(EXAMPLE_4_3) + raw["object"] = {"prior": "tikhonov", "weight": "auto"} + raw["solver"]["calibration"]["score"] = "sharpness" + with pytest.raises(ValueError, match="weight: auto"): + V4Config(**raw) + + +def test_extra_fields_rejected(): + raw = yaml.safe_load(EXAMPLE_4_2) + raw["imaging_system"]["detection"]["numerical_apertur"] = 1.0 + with pytest.raises(ValueError): + V4Config(**raw) + + +def test_beads_prior_requires_a_free_parameter(): + raw = yaml.safe_load(EXAMPLE_4_3) + raw["imaging_system"]["detection"]["numerical_aperture"] = 1.2 + raw["imaging_system"]["detection"].pop("pupil_aberrations") + with pytest.raises(ValueError, match="beads"): + V4Config(**raw) + + +def test_grid_point_count_warning(): + raw = yaml.safe_load(EXAMPLE_4_3) + raw["imaging_system"]["detection"]["pupil_aberrations"] = {"zernike_up_to_order": {"order": 4, "tolerance": 0.05}} + raw["solver"]["calibration"]["search"] = [{"grid": {"points": 7}}] + with pytest.warns(UserWarning, match="grid search stage evaluates"): + V4Config(**raw) # 7^13 points (12 Zernikes + free NA) + + +def test_tiled_all_free_sharpness_warning(): + raw = yaml.safe_load(EXAMPLE_4_3) + raw["imaging_system"]["field_dependence"] = {"tiled": {"size_yx": [64, 64]}} + raw["solver"]["calibration"] = {"score": "sharpness", "search": [{"gradient": {"max_iter": 10}}]} + raw["object"] = {"prior": "tikhonov", "weight": 1e-3} + with pytest.warns(UserWarning, match="overfitting"): + V4Config(**raw) diff --git a/tests/test_benchmark_reference.py b/tests/test_benchmark_reference.py index 3251ea86..6d997e5b 100644 --- a/tests/test_benchmark_reference.py +++ b/tests/test_benchmark_reference.py @@ -63,7 +63,7 @@ def test_defaults_to_none(self): } -def _write_optimized_config(case_dir: Path, tf_values: dict): +def _write_optimized_config(case_dir: Path, tf_values: dict, filename: str = "config_optimized.yml"): cfg = { "reconstruction_dimension": 2, "input_channel_names": ["BF"], @@ -81,7 +81,7 @@ def _write_optimized_config(case_dir: Path, tf_values: dict): }, } case_dir.mkdir(parents=True, exist_ok=True) - (case_dir / "config_optimized.yml").write_text(yaml.dump(cfg)) + (case_dir / filename).write_text(yaml.dump(cfg)) class TestMetricChecks: @@ -138,6 +138,13 @@ def test_one_fails(self, tmp_path): assert result["all_pass"] is False assert result["per_ref"]["parameter.z_focus_offset"]["pass"] is False + def test_migrated_mode_filename_accepted(self, tmp_path): + """In migrated-config mode the optimizer writes config_v4_optimized.yml.""" + _write_optimized_config(tmp_path, {"z_focus_offset": -0.05}, filename="config_v4_optimized.yml") + refs = {"parameter.z_focus_offset": ReferenceBound(min=-0.25, max=0.25)} + result = check_reference(tmp_path, {}, refs) + assert result["all_pass"] is True + def test_raises_when_no_optimized_config(self, tmp_path): refs = {"parameter.z_focus_offset": ReferenceBound(min=-0.25, max=0.25)} with pytest.raises(ValueError, match="not written"): diff --git a/tests/test_benchmark_runner.py b/tests/test_benchmark_runner.py index cbe6f576..46fae8e5 100644 --- a/tests/test_benchmark_runner.py +++ b/tests/test_benchmark_runner.py @@ -2,11 +2,17 @@ import shutil +import numpy as np import pytest import yaml from benchmarks.config import ExperimentConfig, PhantomConfig, load_experiment -from benchmarks.runner import _build_phantom, run_synthetic_case +from benchmarks.runner import ( + _build_phantom, + _migrate_config_for_execution, + _read_zarr, + run_synthetic_case, +) class TestPhantomConfig: @@ -208,3 +214,58 @@ def test_fluorescence_end_to_end(self, tmp_path): assert "image_quality" in metrics assert "with_phantom" in metrics assert (case_dir / "cli_command.sh").exists() + + +def test_migrate_helper_skips_configs_that_stay_3x(tmp_path, pytestconfig): + """Birefringence and already-4.0 configurations run unmigrated.""" + birefringence = pytestconfig.rootpath / "docs" / "examples" / "cli" / "configs" / "birefringence_3d.yml" + config_path = tmp_path / "config.yml" + config_path.write_text(birefringence.read_text()) + assert _migrate_config_for_execution(config_path) is None + + config_path.write_text(yaml.dump({"version": 4})) + assert _migrate_config_for_execution(config_path) is None + + +@pytest.mark.skipif(not shutil.which("wo"), reason="wo CLI not installed") +class TestMigratedConfigMode: + """``--migrate-configs``: the design section 2 regression mode.""" + + def test_migrated_config_reproduces_reconstruction(self, tmp_path): + phantom_config = PhantomConfig( + function="single_bead", + shape=(16, 32, 32), + pixel_sizes=(0.25, 0.1, 0.1), + bead_radius_um=1.0, + fluorescence_intensity=1.0, + ) + recon_config = { + "reconstruction_dimension": 3, + "input_channel_names": ["GFP"], + "fluorescence": { + "transfer_function": { + "yx_pixel_size": 0.1, + "z_pixel_size": 0.25, + "wavelength_emission": 0.532, + "index_of_refraction_media": 1.3, + "numerical_aperture_detection": 1.2, + }, + }, + } + + results = {} + for label, migrate in (("plain", False), ("migrated", True)): + case_dir = tmp_path / label + metrics = run_synthetic_case( + phantom_config, + recon_config, + case_dir, + modality="fluorescence", + migrate_config=migrate, + ) + results[label] = (metrics, _read_zarr(case_dir / "reconstruction.zarr")) + + assert results["migrated"][0]["config_migrated"] is True + assert "config_migrated" not in results["plain"][0] + assert (tmp_path / "migrated" / "config_v4.yml").exists() + np.testing.assert_array_equal(results["plain"][1].numpy(), results["migrated"][1].numpy()) diff --git a/waveorder/calibrate/__init__.py b/waveorder/calibrate/__init__.py new file mode 100644 index 00000000..e441e114 --- /dev/null +++ b/waveorder/calibrate/__init__.py @@ -0,0 +1,79 @@ +"""Imaging-system calibration: parameter vectors, PSF/OTF synthesis, scores. + +This package holds the numerics behind ``calibrate-tf``: + +* :mod:`waveorder.calibrate.zernike_names` — named Zernike terms + (``defocus`` ... ``spherical``, ``z``) to Noll indices. +* :mod:`waveorder.calibrate.synthesis` — differentiable fluorescence + PSF/OTF synthesis from a parameter vector (Zernike coefficients in waves). +* :mod:`waveorder.calibrate.scores` — Whittle/Gaussian evidence and + related calibration scores. +* :mod:`waveorder.calibrate.tfzarr` — the ``tf.zarr`` store. +* :mod:`waveorder.calibrate.run`, :mod:`waveorder.calibrate.evidence_search` + — configuration-driven searches (loaded lazily: they depend on + :mod:`waveorder.config`, which itself imports the leaf modules above, so + importing them here eagerly would be circular). +""" + +import importlib + +from waveorder.calibrate.scores import ( + calibration_penalty, + whittle_evidence, +) +from waveorder.calibrate.synthesis import ( + fluorescence_otf, + fluorescence_psf, +) +from waveorder.calibrate.tfzarr import ( + TransferFunctionStore, + read_tf_zarr, + write_tf_zarr, +) +from waveorder.calibrate.zernike_names import ( + NAMED_TERMS, + name_to_noll, + noll_indices_up_to_order, + noll_to_name, + validate_zernike_name, +) + +_LAZY_EXPORTS = { + "CalibrationResult": "waveorder.calibrate.run", + "build_recovery_settings": "waveorder.calibrate.run", + "run_zernike_calibration": "waveorder.calibrate.run", + "zernike_calibration_requested": "waveorder.calibrate.run", + "EvidenceCalibrationResult": "waveorder.calibrate.evidence_search", + "gcv_score": "waveorder.calibrate.evidence_search", + "run_evidence_calibration": "waveorder.calibrate.evidence_search", +} + +__all__ = [ + "NAMED_TERMS", + "CalibrationResult", + "EvidenceCalibrationResult", + "TransferFunctionStore", + "build_recovery_settings", + "calibration_penalty", + "fluorescence_otf", + "fluorescence_psf", + "gcv_score", + "name_to_noll", + "noll_indices_up_to_order", + "noll_to_name", + "read_tf_zarr", + "run_evidence_calibration", + "run_zernike_calibration", + "validate_zernike_name", + "whittle_evidence", + "write_tf_zarr", + "zernike_calibration_requested", +] + + +def __getattr__(name: str): + """Load the config-dependent calibration entry points on first use.""" + module_name = _LAZY_EXPORTS.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(module_name), name) diff --git a/waveorder/calibrate/beads.py b/waveorder/calibrate/beads.py new file mode 100644 index 00000000..502bd431 --- /dev/null +++ b/waveorder/calibrate/beads.py @@ -0,0 +1,321 @@ +"""Bead calibration: ``object: {prior: {beads: {diameter}}}`` (design example 4.3). + +The object model is a sparse set of spheres of known diameter at unknown +subpixel positions with unknown amplitudes over a constant background:: + + g = H * (sum_i a_i s(r - r_i)) + b + n, + +where ``s`` is the sphere of the configured diameter. With this prior the +object step is point localization and ``solver.reconstruction`` is ignored: +positions, amplitudes, background, and the free imaging-system parameters +are fit jointly by maximum a posteriori (Gaussian noise), with the +calibration penalty from the parameter tolerances. The reported +``parameters_covariance`` is the Laplace curvature with respect to the +imaging-system parameters at the joint optimum, with the bead parameters +held fixed (a lower bound on the marginal covariance). + +Everything runs in Fourier space and is differentiable: the beads enter as +phase ramps ``a_i exp(-2 pi i k . r_i)`` on the sphere's analytic form +factor, so subpixel positions are exact. +""" + +from __future__ import annotations + +import math +import warnings +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as functional + +from waveorder.calibrate.evidence_search import ParameterSpec, build_parameter_specs +from waveorder.calibrate.synthesis import fluorescence_otf +from waveorder.calibrate.zernike_names import name_to_noll +from waveorder.config.schema import ( + FreeParameter, + GaussianNoise, + PriorBeads, + SearchGradient, + V4Config, +) + + +@dataclass +class BeadCalibrationResult: + """Output of :func:`run_bead_calibration`.""" + + parameters: np.ndarray # [1, P] imaging-system parameters (waves) + parameter_names: list[str] + parameters_covariance: np.ndarray # [1, P, P] Laplace (beads held fixed) + positions_zyx: np.ndarray # [N, 3] fitted bead positions in voxels + amplitudes: np.ndarray # [N] + background: float + noise_sigma: float + score_trace: list[float] + best_score: float + + +def sphere_form_factor( + zyx_shape: tuple[int, int, int], + yx_pixel_size: float, + z_pixel_size: float, + diameter: float, +) -> torch.Tensor: + """Fourier transform of a solid sphere, normalized to 1 at DC. + + Parameters + ---------- + zyx_shape : tuple[int, int, int] + Grid shape ``(Z, Y, X)``. + yx_pixel_size, z_pixel_size : float + Voxel sizes in um. + diameter : float + Sphere diameter in um. + + Returns + ------- + Tensor + Real array, FFT frequency order: ``3 (sin u - u cos u) / u^3`` with + ``u = pi * diameter * |k|``. + """ + fz = torch.fft.fftfreq(zyx_shape[0], d=z_pixel_size) + fy = torch.fft.fftfreq(zyx_shape[1], d=yx_pixel_size) + fx = torch.fft.fftfreq(zyx_shape[2], d=yx_pixel_size) + kz, ky, kx = torch.meshgrid(fz, fy, fx, indexing="ij") + u = math.pi * diameter * torch.sqrt(kz**2 + ky**2 + kx**2) + small = u < 1e-4 + u_safe = torch.where(small, torch.ones_like(u), u) + form = 3.0 * (torch.sin(u_safe) - u_safe * torch.cos(u_safe)) / u_safe**3 + return torch.where(small, torch.ones_like(u), form) + + +def detect_beads( + zyx_data: torch.Tensor, + max_beads: int = 50, + min_separation: int = 5, + snr_threshold: float = 6.0, + adaptive: bool = True, +) -> torch.Tensor: + """Initial bead positions: local maxima well above the background. + + Parameters + ---------- + zyx_data : Tensor + 3D stack ``(Z, Y, X)``. + max_beads : int + Keep at most this many peaks, brightest first. + min_separation : int + Non-maximum suppression window in voxels. + snr_threshold : float + Peaks must exceed ``median + snr_threshold * 1.4826 * MAD``. + adaptive : bool + When nothing clears ``snr_threshold``, relax it down to half its + value (with a warning) before giving up. + + Returns + ------- + Tensor + ``[N, 3]`` integer ``(z, y, x)`` peak positions. + """ + median = zyx_data.median() + mad = (zyx_data - median).abs().median().clamp(min=1e-12) + + kernel = 2 * (min_separation // 2) + 1 + pooled = functional.max_pool3d(zyx_data[None, None], kernel_size=kernel, stride=1, padding=kernel // 2)[0, 0] + is_local_max = zyx_data == pooled + + thresholds = [snr_threshold] + if adaptive: + thresholds += [0.75 * snr_threshold, 0.5 * snr_threshold] + positions = torch.zeros((0, 3), dtype=torch.int64) + for i, candidate_snr in enumerate(thresholds): + threshold = median + candidate_snr * 1.4826 * mad + positions = (is_local_max & (zyx_data > threshold)).nonzero() + if positions.shape[0] > 0: + if i > 0: + warnings.warn( + f"no beads above median + {snr_threshold:g} * 1.4826 * MAD; relaxed the " + f"detection threshold to {candidate_snr:g} robust noise units. Set " + f"object.prior.beads.snr_threshold to control this.", + UserWarning, + ) + break + if positions.shape[0] == 0: + raise ValueError( + f"no beads found above median + {thresholds[-1]:g} * 1.4826 * MAD. Lower " + f"object.prior.beads.snr_threshold (currently {snr_threshold:g}), check " + f"data.input_channel_names, or verify the volume contains beads." + ) + intensities = zyx_data[positions[:, 0], positions[:, 1], positions[:, 2]] + order = torch.argsort(intensities, descending=True)[:max_beads] + return positions[order] + + +def run_bead_calibration( + cfg: V4Config, + zyx_data: np.ndarray | torch.Tensor, + gradient_lr: float = 0.01, + max_beads: int | None = None, +) -> BeadCalibrationResult: + """Fit free imaging-system parameters to a bead volume. + + Parameters + ---------- + cfg : V4Config + Fluorescence configuration with ``object: {prior: {beads: ...}}`` + and at least one free ``pupil_aberrations`` term. + zyx_data : ndarray or Tensor + 3D bead stack ``(Z, Y, X)``. + gradient_lr : float + Adam step size for the imaging-system parameters (waves). Bead + positions use ``10 x`` this (voxels move faster than waves). + max_beads : int, optional + At most this many beads are fit; None reads + ``object.prior.beads.max_beads``. + + Returns + ------- + BeadCalibrationResult + """ + if not isinstance(cfg.object.prior, PriorBeads): + raise ValueError("run_bead_calibration requires object.prior: {beads: {diameter: ...}}") + if cfg.imaging_system.contrast != "fluorescence": + raise NotImplementedError("bead calibration supports fluorescence contrast") + if cfg.imaging_system.field_dependence != "shift_invariant": + raise NotImplementedError("bead calibration is shift-invariant for now") + if cfg.noise.kind != "gaussian": + raise NotImplementedError(f"bead calibration supports gaussian noise for now, got {cfg.noise.kind}") + + specs: list[ParameterSpec] = build_parameter_specs(cfg) + names = [spec.name for spec in specs] + + detection = cfg.imaging_system.detection + if detection.wavelength is None: + raise ValueError("imaging_system.detection.wavelength is required") + yx = cfg.data.pixel_size.yx + yx_pixel_size = float(yx) if isinstance(yx, (int, float)) else float(yx.y) + z_pixel_size = float(cfg.data.pixel_size.z) + + data = torch.as_tensor(np.asarray(zyx_data), dtype=torch.float32) + zyx_shape = tuple(data.shape) + + optics_kwargs = dict( + zyx_shape=zyx_shape, + yx_pixel_size=yx_pixel_size, + z_pixel_size=z_pixel_size, + wavelength_emission=detection.wavelength, + index_of_refraction_media=cfg.imaging_system.index_of_refraction_media, + numerical_aperture_detection=( + detection.numerical_aperture.nominal + if isinstance(detection.numerical_aperture, FreeParameter) + else detection.numerical_aperture + ), + ) + form_factor = sphere_form_factor(zyx_shape, yx_pixel_size, z_pixel_size, cfg.object.prior.beads.diameter) + + # Initialization from the data. + beads_prior = cfg.object.prior.beads + peak_positions = detect_beads( + data, + max_beads=max_beads if max_beads is not None else beads_prior.max_beads, + min_separation=beads_prior.min_separation, + snr_threshold=beads_prior.snr_threshold, + ) + background0 = float(data.median()) + amplitudes0 = (data[peak_positions[:, 0], peak_positions[:, 1], peak_positions[:, 2]] - background0).clamp(min=1e-6) + + noise = cfg.noise.gaussian if cfg.noise.gaussian is not None else GaussianNoise() + estimate_sigma = noise.sigma == "estimate" + sigma0 = float((data - data.median()).abs().median() * 1.4826) if estimate_sigma else float(noise.sigma) + sigma0 = max(sigma0, 1e-6) + + # Frequency grids for the bead phase ramps (cycles per voxel). + freqs = [torch.fft.fftfreq(n).to(torch.float64) for n in zyx_shape] + grids = torch.meshgrid(*freqs, indexing="ij") + + values = torch.tensor([spec.nominal for spec in specs], dtype=torch.float64, requires_grad=True) + positions = peak_positions.to(torch.float64).clone().requires_grad_(True) + log_amplitudes = torch.log(amplitudes0.to(torch.float64)).clone().requires_grad_(True) + background = torch.tensor(background0, dtype=torch.float64, requires_grad=True) + log_sigma = torch.tensor(math.log(sigma0), dtype=torch.float64, requires_grad=estimate_sigma) + + n_voxels = data.numel() + + def total_score() -> torch.Tensor: + coefficients = {name_to_noll(spec.name): value for spec, value in zip(specs, values)} + otf = fluorescence_otf(zernike_coefficients=coefficients, **optics_kwargs) + phase = ( + grids[0] * positions[:, 0, None, None, None] + + grids[1] * positions[:, 1, None, None, None] + + grids[2] * positions[:, 2, None, None, None] + ) + ramps = torch.exp(-2j * math.pi * phase) + object_fft = (torch.exp(log_amplitudes)[:, None, None, None] * ramps).sum(dim=0) * form_factor + predicted = torch.real(torch.fft.ifftn(object_fft * otf)) + background + sigma2 = torch.exp(2.0 * log_sigma) + misfit = ((predicted - data) ** 2).sum() / (2.0 * sigma2) + 0.5 * n_voxels * torch.log(2.0 * math.pi * sigma2) + penalty = torch.zeros((), dtype=torch.float64) + for i, spec in enumerate(specs): + if spec.tolerance != "free": + penalty = penalty + (values[i] - spec.nominal) ** 2 / (2.0 * spec.tolerance**2) + return misfit + penalty + + stages = cfg.solver.calibration.search if cfg.solver.calibration is not None else [] + max_iter = 200 + for stage in stages: + if isinstance(stage, SearchGradient): + max_iter = stage.gradient.max_iter + else: + raise NotImplementedError("bead calibration runs a single gradient stage for now") + + optimized = [ + {"params": [values], "lr": gradient_lr}, + {"params": [positions], "lr": 10.0 * gradient_lr}, + {"params": [log_amplitudes, background], "lr": 10.0 * gradient_lr}, + ] + if estimate_sigma: + optimized.append({"params": [log_sigma], "lr": gradient_lr}) + optimizer = torch.optim.Adam(optimized) + + score_trace: list[float] = [] + for _ in range(max_iter): + optimizer.zero_grad() + score = total_score() + score.backward() + optimizer.step() + with torch.no_grad(): + for i, spec in enumerate(specs): + if spec.bounds is not None: + values.data[i] = values.data[i].clamp(spec.bounds[0], spec.bounds[1]) + score_trace.append(float(score.detach())) + + best_score = score_trace[-1] + + def score_of_values(candidate: torch.Tensor) -> torch.Tensor: + nonlocal values + saved = values + values = candidate + try: + return total_score() + finally: + values = saved + + final_values = values.detach() + try: + hessian = torch.autograd.functional.hessian(score_of_values, final_values) + covariance = torch.linalg.inv(hessian).numpy() + except RuntimeError: + covariance = np.full((len(specs), len(specs)), np.nan) + + return BeadCalibrationResult( + parameters=final_values.numpy()[None, :], + parameter_names=names, + parameters_covariance=covariance[None, :, :], + positions_zyx=positions.detach().numpy(), + amplitudes=torch.exp(log_amplitudes.detach()).numpy(), + background=float(background.detach()), + noise_sigma=float(torch.exp(log_sigma.detach())), + score_trace=score_trace, + best_score=best_score, + ) diff --git a/waveorder/calibrate/evidence_search.py b/waveorder/calibrate/evidence_search.py new file mode 100644 index 00000000..6ea2187a --- /dev/null +++ b/waveorder/calibrate/evidence_search.py @@ -0,0 +1,480 @@ +"""Evidence- and GCV-scored calibration of a shift-invariant fluorescence system. + +Implements the MAPk estimator of the design document section 6 for the +Gaussian case: the score is the exact Whittle negative log evidence +(:func:`waveorder.calibrate.scores.whittle_evidence`) or the closed-form +GCV of the Wiener filter, plus the calibration penalty from the parameter +tolerances. The search runs the configured stages in order (``grid``, +``random``, ``gradient``), each starting from the previous stage's best +point, and the parameter covariance is the Laplace approximation (inverse +Hessian of the total score at the optimum). + +Parameters are the free ``pupil_aberrations`` terms (Zernike coefficients +in waves). ``noise.gaussian.sigma: estimate`` and ``object.weight: auto`` +add ``log sigma`` / ``log lambda`` as auxiliary variables of the gradient +stage; they are reported in the provenance, not in ``parameters``. +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, field + +import numpy as np +import torch + +from waveorder.calibrate.scores import whittle_evidence +from waveorder.calibrate.synthesis import fluorescence_otf +from waveorder.calibrate.tfzarr import TransferFunctionStore +from waveorder.calibrate.zernike_names import name_to_noll, noll_indices_up_to_order, noll_to_name +from waveorder.config.schema import ( + FreeParameter, + GaussianNoise, + SearchGradient, + SearchGrid, + SearchRandom, + V4Config, + free_parameters, +) + + +@dataclass +class ParameterSpec: + """One free scalar parameter of the imaging system.""" + + name: str # canonical Zernike name, e.g. "astigmatism_0" + nominal: float + tolerance: float | str # Gaussian prior std, or "free" + bounds: tuple[float, float] | None + + +@dataclass +class EvidenceCalibrationResult: + """Output of :func:`run_evidence_calibration`.""" + + parameters: np.ndarray # [1, P] + parameter_names: list[str] + parameters_covariance: np.ndarray # [1, P, P] Laplace + score_trace: list[float] # per gradient iteration + score_grid: np.ndarray | None # [n_points, P + 1]: coordinates + score + noise_sigma: float + weight: float + best_score: float + extras: dict = field(default_factory=dict) + + +def gcv_score( + data_fft: torch.Tensor, + otf: torch.Tensor, + regularization_strength: float | torch.Tensor, +) -> torch.Tensor: + """Closed-form GCV of the Wiener filter ``A = |H|^2 / (|H|^2 + lambda)``. + + ``GCV = ||g - A g||^2 / (1 - tr(A) / N)^2`` with both terms diagonal per + spatial frequency, so no trace estimation is needed. Independent of the + noise level (Wahba). + + Parameters + ---------- + data_fft : Tensor + ``torch.fft.fftn(g, norm="ortho")``. + otf : Tensor + Transfer function, same shape, FFT order. + regularization_strength : float or Tensor + Tikhonov weight ``lambda``. + + Returns + ------- + Tensor + Scalar GCV score (lower is better). + """ + h2 = torch.abs(otf) ** 2 + lam = regularization_strength + shrink = lam / (h2 + lam) # 1 - A, per frequency + n = data_fft.numel() + residual_power = (torch.abs(data_fft) ** 2 * shrink**2).sum() + denominator = (shrink.sum() / n) ** 2 + return residual_power / denominator + + +def build_parameter_specs(cfg: V4Config) -> list[ParameterSpec]: + """Free Zernike parameters of the configuration, in Noll order. + + Named entries override the ``zernike_up_to_order`` shorthand for their + term (design document section 4.4). + """ + aberrations = cfg.imaging_system.detection.pupil_aberrations + if aberrations is None: + raise ValueError("no pupil_aberrations block; nothing to estimate on the evidence path") + + by_noll: dict[int, ParameterSpec] = {} + shorthand = aberrations.zernike_up_to_order + if shorthand is not None: + for j in noll_indices_up_to_order(shorthand.order): + tolerance = shorthand.tolerance if shorthand.tolerance == "free" else float(shorthand.tolerance) + by_noll[j] = ParameterSpec(noll_to_name(j), 0.0, tolerance, shorthand.bounds) + for name, value in aberrations.terms.items(): + j = name_to_noll(name) + if isinstance(value, FreeParameter): + tolerance = value.tolerance if value.tolerance == "free" else float(value.tolerance) + by_noll[j] = ParameterSpec(noll_to_name(j), value.nominal, tolerance, value.bounds) + else: + # Fixed named term: keep it out of the estimated set; nonzero + # fixed terms are folded into the pupil as constants. + by_noll.pop(j, None) + + non_zernike = { + name for name in free_parameters(cfg.imaging_system) if not name.startswith("detection.pupil_aberrations") + } + if non_zernike: + raise NotImplementedError( + f"the evidence path estimates pupil aberrations only for now; also free: {sorted(non_zernike)}" + ) + if not by_noll: + raise ValueError("no pupil-aberration term has a tolerance") + return [by_noll[j] for j in sorted(by_noll)] + + +def _fixed_coefficients(cfg: V4Config, estimated: list[ParameterSpec]) -> dict[int, float]: + """Fixed nonzero Zernike terms, folded into every candidate pupil.""" + aberrations = cfg.imaging_system.detection.pupil_aberrations + estimated_nolls = {name_to_noll(spec.name) for spec in estimated} + fixed = {} + for name, value in aberrations.terms.items(): + j = name_to_noll(name) + if not isinstance(value, FreeParameter) and j not in estimated_nolls and float(value) != 0.0: + fixed[j] = float(value) + return fixed + + +def _stage_ranges(specs: list[ParameterSpec], centers: torch.Tensor) -> list[tuple[float, float]]: + """Grid range per parameter: bounds, else center +/- 2 tolerance. + + ``centers`` are the (possibly ``-t``-moved) prior centers, so a grid + stage after a ``-t`` prior searches around the prior parameters rather + than the configuration's nominal values. + """ + ranges = [] + for spec, center in zip(specs, centers.tolist()): + if spec.bounds is not None: + ranges.append(spec.bounds) + elif spec.tolerance != "free": + ranges.append((center - 2.0 * spec.tolerance, center + 2.0 * spec.tolerance)) + else: + raise ValueError(f"grid/random stage needs bounds for 'free' parameter {spec.name!r}") + return ranges + + +def run_evidence_calibration( + cfg: V4Config, + zyx_data: np.ndarray | torch.Tensor, + gradient_lr: float = 0.01, + prior_store: TransferFunctionStore | None = None, + init_store: TransferFunctionStore | None = None, +) -> EvidenceCalibrationResult: + """Estimate free pupil-aberration parameters by evidence or GCV. + + Parameters + ---------- + cfg : V4Config + Shift-invariant fluorescence configuration with free + ``pupil_aberrations`` terms and ``solver.calibration.score`` in + {``evidence``, ``gcv``}. + zyx_data : ndarray or Tensor + 3D fluorescence stack ``(Z, Y, X)``. + gradient_lr : float + Adam step size for the gradient stage, in waves. + prior_store : TransferFunctionStore, optional + A ``tf.zarr`` passed with ``-t``: its ``parameters`` move the prior + center (and the search start), and a non-zero + ``parameters_covariance`` replaces the configuration's tolerances + with a correlated Gaussian prior for the parameters it covers + (design document section 3). + init_store : TransferFunctionStore, optional + A ``tf.zarr`` passed with ``--init``: sets the starting point only, + without touching the prior. Takes precedence over ``prior_store`` + for initialization. + + Returns + ------- + EvidenceCalibrationResult + + Notes + ----- + The evidence score drops the DC (k = 0) term (``exclude_dc=True``): + real data carries a constant background the Tikhonov prior does not + model, so the reported score differs from the exact ``-log p(g | H)`` + by the DC term. Scores are comparable across candidate systems on the + same data. + + ``parameters_covariance`` is the inverse Hessian of the total score at + the optimum with ``sigma`` and ``lambda`` held fixed at their final + values (a profile-curvature approximation); when ``sigma: estimate`` + or ``weight: auto`` is active the true marginal covariance is somewhat + wider. + """ + system = cfg.imaging_system + if system.contrast != "fluorescence": + raise NotImplementedError("the evidence path supports fluorescence contrast for now") + if system.field_dependence != "shift_invariant": + raise NotImplementedError("the evidence path is shift-invariant for now; tiled follows") + calibration = cfg.solver.calibration + score_name = calibration.score if calibration is not None else "evidence" + if score_name not in ("evidence", "gcv"): + raise ValueError(f"run_evidence_calibration handles score evidence|gcv, got {score_name}") + + specs = build_parameter_specs(cfg) + names = [spec.name for spec in specs] + nominals = torch.tensor([spec.nominal for spec in specs], dtype=torch.float64) + + # -t: move the prior center (and start) to the store's parameters; a + # non-zero covariance replaces the per-parameter tolerances with a + # correlated Gaussian prior for the parameters it covers. An all-zero + # covariance block (e.g. from simulate) carries no uncertainty + # information and leaves the configuration's tolerances in force. + name_to_index = {name: i for i, name in enumerate(names)} + covered: list[int] = [] + prior_precision: torch.Tensor | None = None + if prior_store is not None: + if prior_store.n_tiles != 1: + raise NotImplementedError("a tiled tf.zarr prior (-t) is not supported on the evidence path yet") + store_indices = [] + for store_index, store_name in enumerate(prior_store.parameter_names): + index = name_to_index.get(store_name) + if index is not None: + nominals[index] = float(prior_store.parameters[0, store_index]) + store_indices.append(store_index) + covered.append(index) + if covered and prior_store.parameters_covariance is not None: + block = prior_store.parameters_covariance[0][np.ix_(store_indices, store_indices)] + if np.any(block): + block_tensor = torch.as_tensor(block, dtype=torch.float64) + try: + prior_precision = torch.linalg.inv(block_tensor) + except torch.linalg.LinAlgError as error: + raise ValueError( + f"the -t parameters_covariance block over {len(store_indices)} covered " + f"parameters is singular and cannot serve as a Gaussian prior; " + f"re-estimate it, or drop the singular parameters from the store" + ) from error + covariance_covered = covered if prior_precision is not None else [] + covered_set = set(covariance_covered) + + start = nominals.clone() + if init_store is not None: + if init_store.n_tiles != 1: + raise NotImplementedError("a tiled --init store is not supported on the evidence path yet") + for store_index, store_name in enumerate(init_store.parameter_names): + index = name_to_index.get(store_name) + if index is not None: + start[index] = float(init_store.parameters[0, store_index]) + + detection = system.detection + if detection.wavelength is None: + raise ValueError("imaging_system.detection.wavelength is required") + yx = cfg.data.pixel_size.yx + optics_kwargs = dict( + zyx_shape=tuple(zyx_data.shape), + yx_pixel_size=float(yx) if isinstance(yx, (int, float)) else float(yx.y), + z_pixel_size=cfg.data.pixel_size.z, + wavelength_emission=detection.wavelength, + index_of_refraction_media=system.index_of_refraction_media, + numerical_aperture_detection=( + detection.numerical_aperture.nominal + if isinstance(detection.numerical_aperture, FreeParameter) + else detection.numerical_aperture + ), + ) + fixed_coefficients = _fixed_coefficients(cfg, specs) + + data = torch.as_tensor(np.asarray(zyx_data), dtype=torch.float32) + data_fft = torch.fft.fftn(data, norm="ortho") + + # sigma and lambda: numbers, or jointly maximized ("estimate" / "auto"). + if cfg.noise.kind != "gaussian": + raise NotImplementedError(f"the evidence path supports gaussian noise for now, got {cfg.noise.kind}") + noise = cfg.noise.gaussian if cfg.noise.gaussian is not None else GaussianNoise() + estimate_sigma = noise.sigma == "estimate" + if estimate_sigma: + sigma0 = float(data.std()) # refined by _profile_sigma_weight below + elif isinstance(noise.sigma, (int, float)): + sigma0 = float(noise.sigma) + else: + raise NotImplementedError("noise.gaussian.sigma from dark_frames is not implemented yet") + + auto_weight = cfg.object.weight == "auto" + if auto_weight: + weight0 = 1e-3 + elif isinstance(cfg.object.weight, (int, float)): + weight0 = float(cfg.object.weight) + else: + raise ValueError(f"object.weight: {cfg.object.weight} is not usable on the evidence path") + + def total_score(values: torch.Tensor, log_sigma: torch.Tensor, log_weight: torch.Tensor) -> torch.Tensor: + coefficients: dict[int, float | torch.Tensor] = dict(fixed_coefficients) + for spec, value in zip(specs, values): + coefficients[name_to_noll(spec.name)] = value + otf = fluorescence_otf(zernike_coefficients=coefficients, **optics_kwargs) + if score_name == "evidence": + data_term = whittle_evidence( + data_fft=data_fft, + otf=otf, + noise_sigma=torch.exp(log_sigma), + regularization_strength=torch.exp(log_weight), + exclude_dc=True, + ) + else: + data_term = gcv_score(data_fft, otf, torch.exp(log_weight)) + penalty = torch.zeros((), dtype=torch.float64) + for i, spec in enumerate(specs): + if i in covered_set: + continue # covered by the -t covariance's matrix term below + if spec.tolerance != "free": + penalty = penalty + (values[i] - nominals[i]) ** 2 / (2.0 * spec.tolerance**2) + if prior_precision is not None: + residual = values[covariance_covered] - nominals[covariance_covered] + penalty = penalty + 0.5 * residual @ prior_precision @ residual + return data_term + penalty + + def clamp(values: torch.Tensor) -> torch.Tensor: + clamped = values.clone() + for i, spec in enumerate(specs): + if spec.bounds is not None: + clamped[i] = clamped[i].clamp(spec.bounds[0], spec.bounds[1]) + return clamped + + current = start.clone() + log_sigma = torch.tensor(math.log(sigma0), dtype=torch.float64) + log_weight = torch.tensor(math.log(weight0), dtype=torch.float64) + score_trace: list[float] = [] + score_grid: np.ndarray | None = None + + def profile_sigma_weight(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Coarse log-grid profile of (sigma, lambda) at fixed coefficients. + + A joint gradient from a bad (sigma, lambda) start settles into poor + local optima (sigma absorbs the model mismatch before the + coefficients can move), so the estimated nuisance scales are + re-profiled before each gradient stage. + """ + data_scale = float(data.std()) + sigma_candidates = [math.exp(float(log_sigma))] if not estimate_sigma else [] + weight_candidates = [math.exp(float(log_weight))] if not auto_weight else [] + if estimate_sigma: + sigma_candidates = list(data_scale * np.geomspace(1e-3, 1.0, 13)) + if auto_weight: + weight_candidates = list(np.geomspace(1e-6, 1.0, 13)) + best = (math.inf, float(log_sigma), float(log_weight)) + with torch.no_grad(): + for sigma_value in sigma_candidates: + for weight_value in weight_candidates: + s = float( + total_score( + values, + torch.tensor(math.log(sigma_value), dtype=torch.float64), + torch.tensor(math.log(weight_value), dtype=torch.float64), + ) + ) + if s < best[0]: + best = (s, math.log(sigma_value), math.log(weight_value)) + return ( + torch.tensor(best[1], dtype=torch.float64), + torch.tensor(best[2], dtype=torch.float64), + ) + + if estimate_sigma or auto_weight: + log_sigma, log_weight = profile_sigma_weight(current) + + stages = calibration.search if calibration is not None else [SearchGradient.model_validate({"gradient": {}})] + for stage in stages: + if isinstance(stage, SearchGrid): + ranges = _stage_ranges(specs, nominals) + points = stage.grid.points + axes = [] + for spec, (lo, hi) in zip(specs, ranges): + n = points if isinstance(points, int) else points.get(spec.name, 7) + axes.append(np.linspace(lo, hi, n)) + rows = [] + best = (None, math.inf) + with torch.no_grad(): + for combo in itertools.product(*axes): + values = torch.tensor(combo, dtype=torch.float64) + s = float(total_score(values, log_sigma, log_weight)) + rows.append([*combo, s]) + if s < best[1]: + best = (values, s) + current = best[0] + score_grid = np.asarray(rows) + elif isinstance(stage, SearchRandom): + generator = torch.Generator().manual_seed(stage.random.seed) + ranges = _stage_ranges(specs, nominals) + best = (current, float(total_score(current, log_sigma, log_weight))) + with torch.no_grad(): + for _ in range(stage.random.samples): + draw = current.clone() + for i, spec in enumerate(specs): + if spec.tolerance == "free": + lo, hi = ranges[i] + draw[i] = lo + (hi - lo) * torch.rand((), generator=generator, dtype=torch.float64) + else: + draw[i] = spec.nominal + spec.tolerance * torch.randn( + (), generator=generator, dtype=torch.float64 + ) + draw = clamp(draw) + s = float(total_score(draw, log_sigma, log_weight)) + if s < best[1]: + best = (draw, s) + current = best[0] + elif isinstance(stage, SearchGradient): + if stage.gradient.method != "split": + raise NotImplementedError(f"gradient method {stage.gradient.method!r} (use split)") + if estimate_sigma or auto_weight: + log_sigma, log_weight = profile_sigma_weight(current) + values = current.clone().requires_grad_(True) + optimized = [values] + if estimate_sigma and score_name == "evidence": + log_sigma = log_sigma.clone().requires_grad_(True) + optimized.append(log_sigma) + if auto_weight: + log_weight = log_weight.clone().requires_grad_(True) + optimized.append(log_weight) + optimizer = torch.optim.Adam(optimized, lr=gradient_lr) + for _ in range(stage.gradient.max_iter): + optimizer.zero_grad() + s = total_score(values, log_sigma, log_weight) + s.backward() + optimizer.step() + with torch.no_grad(): + values.data = clamp(values.data) + score_trace.append(float(s.detach())) + current = values.detach() + log_sigma = log_sigma.detach() + log_weight = log_weight.detach() + else: # pragma: no cover - schema normalizes stages + raise ValueError(f"unknown search stage {stage!r}") + + best_score = float(total_score(current, log_sigma, log_weight)) + + # Laplace covariance: inverse Hessian of the total score at the optimum, + # with sigma and lambda held fixed (profile curvature; see Notes above). + def score_of(values: torch.Tensor) -> torch.Tensor: + return total_score(values, log_sigma, log_weight) + + try: + hessian = torch.autograd.functional.hessian(score_of, current) + covariance = torch.linalg.inv(hessian).numpy() + except RuntimeError: + covariance = np.full((len(specs), len(specs)), np.nan) + + return EvidenceCalibrationResult( + parameters=current.numpy()[None, :], + parameter_names=names, + parameters_covariance=covariance[None, :, :], + score_trace=score_trace, + score_grid=score_grid, + noise_sigma=float(torch.exp(log_sigma)), + weight=float(torch.exp(log_weight)), + best_score=best_score, + extras={"score": score_name}, + ) diff --git a/waveorder/calibrate/run.py b/waveorder/calibrate/run.py new file mode 100644 index 00000000..b3c3f019 --- /dev/null +++ b/waveorder/calibrate/run.py @@ -0,0 +1,244 @@ +"""Map a 4.0 configuration onto the existing calibration numerics. + +This is the engine behind ``wo calibrate-tf`` at the "fold the existing +blind paths into the configuration" stage (design document section 9, +step 2): + +* ``pupil_aberrations`` entries with a tolerance, ``score: sharpness`` and a + gradient search run the existing per-tile Zernike recovery + (:func:`waveorder.api.shift_variant_recovery.recover_zernikes`). + ``field_dependence: shift_invariant`` is the same recovery with a single + tile spanning the field of view; ``tiled`` uses the configured tile grid. +* Scores other than ``sharpness``, numeric tolerances, and grid/random + search stages belong to later steps and raise ``NotImplementedError`` + with a pointer to what is missing. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from waveorder.api.shift_variant_recovery import RecoverySettings, recover_zernikes +from waveorder.calibrate.zernike_names import ( + name_to_noll, + noll_indices_up_to_order, + noll_to_name, +) +from waveorder.config.schema import ( + FreeParameter, + SearchGradient, + V4Config, + free_parameters, + nominal_of, +) + +#: Map from the 3.x sharpness loss names to the recovery's loss names. +_LOSS_NAME_MAP = { + "midband_power": "midband", + "total_variation": "tv", + "laplacian_variance": "laplacian_var", + "normalized_variance": "normalized_var", + "spectral_flatness": "spectral_flatness", +} + + +@dataclass +class CalibrationResult: + """Output of :func:`run_zernike_calibration`, ready for ``tf.zarr``.""" + + parameters: np.ndarray # [N_tiles, P] Zernike coefficients in waves + parameter_names: list[str] + tile_grid: dict[str, Any] + reconstruction_zyx: np.ndarray # stitched Wiener deconvolution + loss_history: list[float] + noll_indices: list[int] + recovery_settings: dict + + +def zernike_calibration_requested(cfg: V4Config) -> bool: + """Whether the configuration frees any pupil-aberration term.""" + aberrations = cfg.imaging_system.detection.pupil_aberrations + if aberrations is None: + return False + if aberrations.zernike_up_to_order is not None: + return True + return any(isinstance(value, FreeParameter) for value in aberrations.terms.values()) + + +def _not_yet(reason: str) -> NotImplementedError: + return NotImplementedError(f"calibrate-tf does not support this yet: {reason}") + + +def _resolve_noll_indices(cfg: V4Config) -> list[int]: + """Noll indices to estimate, from the shorthand and the named terms.""" + aberrations = cfg.imaging_system.detection.pupil_aberrations + indices: set[int] = set() + if aberrations.zernike_up_to_order is not None: + shorthand = aberrations.zernike_up_to_order + if shorthand.tolerance != "free": + raise _not_yet( + "numeric tolerances on pupil aberrations (the sharpness score has no " + "calibration penalty; use tolerance: free until the evidence score lands)" + ) + indices.update(noll_indices_up_to_order(shorthand.order)) + for name, value in aberrations.terms.items(): + if isinstance(value, FreeParameter): + if value.tolerance != "free": + raise _not_yet(f"numeric tolerance on pupil aberration {name!r}") + if value.nominal != 0.0: + raise _not_yet(f"nonzero nominal on free pupil aberration {name!r} (the recovery starts at 0)") + indices.add(name_to_noll(name)) + elif float(value) != 0.0: + raise _not_yet(f"fixed nonzero pupil aberration {name!r}") + return sorted(indices) + + +def build_recovery_settings(cfg: V4Config, yx_shape: tuple[int, int]) -> tuple[RecoverySettings, list[int]]: + """Translate a 4.0 configuration into :class:`RecoverySettings`. + + Parameters + ---------- + cfg : V4Config + Configuration with at least one free pupil-aberration term. + yx_shape : tuple[int, int] + ``(Y, X)`` shape of the data, used for the shift-invariant + single-tile grid. + + Returns + ------- + settings : RecoverySettings + Settings for :func:`waveorder.api.shift_variant_recovery.recover_zernikes`. + noll_indices : list[int] + The estimated Noll indices, in ``parameters`` column order. + """ + system = cfg.imaging_system + if system.contrast != "fluorescence": + raise _not_yet("pupil-aberration calibration for phase contrast") + if cfg.object.dimension != 3: + raise _not_yet("pupil-aberration calibration for 2D (thin) objects") + + # Other free imaging-system numbers run through the 3.x optimizer via + # `wo reconstruct`; mixing them with pupil aberrations in one calibration + # is a later step. + non_zernike = { + name: p for name, p in free_parameters(system).items() if not name.startswith("detection.pupil_aberrations") + } + if non_zernike: + raise _not_yet(f"estimating {sorted(non_zernike)} jointly with pupil aberrations") + + calibration = cfg.solver.calibration + score = calibration.score if calibration is not None else "sharpness" + if score != "sharpness": + raise _not_yet(f"score: {score} on the sharpness path (evidence and gcv run through the evidence calibrator)") + + max_iter = 250 + if calibration is not None: + if len(calibration.search) != 1 or not isinstance(calibration.search[0], SearchGradient): + raise _not_yet("search stages other than a single gradient stage") + max_iter = calibration.search[0].gradient.max_iter + if calibration.final_estimate != "best_H": + raise _not_yet("final_estimate: average_near_best_H") + + loss = "midband" + midband_fractions = (0.125, 0.25) + optimizer = "adam" + advanced = calibration.advanced if calibration is not None else None + if advanced is not None: + if advanced.loss is not None: + loss_type = advanced.loss.get("type", "midband_power") + loss = _LOSS_NAME_MAP.get(loss_type) + if loss is None: + raise _not_yet(f"sharpness loss {loss_type!r}") + midband_fractions = tuple(advanced.loss.get("midband_fractions", midband_fractions)) + if advanced.optimizer is not None: + if advanced.optimizer in ("nelder_mead", "grid_search"): + raise _not_yet(f"optimizer {advanced.optimizer!r} for pupil aberrations") + optimizer = advanced.optimizer + + if cfg.object.prior != "tikhonov" or cfg.reconstruction_kind != "closed_form": + raise _not_yet("pupil-aberration calibration scores a Tikhonov closed-form reconstruction") + weight = cfg.object.weight + if not isinstance(weight, (int, float)): + raise _not_yet(f"object.weight: {weight} during calibration (use a number)") + + if system.field_dependence == "shift_invariant": + tile_size_yx = {"y": int(yx_shape[0]), "x": int(yx_shape[1])} + tile_overlap_yx = {"y": 0.0, "x": 0.0} + else: + tiled = system.field_dependence.tiled + tile_size_yx = {"y": int(tiled.size_yx[0]), "x": int(tiled.size_yx[1])} + tile_overlap_yx = {"y": float(tiled.overlap_yx[0]), "x": float(tiled.overlap_yx[1])} + + noll_indices = _resolve_noll_indices(cfg) + if not noll_indices: + raise ValueError("nothing to estimate: no pupil-aberration term has a tolerance") + + settings = RecoverySettings( + noll_indices=tuple(noll_indices), + tile_size_yx=tile_size_yx, + tile_overlap_yx=tile_overlap_yx, + loss=loss, + midband_fractions=midband_fractions, + optimizer=optimizer, + n_iter=max_iter, + wiener_regularization=float(weight), + ) + return settings, noll_indices + + +def run_zernike_calibration(cfg: V4Config, zyx_data: np.ndarray | torch.Tensor) -> CalibrationResult: + """Estimate pupil-aberration parameters from one volume. + + Parameters + ---------- + cfg : V4Config + Configuration with free pupil-aberration terms + (:func:`zernike_calibration_requested` is True). + zyx_data : ndarray or Tensor + 3D fluorescence stack ``(Z, Y, X)``. + + Returns + ------- + CalibrationResult + Per-tile Zernike coefficients (waves), tile grid, and the stitched + Wiener reconstruction the sharpness score was evaluated on. + """ + zyx_shape = tuple(zyx_data.shape) + settings, noll_indices = build_recovery_settings(cfg, zyx_shape[1:]) + + system = cfg.imaging_system + detection = system.detection + if detection.wavelength is None: + raise ValueError("imaging_system.detection.wavelength is required for fluorescence calibration") + + yx = cfg.data.pixel_size.yx + yx_pixel_size = float(yx) if isinstance(yx, (int, float)) else float(yx.y) + + result = recover_zernikes( + zyx_data, + yx_pixel_size_um=yx_pixel_size, + z_pixel_size_um=cfg.data.pixel_size.z, + wavelength_emission_um=detection.wavelength, + numerical_aperture_detection=nominal_of(detection.numerical_aperture), + index_of_refraction_media=system.index_of_refraction_media, + settings=settings, + ) + + tile_grid = { + "size_yx": [settings.tile_size_yx["y"], settings.tile_size_yx["x"]], + "overlap_yx": [settings.tile_overlap_yx["y"], settings.tile_overlap_yx["x"]], + "centers_yx": [[float(y), float(x)] for (y, x) in result.tile_centers_um], + } + return CalibrationResult( + parameters=result.coefs.numpy().astype(np.float64), + parameter_names=[noll_to_name(j) for j in noll_indices], + tile_grid=tile_grid, + reconstruction_zyx=result.reconstruction_zyx, + loss_history=result.loss_history, + noll_indices=noll_indices, + recovery_settings=result.settings, + ) diff --git a/waveorder/calibrate/scores.py b/waveorder/calibrate/scores.py new file mode 100644 index 00000000..825730d6 --- /dev/null +++ b/waveorder/calibrate/scores.py @@ -0,0 +1,132 @@ +"""Calibration scores: Whittle/Gaussian evidence and the calibration penalty. + +The evidence score is the exact negative log marginal likelihood of the +imaging system H for the Gaussian model (design document section 6, steps +3-5):: + + g = H f + n, f ~ N(0, S), n ~ N(0, sigma^2 I), + +which is diagonal per spatial frequency k with variance +``v_k = |H_k|^2 S_k + sigma^2``:: + + -log p(g | H) = sum_k [ |G_k|^2 / (2 v_k) + (1/2) log(2 pi v_k) ], + +the Whittle likelihood, with ``G = FFT(g) / sqrt(N)`` (unitary +normalization). The Tikhonov prior with weight ``lambda`` is +``S_k = sigma^2 / lambda`` (the prior implied by the Wiener filter +``H* / (|H|^2 + lambda)``). The first term rewards systems that explain the +observed power; the log-det "volume" term stops the score from always +preferring the largest transfer function. + +All terms are differentiable by torch autograd through ``otf`` (and +through anything ``otf`` was built from, e.g. Zernike coefficients via +:mod:`waveorder.calibrate.synthesis`). +""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor + + +def whittle_evidence( + data: Tensor | None = None, + otf: Tensor | None = None, + noise_sigma: float | Tensor = 1.0, + regularization_strength: float | Tensor = 1e-3, + object_spectrum: Tensor | None = None, + data_fft: Tensor | None = None, + exclude_dc: bool = False, +) -> Tensor: + """Negative log evidence ``-log p(g | H)`` under the Gaussian model. + + Lower is better. Compare values only across calls with identical data + and identical ``exclude_dc``. + + Parameters + ---------- + data : Tensor, optional + Real data ``g``, shape ``(Z, Y, X)`` (any number of dims works). + Either ``data`` or ``data_fft`` must be given. + otf : Tensor + Transfer function ``H_k``, same shape as the data, FFT frequency + order (e.g. from :func:`waveorder.calibrate.synthesis.fluorescence_otf`). + noise_sigma : float or Tensor + Noise standard deviation ``sigma`` in data units. + regularization_strength : float or Tensor + Tikhonov weight ``lambda``; the object prior spectrum is + ``S_k = sigma^2 / lambda``. Ignored when ``object_spectrum`` is given. + object_spectrum : Tensor, optional + Explicit object power spectrum ``S_k`` (same shape as ``otf``). When + given it overrides the flat Tikhonov ``S_k = sigma^2 / lambda``. + data_fft : Tensor, optional + Precomputed ``torch.fft.fftn(data, norm="ortho")``; pass this when + scoring many candidate systems against the same data. + exclude_dc : bool + Drop the ``k = 0`` term. Useful on real data with a large constant + background that the object prior does not model. + + Returns + ------- + Tensor + Scalar negative log evidence (sums over all frequencies). + """ + if data_fft is None: + if data is None: + raise ValueError("provide either 'data' or 'data_fft'") + data_fft = torch.fft.fftn(data, norm="ortho") + if otf is None: + raise ValueError("'otf' is required") + + sigma2 = torch.as_tensor(noise_sigma, dtype=torch.float64) ** 2 + power = torch.abs(data_fft) ** 2 + h2 = torch.abs(otf) ** 2 + if object_spectrum is not None: + s_k = object_spectrum + else: + s_k = sigma2 / torch.as_tensor(regularization_strength, dtype=torch.float64) + v_k = h2 * s_k + sigma2 + + terms = power / (2.0 * v_k) + 0.5 * torch.log(2.0 * math.pi * v_k) + if exclude_dc: + flat = terms.reshape(-1) + return flat[1:].sum() + return terms.sum() + + +def calibration_penalty( + values: dict[str, float | Tensor], + nominals: dict[str, float], + tolerances: dict[str, float | str], +) -> Tensor: + """The calibration term ``C = sum_i (h_i - h_i0)^2 / (2 tau_i^2)``. + + Always added to the chosen score; the tolerances are its weights + (design document section 4.4, "Score"). Parameters with tolerance + ``"free"`` contribute nothing. + + Parameters + ---------- + values : dict[str, float or Tensor] + Current parameter values by name. + nominals : dict[str, float] + Prior centers ``h_i0`` by name. + tolerances : dict[str, float or "free"] + Gaussian prior standard deviations ``tau_i`` by name, or ``"free"``. + + Returns + ------- + Tensor + Scalar penalty. + """ + total = torch.zeros((), dtype=torch.float64) + for name, value in values.items(): + tau = tolerances.get(name, "free") + if tau == "free": + continue + nominal = nominals[name] + residual = torch.as_tensor(value, dtype=torch.float64) - nominal + total = total + residual**2 / (2.0 * float(tau) ** 2) + return total diff --git a/waveorder/calibrate/synthesis.py b/waveorder/calibrate/synthesis.py new file mode 100644 index 00000000..868d9003 --- /dev/null +++ b/waveorder/calibrate/synthesis.py @@ -0,0 +1,221 @@ +"""Differentiable fluorescence PSF/OTF synthesis from imaging-system parameters. + +Builds the widefield fluorescence point spread function from a pupil with +named Zernike aberrations:: + + P(nu) = P0(nu) * exp(i 2 pi sum_j c_j Z_j(nu)), + h(z, y, x) = | IFFT2[ P(nu) exp(i 2 pi z n_z(nu)) ] |^2, + +with the Zernike coefficients ``c_j`` in **waves** (orthonormal Noll +normalization; see :mod:`waveorder.calibrate.zernike_names`). If your +coefficients are in radians of pupil phase (the cross-track exchange +currency of the 2026-10-30 sprint), divide by ``2 * pi`` before calling +in here. This mirrors +:func:`waveorder.models.shift_variant_fluorescent_3d.calculate_psf_at_field_position` +but takes the coefficients directly and preserves torch autograd through the +coefficients, the detection NA (via the sigmoid pupil edge) and the focus +offset, so the PSF can sit inside a gradient-based calibration loop. + +Notes +----- +The transfer function of ``wo compute-tf`` +(:func:`waveorder.models.isotropic_fluorescent_thick_3d.calculate_transfer_function`) +supersamples to Nyquist before cropping back; this module evaluates the +pupil on the native grid (like the shift-variant model), so the two agree +to a small discretization difference when the data is undersampled. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping + +import torch +from torch import Tensor + +from waveorder import optics +from waveorder.calibrate.zernike_names import name_to_noll +from waveorder.zernike import make_pupil_grid, zernike_noll + + +def _as_noll_coefficients(zernike_coefficients: Mapping[str | int, float | Tensor] | None) -> dict[int, float | Tensor]: + """Normalize a name- or Noll-keyed coefficient mapping to Noll keys.""" + if zernike_coefficients is None: + return {} + out: dict[int, float | Tensor] = {} + for key, value in zernike_coefficients.items(): + j = key if isinstance(key, int) else name_to_noll(key) + if j in out: + raise ValueError(f"Zernike term {key!r} (Noll {j}) given more than once") + out[j] = value + return out + + +def aberrated_pupil( + yx_shape: tuple[int, int], + yx_pixel_size: float, + wavelength_emission: float, + numerical_aperture_detection: float | Tensor, + zernike_coefficients: Mapping[str | int, float | Tensor] | None = None, +) -> Tensor: + """Complex detection pupil with named Zernike aberrations. + + Parameters + ---------- + yx_shape : tuple[int, int] + ``(Y, X)`` grid shape. + yx_pixel_size : float + Transverse pixel size in um. + wavelength_emission : float + Emission wavelength in um (free space). + numerical_aperture_detection : float or Tensor + Detection NA. A tensor keeps autograd through the (sigmoid) pupil + edge. + zernike_coefficients : mapping, optional + ``{name_or_noll_index: coefficient}`` with coefficients in waves; + float or 0-dim tensor values (tensors keep autograd). + + Returns + ------- + Tensor + Complex pupil, shape ``(Y, X)``, FFT frequency order (DC at + ``[0, 0]``). + """ + _, _, frr, theta = make_pupil_grid(yx_shape, yx_pixel_size, fft_order=True) + pupil_amplitude = optics.generate_pupil(frr, numerical_aperture_detection, wavelength_emission) + + coefficients = _as_noll_coefficients(zernike_coefficients) + if not coefficients: + return pupil_amplitude.to(torch.complex64) + + na = ( + numerical_aperture_detection + if isinstance(numerical_aperture_detection, Tensor) + else torch.tensor(float(numerical_aperture_detection)) + ) + rho = (frr * wavelength_emission / na).clamp(max=1.0) + phase_waves = torch.zeros_like(frr) + for j, c in coefficients.items(): + phase_waves = phase_waves + c * zernike_noll(j, rho, theta) + return pupil_amplitude * torch.exp(2j * math.pi * phase_waves) + + +def fluorescence_psf( + zyx_shape: tuple[int, int, int], + yx_pixel_size: float, + z_pixel_size: float, + wavelength_emission: float, + index_of_refraction_media: float, + numerical_aperture_detection: float | Tensor, + zernike_coefficients: Mapping[str | int, float | Tensor] | None = None, + z_focus_offset: float | Tensor = 0.0, +) -> Tensor: + """Widefield fluorescence intensity PSF from imaging-system parameters. + + Parameters + ---------- + zyx_shape : tuple[int, int, int] + ``(Z, Y, X)`` output shape. + yx_pixel_size : float + Transverse pixel size in um. + z_pixel_size : float + Axial pixel size in um. + wavelength_emission : float + Emission wavelength in um (free space). + index_of_refraction_media : float + Refractive index of the imaging medium. + numerical_aperture_detection : float or Tensor + Detection NA. + zernike_coefficients : mapping, optional + ``{name_or_noll_index: coefficient}``, coefficients in waves + (see :mod:`waveorder.calibrate.zernike_names`). + z_focus_offset : float or Tensor + Focus offset in slices; shifts the defocus axis. + + Returns + ------- + Tensor + Real PSF, shape ``(Z, Y, X)``, sum-normalized to 1, in FFT order on + every axis (matching the shift-variant model): the in-focus slice + is at index 0 along z and the peak sits at ``[0, 0, 0]``, so + ``fftn(psf)`` is a phase-ramp-free convolution kernel. Apply + ``torch.fft.fftshift`` for a center-origin view. + """ + z_total, _, _ = zyx_shape + yx_shape = zyx_shape[1:] + + pupil = aberrated_pupil( + yx_shape, + yx_pixel_size, + wavelength_emission, + numerical_aperture_detection, + zernike_coefficients, + ) + _, _, frr, _ = make_pupil_grid(yx_shape, yx_pixel_size, fft_order=True) + pupil_support = optics.generate_pupil(frr, numerical_aperture_detection, wavelength_emission) + + z_steps = torch.arange(z_total) - z_total // 2 + if isinstance(z_focus_offset, Tensor) or z_focus_offset != 0.0: + z_steps = z_steps - z_focus_offset + z_position_list = torch.fft.ifftshift(z_steps * z_pixel_size) + + propagation_kernel = optics.generate_propagation_kernel( + frr, + pupil_support, + wavelength_emission / index_of_refraction_media, + z_position_list, + ) + coherent_psf = torch.fft.ifft2(propagation_kernel * pupil[None, :, :], dim=(1, 2)) + psf = torch.abs(coherent_psf) ** 2 + + total = psf.sum() + return psf / total.clamp(min=1e-30) + + +def fluorescence_otf( + zyx_shape: tuple[int, int, int], + yx_pixel_size: float, + z_pixel_size: float, + wavelength_emission: float, + index_of_refraction_media: float, + numerical_aperture_detection: float | Tensor, + zernike_coefficients: Mapping[str | int, float | Tensor] | None = None, + z_focus_offset: float | Tensor = 0.0, + normalize: str = "peak", +) -> Tensor: + """Widefield fluorescence OTF from imaging-system parameters. + + Same parameters as :func:`fluorescence_psf`, plus: + + Parameters + ---------- + normalize : {"peak", "none"} + ``"peak"`` scales to ``max |OTF| = 1`` (the WaveOrder transfer + function convention, matching + ``isotropic_fluorescent_thick_3d.calculate_transfer_function``); + ``"none"`` returns the FFT of the sum-normalized PSF, whose DC + value is exactly 1 up to the pupil discretization. + + Returns + ------- + Tensor + Complex OTF, shape ``(Z, Y, X)``, FFT frequency order on every + axis — directly usable in ``g_fft = f_fft * otf`` forward models + and in :func:`waveorder.calibrate.scores.whittle_evidence`. + """ + psf = fluorescence_psf( + zyx_shape, + yx_pixel_size, + z_pixel_size, + wavelength_emission, + index_of_refraction_media, + numerical_aperture_detection, + zernike_coefficients, + z_focus_offset, + ) + otf = torch.fft.fftn(psf, dim=(0, 1, 2)) + if normalize == "peak": + otf = otf / torch.clamp(torch.max(torch.abs(otf)), min=1e-12) + elif normalize != "none": + raise ValueError(f"normalize must be 'peak' or 'none', got {normalize!r}") + return otf diff --git a/waveorder/calibrate/tfzarr.py b/waveorder/calibrate/tfzarr.py new file mode 100644 index 00000000..125b816c --- /dev/null +++ b/waveorder/calibrate/tfzarr.py @@ -0,0 +1,149 @@ +"""The ``tf.zarr`` transfer-function store (design document section 3). + +Layout:: + + tf.zarr/ + parameters [N_tiles, P] always + parameters_covariance [N_tiles, P, P] zeros when the score has no curvature + transfer_function [Z, Y, X] or [N_tiles, Z, Y, X], one tile per chunk + attrs: + parameter_names [P] + tile_grid {size_yx, overlap_yx, centers_yx} + provenance {kind, sources, config_hash, ...} + +``parameters`` rows are per tile; a single row is broadcast across tiles by +consumers. ``transfer_function`` is optional for tiled stores, where it is +regenerated per tile from ``parameters`` at apply time. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import zarr + + +@dataclass +class TransferFunctionStore: + """In-memory view of a ``tf.zarr``.""" + + parameters: np.ndarray # [N_tiles, P] + parameter_names: list[str] + tile_grid: dict[str, Any] + provenance: dict[str, Any] = field(default_factory=dict) + parameters_covariance: np.ndarray | None = None # [N_tiles, P, P] + transfer_function: np.ndarray | None = None # [Z, Y, X] or [N_tiles, Z, Y, X] + + @property + def n_tiles(self) -> int: + return self.parameters.shape[0] + + +def write_tf_zarr( + path: Path | str, + parameters: np.ndarray, + parameter_names: list[str], + tile_grid: dict[str, Any], + provenance: dict[str, Any] | None = None, + parameters_covariance: np.ndarray | None = None, + transfer_function: np.ndarray | None = None, + extra_arrays: dict[str, np.ndarray] | None = None, +) -> None: + """Write a ``tf.zarr`` store. + + Parameters + ---------- + path : Path or str + Output store path. + parameters : ndarray + ``[N_tiles, P]`` estimated (or nominal) parameter values. + parameter_names : list[str] + ``P`` names, e.g. ``["defocus", "astigmatism_0", ...]``. + tile_grid : dict + ``{"size_yx": [y, x], "overlap_yx": [y, x], "centers_yx": [[y_um, x_um], ...]}``; + one center when shift-invariant. + provenance : dict, optional + ``{"kind": "nominal" | "beads" | "refined" | "simulated", "sources": [...], + "config_hash": ...}`` plus free-form entries. + parameters_covariance : ndarray, optional + ``[N_tiles, P, P]``; written as zeros when None. + transfer_function : ndarray, optional + ``[Z, Y, X]`` (shift-invariant) or ``[N_tiles, Z, Y, X]`` (tiled, + chunked one tile per chunk). + extra_arrays : dict, optional + Additional named arrays, e.g. ``score_trace`` (per gradient + iteration) and ``score_grid`` (one row per evaluated point: + parameter coordinates followed by the score). + """ + parameters = np.asarray(parameters, dtype=np.float64) + if parameters.ndim != 2: + raise ValueError(f"parameters must be [N_tiles, P], got shape {parameters.shape}") + n_tiles, n_params = parameters.shape + if len(parameter_names) != n_params: + raise ValueError(f"{len(parameter_names)} parameter_names for {n_params} parameters") + + if parameters_covariance is None: + parameters_covariance = np.zeros((n_tiles, n_params, n_params), dtype=np.float64) + parameters_covariance = np.asarray(parameters_covariance, dtype=np.float64) + if parameters_covariance.shape != (n_tiles, n_params, n_params): + raise ValueError( + f"parameters_covariance must be [N_tiles, P, P] = {(n_tiles, n_params, n_params)}, " + f"got {parameters_covariance.shape}" + ) + + group = zarr.open_group(str(path), mode="w") + group.create_array("parameters", shape=parameters.shape, dtype="float64")[:] = parameters + group.create_array("parameters_covariance", shape=parameters_covariance.shape, dtype="float64")[:] = ( + parameters_covariance + ) + + if transfer_function is not None: + transfer_function = np.asarray(transfer_function) + if transfer_function.ndim == 4: + chunks = (1, *transfer_function.shape[1:]) + else: + chunks = transfer_function.shape + tf_array = group.create_array( + "transfer_function", + shape=transfer_function.shape, + dtype=str(transfer_function.dtype), + chunks=chunks, + ) + tf_array[:] = transfer_function + + for name, array in (extra_arrays or {}).items(): + array = np.asarray(array) + group.create_array(name, shape=array.shape, dtype=str(array.dtype))[:] = array + + group.attrs["parameter_names"] = list(parameter_names) + group.attrs["tile_grid"] = tile_grid + group.attrs["provenance"] = provenance or {} + + +def read_tf_zarr(path: Path | str) -> TransferFunctionStore: + """Read a ``tf.zarr`` store written by :func:`write_tf_zarr`. + + Parameters + ---------- + path : Path or str + Store path. + + Returns + ------- + TransferFunctionStore + """ + group = zarr.open_group(str(path), mode="r") + transfer_function = None + if "transfer_function" in group: + transfer_function = group["transfer_function"][:] + return TransferFunctionStore( + parameters=group["parameters"][:], + parameters_covariance=group["parameters_covariance"][:], + parameter_names=list(group.attrs["parameter_names"]), + tile_grid=dict(group.attrs["tile_grid"]), + provenance=dict(group.attrs.get("provenance", {})), + transfer_function=transfer_function, + ) diff --git a/waveorder/calibrate/zernike_names.py b/waveorder/calibrate/zernike_names.py new file mode 100644 index 00000000..de8795cb --- /dev/null +++ b/waveorder/calibrate/zernike_names.py @@ -0,0 +1,123 @@ +"""Named Zernike terms and their Noll indices. + +Convention (fixed for WaveOrder 4.0, see the design document section 4.4): + +* Indexing is Noll (``j = 1`` is piston), evaluated by + :func:`waveorder.zernike.zernike_noll` with the orthonormal + normalization, so a coefficient is the RMS wavefront error of its term. +* Coefficients are in **waves** of pupil phase: the pupil is + ``P0 * exp(i 2 pi sum_j c_j Z_j)``. + + .. warning:: UNIT CONVERSION. Some collaborating pipelines (e.g. the + blind-deconvolution track) exchange Zernike coefficients in + **radians** of pupil phase: ``c_radians = 2 * pi * c_waves``. Every + function in :mod:`waveorder.calibrate` takes and returns **waves**; + convert explicitly at the boundary and never mix the two silently. +* The pupil azimuth ``theta = atan2(f_y, f_x)`` is zero along the +x + frequency axis, matching :func:`waveorder.zernike.make_pupil_grid`. +* Named terms: ``defocus`` (Noll 4), ``astigmatism_45`` (5, the + ``sin 2theta`` term), ``astigmatism_0`` (6, ``cos 2theta``), ``coma_y`` + (7, ``sin theta``), ``coma_x`` (8, ``cos theta``), ``trefoil_y`` (9, + ``sin 3theta``), ``trefoil_x`` (10, ``cos 3theta``), ``spherical`` (11). + Any other term is addressed as ``z`` by Noll index, e.g. ``z12``. +* ``zernike_up_to_order: {order: n}`` frees every Noll index through + radial order ``n`` **excluding piston, tip and tilt** (Noll 1-3): piston + does not change an intensity PSF, and tip/tilt are degenerate with a + lateral translation of the object. +""" + +from __future__ import annotations + +import re + +from waveorder.zernike import _noll_to_nm + +NAMED_TERMS: dict[str, int] = { + "defocus": 4, + "astigmatism_45": 5, + "astigmatism_0": 6, + "coma_y": 7, + "coma_x": 8, + "trefoil_y": 9, + "trefoil_x": 10, + "spherical": 11, +} + +_NOLL_TO_NAME: dict[int, str] = {j: name for name, j in NAMED_TERMS.items()} + +_ZN_PATTERN = re.compile(r"^z([0-9]+)$") + + +def name_to_noll(name: str) -> int: + """Convert a named Zernike term to its Noll index. + + Parameters + ---------- + name : str + A named term (``defocus``, ``astigmatism_0``, ...) or ``z`` by + Noll index. + + Returns + ------- + int + Noll index (1-based). + + Raises + ------ + ValueError + If the name is not a recognized term. + """ + if name in NAMED_TERMS: + return NAMED_TERMS[name] + match = _ZN_PATTERN.match(name) + if match: + j = int(match.group(1)) + if j < 1: + raise ValueError(f"Noll index must be >= 1, got {name!r}") + return j + raise ValueError(f"Unknown Zernike term {name!r}. Use one of {sorted(NAMED_TERMS)} or 'z' with a Noll index n.") + + +def noll_to_name(j: int) -> str: + """Canonical name of a Noll index: the named term if it has one, else ``z``.""" + return _NOLL_TO_NAME.get(j, f"z{j}") + + +def validate_zernike_name(name: str) -> str: + """Validate a Zernike term name, returning it unchanged. + + Raises + ------ + ValueError + If the name is neither a named term nor ``z``. + """ + name_to_noll(name) + return name + + +def noll_indices_up_to_order(order: int) -> list[int]: + """Noll indices of every term through radial order ``order``. + + Piston, tip and tilt (Noll 1-3) are excluded; see the module docstring. + + Parameters + ---------- + order : int + Highest radial order ``n`` to include, ``order >= 2``. + + Returns + ------- + list[int] + Sorted Noll indices with ``4 <= j`` and radial order ``<= order``. + """ + if order < 2: + raise ValueError(f"order must be >= 2, got {order}") + indices = [] + j = 4 + while True: + n, _ = _noll_to_nm(j) + if n > order: + break + indices.append(j) + j += 1 + return indices diff --git a/waveorder/cli/apply_inverse_transfer_function.py b/waveorder/cli/apply_inverse_transfer_function.py index 41e023fe..15f7688a 100644 --- a/waveorder/cli/apply_inverse_transfer_function.py +++ b/waveorder/cli/apply_inverse_transfer_function.py @@ -159,8 +159,7 @@ def get_reconstruction_output_metadata( import numpy as np from iohub import open_ome_zarr - from waveorder.cli.settings import ReconstructionSettings - from waveorder.io import utils + from waveorder.config.loader import load_reconstruction_settings # Get non-OME-Zarr plate-level metadata if it's available plate_metadata = {} @@ -181,7 +180,7 @@ def get_reconstruction_output_metadata( T, _, Z, Y, X = input_dataset.data.shape scale = input_dataset.scale - settings = utils.yaml_to_model(config_path, ReconstructionSettings) + settings = load_reconstruction_settings(config_path) channel_names = settings.output_channel_names output_z_shape = 1 if settings.output_z_is_singleton else Z @@ -250,12 +249,12 @@ def apply_inverse_transfer_function_single_position( phase, ) from waveorder.cli.printing import echo_headline, echo_settings - from waveorder.cli.settings import ReconstructionSettings from waveorder.cli.utils import ( apply_inverse_czyx, reconstruction_fingerprint, resolve_time_indices, ) + from waveorder.config.loader import load_reconstruction_settings from waveorder.io import utils if verbose: @@ -270,7 +269,7 @@ def apply_inverse_transfer_function_single_position( input_xa = input_dataset.to_xarray() # Load config file - settings = utils.yaml_to_model(config_filepath, ReconstructionSettings) + settings = load_reconstruction_settings(config_filepath) # Check input channel names if not set(settings.input_channel_names).issubset(input_dataset.channel_names): @@ -516,12 +515,11 @@ def apply_inverse_transfer_function_cli( from iohub import open_ome_zarr from iohub.ngff.utils import create_empty_plate - from waveorder.cli.settings import ReconstructionSettings from waveorder.cli.utils import ( generate_valid_position_key, is_single_position_store, ) - from waveorder.io import utils + from waveorder.config.loader import load_reconstruction_settings per_position_transfer_function_dirpaths = _resolve_transfer_function_dirpaths( transfer_function_dirpath, input_position_dirpaths @@ -552,7 +550,7 @@ def apply_inverse_transfer_function_cli( position_keys.append(position_key) # Every position gets the same provenance, recorded before any timepoint runs. - provenance = get_reconstruction_provenance(utils.yaml_to_model(config_filepath, ReconstructionSettings)) + provenance = get_reconstruction_provenance(load_reconstruction_settings(config_filepath)) # `create_empty_plate` allocates one shape and scale per call. for position_key, output_metadata in zip(position_keys, output_metadatas, strict=True): diff --git a/waveorder/cli/bench.py b/waveorder/cli/bench.py index 967358f8..9a082fdf 100644 --- a/waveorder/cli/bench.py +++ b/waveorder/cli/bench.py @@ -78,6 +78,18 @@ def run( ), ), ] = False, + migrate_configs: Annotated[ + bool, + typer.Option( + "--migrate-configs", + help=( + "Migrate each 3.x case 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." + ), + ), + ] = False, ): """Run benchmark cases. @@ -111,6 +123,7 @@ def run( n_skipped_hpc = len(exp.cases) - len(selected) metadata = collect_metadata() + metadata["migrate_configs"] = migrate_configs timestamp = datetime.now().strftime("%Y-%m-%dT%H-%M-%S") run_name = f"{timestamp}_{exp.name}" @@ -147,6 +160,8 @@ def run( reference=case.reference, simulation=case.simulation, recovery=case.recovery, + # Forward-compatible with runners that predate the flag. + **({"migrate_config": True} if migrate_configs else {}), ) elif case.type == "hpc": metrics = run_hpc_case( @@ -158,6 +173,8 @@ def run( save_all=save_all, reference=case.reference, crop=case.crop, + # Forward-compatible with runners that predate the flag. + **({"migrate_config": True} if migrate_configs else {}), ) results[case_name] = metrics diff --git a/waveorder/cli/calibrate_tf.py b/waveorder/cli/calibrate_tf.py new file mode 100644 index 00000000..400051cd --- /dev/null +++ b/waveorder/cli/calibrate_tf.py @@ -0,0 +1,263 @@ +from pathlib import Path +from typing import Annotated, Optional + +import typer + +from waveorder.cli.parsing import ( + ConfigFilepath, + InputPositionDirpaths, + OutputDirpath, +) + +PriorTransferFunctionDirpath = Annotated[ + Optional[Path], + typer.Option( + "--transfer-function-dirpath", + "-t", + help="tf.zarr whose parameters (and covariance, if non-zero) set the prior on the " + "estimated parameters, overriding the configuration's tolerances for the parameters " + "it covers.", + ), +] +InitTransferFunctionDirpath = Annotated[ + Optional[Path], + typer.Option( + "--init", + help="tf.zarr whose parameters set the starting point of the search without changing the prior.", + ), +] + + +def calibrate_transfer_function_cli( + input_position_dirpath: Path, + config_filepath: Path, + output_dirpath: Path, + prior_tf_dirpath: Path | None = None, + init_tf_dirpath: Path | None = None, +) -> None: + """Estimate the free imaging-system parameters from one position. + + Writes a ``tf.zarr`` with ``parameters``, ``parameters_covariance``, + ``parameter_names``, ``tile_grid`` and ``provenance`` (design document + section 3). + """ + # Deferred imports: these pull in torch, iohub, numpy, etc. + # Only loaded when the command runs, keeping wo calibrate-tf -h fast. + import hashlib + import warnings + + import numpy as np + import yaml + from iohub.ngff import open_ome_zarr + + from waveorder.calibrate.beads import run_bead_calibration + from waveorder.calibrate.evidence_search import run_evidence_calibration + from waveorder.calibrate.run import ( + run_zernike_calibration, + zernike_calibration_requested, + ) + from waveorder.calibrate.synthesis import fluorescence_otf + from waveorder.calibrate.tfzarr import read_tf_zarr, write_tf_zarr + from waveorder.calibrate.zernike_names import name_to_noll + from waveorder.cli.printing import echo_headline + from waveorder.cli.utils import resolve_time_indices + from waveorder.config.loader import load_v4_config + from waveorder.config.migrate import detect_version, migrate_3x_dict + from waveorder.config.schema import PriorBeads, free_parameters + + config_text = Path(config_filepath).read_text() + raw = yaml.safe_load(config_text) + if detect_version(raw) >= 4: + cfg = load_v4_config(config_filepath) + else: + warnings.warn( + f"{config_filepath} is a 3.x configuration; converting in memory. " + f"Convert it with: waveorder migrate-config {config_filepath}", + DeprecationWarning, + stacklevel=2, + ) + cfg = migrate_3x_dict(raw) + + if not free_parameters(cfg.imaging_system): + raise ValueError( + "calibrate-tf with no tolerances: nothing to estimate. Give at least one " + "imaging_system number as {nominal: x, tolerance: t}." + ) + if not zernike_calibration_requested(cfg): + raise NotImplementedError( + "calibrate-tf currently estimates pupil_aberrations terms; estimate " + "numerical_aperture / tilt / z_focus_offset through 'wo reconstruct' " + "with solver.calibration for now." + ) + + echo_headline(f"Calibrating transfer function from {input_position_dirpath}\n") + + input_dataset = open_ome_zarr(input_position_dirpath, layout="fov", mode="r") + try: + if not set(cfg.data.input_channel_names).issubset(input_dataset.channel_names): + raise ValueError( + f"input_channel_names = {cfg.data.input_channel_names} must appear in " + f"{input_position_dirpath} (channels: {input_dataset.channel_names})" + ) + t_index = resolve_time_indices(cfg.data.time_indices, input_dataset.data.shape[0])[0] + czyx = input_dataset.to_xarray().isel(t=t_index).sel(c=cfg.data.input_channel_names) + zyx_data = np.asarray(czyx.values[0], dtype=np.float32) + print( + f" calibrating on t={t_index} (first requested time index), channel {cfg.data.input_channel_names[0]!r}" + ) + finally: + input_dataset.close() + + calibration = cfg.solver.calibration + score = calibration.score if calibration is not None else "evidence" + extra_provenance = {} + score_arrays = {} + covariance = None + provenance_kind = "refined" + prior_store = read_tf_zarr(prior_tf_dirpath) if prior_tf_dirpath is not None else None + init_store = read_tf_zarr(init_tf_dirpath) if init_tf_dirpath is not None else None + if score not in ("evidence", "gcv") and (prior_store is not None or init_store is not None): + raise NotImplementedError("-t / --init are supported with score: evidence or gcv for now") + if isinstance(cfg.object.prior, PriorBeads): + if prior_store is not None or init_store is not None: + raise NotImplementedError("-t / --init with prior: beads is not supported yet") + bead_result = run_bead_calibration(cfg, zyx_data) + parameters = bead_result.parameters + parameter_names = bead_result.parameter_names + covariance = bead_result.parameters_covariance + tile_grid = { + "size_yx": [int(zyx_data.shape[1]), int(zyx_data.shape[2])], + "overlap_yx": [0, 0], + "centers_yx": [[0.0, 0.0]], + } + extra_provenance = { + "best_score": bead_result.best_score, + "noise_sigma": bead_result.noise_sigma, + "background": bead_result.background, + "n_beads": int(bead_result.positions_zyx.shape[0]), + } + provenance_kind = "beads" + if "score_trace" in (cfg.outputs if cfg.outputs != "all" else ["score_trace"]): + score_arrays["score_trace"] = np.asarray(bead_result.score_trace) + score_arrays["bead_positions_zyx"] = bead_result.positions_zyx + score_arrays["bead_amplitudes"] = bead_result.amplitudes + final_loss = bead_result.best_score + elif score in ("evidence", "gcv"): + evidence_result = run_evidence_calibration(cfg, zyx_data, prior_store=prior_store, init_store=init_store) + parameters = evidence_result.parameters + parameter_names = evidence_result.parameter_names + covariance = evidence_result.parameters_covariance + tile_grid = { + "size_yx": [int(zyx_data.shape[1]), int(zyx_data.shape[2])], + "overlap_yx": [0, 0], + "centers_yx": [[0.0, 0.0]], + } + extra_provenance = { + "best_score": evidence_result.best_score, + "noise_sigma": evidence_result.noise_sigma, + "weight": evidence_result.weight, + } + wants = cfg.outputs if cfg.outputs != "all" else ["score_trace", "score_grid"] + if "score_trace" in wants and evidence_result.score_trace: + score_arrays["score_trace"] = np.asarray(evidence_result.score_trace) + if "score_grid" in wants and evidence_result.score_grid is not None: + score_arrays["score_grid"] = evidence_result.score_grid + final_loss = evidence_result.best_score + elif score == "sharpness": + sharpness_result = run_zernike_calibration(cfg, zyx_data) + parameters = sharpness_result.parameters + parameter_names = sharpness_result.parameter_names + tile_grid = sharpness_result.tile_grid + final_loss = sharpness_result.loss_history[-1] + else: + raise NotImplementedError(f"score: {score} is not implemented in calibrate-tf") + + # Shift-invariant stores always carry the transfer function; tiled stores + # only on request, and are otherwise regenerated per tile at apply time. + transfer_function = None + n_tiles = parameters.shape[0] + wants_tf = cfg.outputs == "all" or "transfer_function" in cfg.outputs + if n_tiles == 1 and wants_tf: + detection = cfg.imaging_system.detection + yx = cfg.data.pixel_size.yx + coefficients = {name_to_noll(name): value for name, value in zip(parameter_names, parameters[0])} + transfer_function = ( + fluorescence_otf( + zyx_shape=zyx_data.shape, + yx_pixel_size=float(yx) if isinstance(yx, (int, float)) else float(yx.y), + z_pixel_size=cfg.data.pixel_size.z, + wavelength_emission=detection.wavelength, + index_of_refraction_media=cfg.imaging_system.index_of_refraction_media, + numerical_aperture_detection=( + detection.numerical_aperture.nominal + if hasattr(detection.numerical_aperture, "nominal") + else detection.numerical_aperture + ), + zernike_coefficients=coefficients, + ) + .numpy() + .astype(np.complex64) + ) + + provenance = { + "kind": provenance_kind, + "sources": [str(input_position_dirpath)], + "config_hash": hashlib.sha256(config_text.encode()).hexdigest(), + "score": score, + "time_index": int(t_index), + **extra_provenance, + } + write_tf_zarr( + output_dirpath, + parameters=parameters, + parameter_names=parameter_names, + tile_grid=tile_grid, + provenance=provenance, + parameters_covariance=covariance, + transfer_function=transfer_function, + extra_arrays=score_arrays, + ) + + echo_headline( + f"Estimated {len(parameter_names)} parameters on {n_tiles} tile(s); final {score} score {final_loss:.4g}" + ) + if n_tiles == 1: + for name, value in zip(parameter_names, parameters[0]): + print(f" {name:16s} {value:+.4f} waves") + echo_headline(f"Wrote {output_dirpath}") + + +def _calibrate_tf_cli( + input_position_dirpaths: InputPositionDirpaths, + config_filepath: ConfigFilepath, + output_dirpath: OutputDirpath, + prior_tf_dirpath: PriorTransferFunctionDirpath = None, + init_tf_dirpath: InitTransferFunctionDirpath = None, +) -> None: + """Estimate free imaging-system parameters (tolerances) from data. + + Every imaging_system number written as {nominal: x, tolerance: t} is + estimated from the first position in `input-position-dirpaths`, and the + result is written as a tf.zarr with `parameters`, `parameters_covariance`, + `parameter_names`, `tile_grid`, and `provenance`. + + A tf.zarr passed with `-t` sets the prior: its parameters move the prior + center, and its covariance (when non-zero) replaces the configuration's + tolerances for the parameters it covers. `--init` sets the starting point + only. + + \b + Example: + \033[92mwo calibrate-tf -i ./input.zarr/0/0/0 -c ./config.yml -o ./tf_hat.zarr\033[0m + """ + if len(input_position_dirpaths) > 1: + raise NotImplementedError( + "calibrate-tf with several positions (one shared H) is not implemented yet; pass one position." + ) + calibrate_transfer_function_cli( + input_position_dirpaths[0], + config_filepath, + output_dirpath, + prior_tf_dirpath=prior_tf_dirpath, + init_tf_dirpath=init_tf_dirpath, + ) diff --git a/waveorder/cli/compute_transfer_function.py b/waveorder/cli/compute_transfer_function.py index e365a09f..1e974354 100644 --- a/waveorder/cli/compute_transfer_function.py +++ b/waveorder/cli/compute_transfer_function.py @@ -104,11 +104,10 @@ def compute_transfer_function_cli( phase, ) from waveorder.cli.printing import echo_headline, echo_settings - from waveorder.cli.settings import ReconstructionSettings - from waveorder.io import utils + from waveorder.config.loader import load_reconstruction_settings # Load config file - settings = utils.yaml_to_model(config_filepath, ReconstructionSettings) + settings = load_reconstruction_settings(config_filepath) echo_headline(f"Generating transfer functions and storing in {output_dirpath}\n") diff --git a/waveorder/cli/main.py b/waveorder/cli/main.py index 64c2e947..66431c2c 100644 --- a/waveorder/cli/main.py +++ b/waveorder/cli/main.py @@ -19,8 +19,10 @@ from waveorder.cli.apply_inverse_transfer_function import _apply_inverse_transfer_function_cli from waveorder.cli.bench import benchmark +from waveorder.cli.calibrate_tf import _calibrate_tf_cli from waveorder.cli.compute_transfer_function import _compute_transfer_function_cli from waveorder.cli.gui_widget import gui +from waveorder.cli.migrate_config import _migrate_config_cli from waveorder.cli.reconstruct import _reconstruct_cli from waveorder.cli.simulate import _simulate_cli from waveorder.cli.tile_stitch import _tile_stitch_cli @@ -92,8 +94,10 @@ def _main( "reconstruct": (_reconstruct_cli, "rec"), "compute-transfer-function": (_compute_transfer_function_cli, "compute-tf"), "apply-inverse-transfer-function": (_apply_inverse_transfer_function_cli, "apply-inv-tf"), + "calibrate-transfer-function": (_calibrate_tf_cli, "calibrate-tf"), "tile-stitch": (_tile_stitch_cli, "ts"), "simulate": (_simulate_cli, "sim"), + "migrate-config": (_migrate_config_cli, "mc"), "view": (_view_cli, "v"), "interactive": (gui, "gui"), } @@ -104,6 +108,7 @@ def _main( "reconstruct", "compute-transfer-function", "apply-inverse-transfer-function", + "calibrate-transfer-function", "tile-stitch", }: command_options["cls"] = _InputPathsCommand diff --git a/waveorder/cli/migrate_config.py b/waveorder/cli/migrate_config.py new file mode 100644 index 00000000..d4116325 --- /dev/null +++ b/waveorder/cli/migrate_config.py @@ -0,0 +1,66 @@ +from pathlib import Path +from typing import Annotated, Optional + +import typer + + +def _migrate_config_cli( + config_filepath: Annotated[ + Path, + typer.Argument( + exists=True, + file_okay=True, + dir_okay=False, + help="Path to a 3.x YAML configuration file.", + ), + ], + output_filepath: Annotated[ + Optional[Path], + typer.Option( + "--output-filepath", + "-o", + help="Path for the 4.0 configuration. Defaults to '_v4.yml' beside the input.", + ), + ] = None, +): + """ + Migrate a 3.x configuration file to the 4.0 format. + + The migration is deterministic and lossless: the 4.0 configuration + selects the same code path with the same numerical parameters and + produces the same reconstruction. + + \b + Example: + \033[92mwaveorder migrate-config ./config.yml -o ./config_v4.yml\033[0m + """ + # Deferred imports: these pull in torch via the settings models. + # Only loaded when the command runs, keeping the help fast. + import yaml + + from waveorder.config.migrate import ( + MigrationError, + detect_version, + migrate_3x_dict, + v4_to_yaml_text, + ) + + with open(config_filepath) as f: + raw = yaml.safe_load(f) + + if detect_version(raw) >= 4: + typer.echo(typer.style(f"{config_filepath} is already a 4.0 configuration.", fg=typer.colors.YELLOW)) + raise typer.Exit() + + try: + migrated = migrate_3x_dict(raw) + except MigrationError as e: + raise typer.BadParameter(str(e)) from e + + if output_filepath is None: + output_filepath = config_filepath.parent / f"{config_filepath.stem}_v4.yml" + + text = v4_to_yaml_text(migrated) + output_filepath.write_text(text) + typer.echo(text) + typer.echo(typer.style(f"Wrote 4.0 configuration to {output_filepath}", fg=typer.colors.GREEN)) diff --git a/waveorder/cli/reconstruct.py b/waveorder/cli/reconstruct.py index 3b49972e..9dbfa049 100644 --- a/waveorder/cli/reconstruct.py +++ b/waveorder/cli/reconstruct.py @@ -51,12 +51,11 @@ def _reconstruct_cli( # Only loaded when the command runs, keeping wo rec -h fast. from waveorder.cli.apply_inverse_transfer_function import apply_inverse_transfer_function_cli from waveorder.cli.compute_transfer_function import compute_transfer_function_cli - from waveorder.cli.settings import ReconstructionSettings from waveorder.cli.utils import read_zyx_shapes - from waveorder.io import utils + from waveorder.config.loader import load_reconstruction_settings from waveorder.optim import has_optimizable_params - settings = utils.yaml_to_model(config_filepath, ReconstructionSettings) + settings = load_reconstruction_settings(config_filepath) # Check for optimizable parameters and run optimization if needed if has_optimizable_params(settings): diff --git a/waveorder/cli/simulate.py b/waveorder/cli/simulate.py index dffcc25c..d19b6c5a 100644 --- a/waveorder/cli/simulate.py +++ b/waveorder/cli/simulate.py @@ -8,9 +8,14 @@ def _simulate_cli(config_filepath: ConfigFilepath, output_dirpath: OutputDirpath): """Simulate phantom data matching a configuration file. - Writes a single zarr with both the ground-truth phantom and + For a 3.x configuration (or a 4.0 configuration without a `simulate` + block), writes a single zarr with both the ground-truth phantom and simulated measurement as separate channels. + A 4.0 configuration with a `simulate: {seed, shape}` block describes the + true imaging system (design section 5) and writes three stores into the + output directory: `data.zarr`, `object.zarr`, and `tf.zarr`. + \b Example: \033[92mwo sim -c ./phase.yml -o ./phase.zarr\033[0m @@ -21,6 +26,7 @@ def _simulate_cli(config_filepath: ConfigFilepath, output_dirpath: OutputDirpath # Only loaded when the command runs, keeping wo sim -h fast. import numpy as np import xarray as xr + import yaml from iohub.ngff import open_ome_zarr from iohub.ngff.models import TransformationMeta @@ -31,8 +37,19 @@ def _simulate_cli(config_filepath: ConfigFilepath, output_dirpath: OutputDirpath phase, ) from waveorder.cli.printing import echo_headline - from waveorder.cli.settings import ReconstructionSettings - from waveorder.io import utils + from waveorder.config.loader import load_reconstruction_settings, load_v4_config + from waveorder.config.migrate import detect_version + from waveorder.config.simulate import run_simulation + + with open(config_filepath) as f: + raw = yaml.safe_load(f) + if detect_version(raw) >= 4 and raw.get("simulate") is not None: + cfg = load_v4_config(config_filepath) + truth = run_simulation(cfg, output_dirpath) + echo_headline(f"Wrote data.zarr, object.zarr, tf.zarr to {output_dirpath}") + for name, value in truth.items(): + print(f" true {name:16s} {value:+.4f} waves") + return def _write_czyx(czyx: xr.DataArray, path: Path): """Write a CZYX xr.DataArray to an HCS OME-Zarr.""" @@ -59,7 +76,7 @@ def _write_czyx(czyx: xr.DataArray, path: Path): position["0"][0] = czyx.values dataset.close() - settings = utils.yaml_to_model(config_filepath, ReconstructionSettings) + settings = load_reconstruction_settings(config_filepath) recon_dim = settings.reconstruction_dimension diff --git a/waveorder/config/__init__.py b/waveorder/config/__init__.py new file mode 100644 index 00000000..e8f4046b --- /dev/null +++ b/waveorder/config/__init__.py @@ -0,0 +1,63 @@ +"""WaveOrder 4.0 configuration: schema, migration, and loading. + +The 4.0 configuration regroups a reconstruction into five blocks (``data``, +``imaging_system``, ``object``, ``noise``, ``solver``) plus ``outputs``, with +a required ``version`` field. A configuration without ``version`` is 3.x and +is migrated mechanically; see :mod:`waveorder.config.migrate`. +""" + +from waveorder.config.loader import ( + load_reconstruction_settings, + load_v4_config, +) +from waveorder.config.migrate import ( + MigrationError, + NotRunnableInLegacyError, + detect_version, + migrate_3x_dict, + migrate_3x_settings, + v4_to_legacy, + v4_to_yaml_text, +) +from waveorder.config.schema import ( + Calibration, + DataSettings, + Detection, + FreeParameter, + Illumination, + ImagingSystem, + NoiseModel, + ObjectSettings, + PupilAberrations, + SolverSettings, + V4Config, + free_parameters, + is_free, + nominal_of, +) + +__all__ = [ + "Calibration", + "MigrationError", + "NotRunnableInLegacyError", + "load_reconstruction_settings", + "load_v4_config", + "migrate_3x_settings", + "v4_to_yaml_text", + "DataSettings", + "Detection", + "FreeParameter", + "Illumination", + "ImagingSystem", + "NoiseModel", + "ObjectSettings", + "PupilAberrations", + "SolverSettings", + "V4Config", + "detect_version", + "free_parameters", + "is_free", + "migrate_3x_dict", + "nominal_of", + "v4_to_legacy", +] diff --git a/waveorder/config/loader.py b/waveorder/config/loader.py new file mode 100644 index 00000000..7c54d173 --- /dev/null +++ b/waveorder/config/loader.py @@ -0,0 +1,92 @@ +"""Version-aware loading of reconstruction configurations. + +Every CLI verb loads its configuration through +:func:`load_reconstruction_settings`: + +* a configuration with ``version: 4`` is parsed against the 4.0 schema and + converted to the 3.x settings that select the same code path; +* a configuration without ``version`` is 3.x: it runs as before, and (when + it is migratable) the 4.0 equivalent is printed once per process with a + deprecation warning, per the 4.0 design document section 2. + +Sunset policy (design document section 2): 3.x configurations are accepted +with the deprecation warning for two releases (4.0 and 4.1); after that +they are rejected with a message pointing to ``waveorder migrate-config``. +Birefringence configurations stay 3.x this release and are exempt. +""" + +from __future__ import annotations + +import warnings +from pathlib import Path + +import yaml + +from waveorder.cli.settings import ReconstructionSettings +from waveorder.config.migrate import ( + MigrationError, + detect_version, + migrate_3x_settings, + v4_to_legacy, + v4_to_yaml_text, +) +from waveorder.config.schema import V4Config + +_already_warned_paths: set[str] = set() + + +def load_v4_config(config_path: Path | str) -> V4Config: + """Load a ``version: 4`` configuration file against the 4.0 schema.""" + with open(config_path) as f: + raw = yaml.safe_load(f) + if detect_version(raw) < 4: + raise ValueError(f"{config_path} is a 3.x configuration; run 'waveorder migrate-config' first") + return V4Config(**raw) + + +def load_reconstruction_settings(config_path: Path | str) -> ReconstructionSettings: + """Load a 3.x or 4.0 configuration file into 3.x execution settings. + + Parameters + ---------- + config_path : Path or str + Path to the YAML configuration. + + Returns + ------- + ReconstructionSettings + Settings driving the existing execution paths. For a 4.0 + configuration these are derived by + :func:`waveorder.config.migrate.v4_to_legacy`; for a 3.x + configuration they are parsed directly (and the 4.0 equivalent is + printed once with a deprecation warning). + """ + with open(config_path) as f: + raw = yaml.safe_load(f) + + if detect_version(raw) >= 4: + return v4_to_legacy(V4Config(**raw)) + + settings = ReconstructionSettings(**raw) + _warn_once_about_3x(settings, str(config_path)) + return settings + + +def _warn_once_about_3x(settings: ReconstructionSettings, config_path: str) -> None: + """Print the 4.0 equivalent of a migratable 3.x configuration, once per path.""" + if config_path in _already_warned_paths: + return + _already_warned_paths.add(config_path) + try: + migrated = migrate_3x_settings(settings) + except MigrationError: + # Birefringence configurations stay 3.x in this release; no warning. + return + warnings.warn( + f"{config_path} is a 3.x configuration; it still runs, but 3.x acceptance ends " + f"two releases from now (after 4.1), when it will be rejected. Convert it with: " + f"waveorder migrate-config {config_path}", + DeprecationWarning, + stacklevel=3, + ) + print(f"The 4.0 equivalent of {config_path} is:\n\n{v4_to_yaml_text(migrated)}") diff --git a/waveorder/config/migrate.py b/waveorder/config/migrate.py new file mode 100644 index 00000000..e3622c04 --- /dev/null +++ b/waveorder/config/migrate.py @@ -0,0 +1,423 @@ +"""Deterministic migration between 3.x and 4.0 configurations. + +Three public entry points: + +* :func:`detect_version` — a configuration without ``version`` is 3.x. +* :func:`migrate_3x_dict` — 3.x dict -> :class:`~waveorder.config.schema.V4Config`, + deterministic and lossless for phase and fluorescence configurations. +* :func:`v4_to_legacy` — 4.0 -> the 3.x :class:`ReconstructionSettings` + that selects the same code path with the same numerical parameters. + +The round trip ``v4_to_legacy(migrate_3x_dict(cfg))`` reproduces the +directly-parsed 3.x settings field for field (enforced by unit tests and, +array by array, by the benchmark regression suite). + +Birefringence configurations are not migratable: WaveOrder 4.0 defers the +birefringence refactor, so they continue to run as 3.x without a +deprecation warning. +""" + +from __future__ import annotations + +from typing import Any + +import yaml + +from waveorder.cli.settings import OptimizationSettings, ReconstructionSettings +from waveorder.config import schema +from waveorder.optim._types import OptimizableFloat + +#: Default learning rate used when a 4.0 free parameter has no recorded +#: step size (matches ``OptimizableFloat.learning_rate``). +_DEFAULT_LEARNING_RATE = 0.01 + + +class MigrationError(ValueError): + """A 3.x configuration that cannot be expressed in the 4.0 schema.""" + + +class NotRunnableInLegacyError(ValueError): + """A 4.0 configuration that has no 3.x-equivalent execution path yet.""" + + +def detect_version(raw: dict) -> int: + """Schema version of a raw configuration dict; 3 when ``version`` is absent.""" + if not isinstance(raw, dict): + raise ValueError(f"configuration must be a YAML mapping, got {type(raw).__name__} (empty file?)") + return int(raw.get("version", 3)) + + +# --- 3.x -> 4.0 ------------------------------------------------------------ + + +def _convert_calibratable(value: Any, dotted_name: str, learning_rates: dict[str, float]) -> Any: + """Convert a 3.x number-or-OptimizableFloat into a 4.0 calibratable value. + + An optimizable float with a positive learning rate becomes a + :class:`~waveorder.config.schema.FreeParameter` with a flat prior, and + its learning rate is recorded under ``calibration.advanced.learning_rates``. + A zero learning rate is a fixed value and becomes a plain float. + """ + if isinstance(value, OptimizableFloat): + if value.is_optimizable: + learning_rates[dotted_name] = value.learning_rate + return schema.FreeParameter(nominal=value.initial_value, tolerance="free") + return value.initial_value + return value + + +def migrate_3x_dict(raw: dict) -> schema.V4Config: + """Migrate a 3.x configuration dict to a 4.0 configuration. + + Parameters + ---------- + raw : dict + A parsed 3.x YAML configuration. + + Returns + ------- + V4Config + The equivalent 4.0 configuration. + + Raises + ------ + MigrationError + For birefringence and tile-stitch configurations, which stay 3.x in + this release. + """ + if "recon" in raw and ("tile" in raw or "blend" in raw): + raise MigrationError( + "tile-stitch configurations (tile/blend/recon) run through 'wo tile-stitch' in " + "the 3.x format; there is no 4.0 tile-stitch execution path yet, so they stay 3.x " + "this release (as birefringence configurations do). Migrate the inner 'recon' " + "block on its own if you want a 4.0 reconstruction config." + ) + settings = ReconstructionSettings(**raw) + return migrate_3x_settings(settings) + + +def migrate_3x_settings(settings: ReconstructionSettings) -> schema.V4Config: + """Migrate parsed 3.x settings to a 4.0 configuration. See :func:`migrate_3x_dict`.""" + if settings.birefringence is not None: + raise MigrationError( + "birefringence configurations are not migratable to 4.0 yet; they continue to run in the 3.x format" + ) + + contrast = "phase" if settings.phase is not None else "fluorescence" + block = settings.phase if settings.phase is not None else settings.fluorescence + tf = block.transfer_function + ai = block.apply_inverse + + learning_rates: dict[str, float] = {} + + detection: dict[str, Any] = { + "numerical_aperture": _convert_calibratable( + tf.numerical_aperture_detection, "detection.numerical_aperture", learning_rates + ), + "z_focus_offset": _convert_calibratable(tf.z_focus_offset, "detection.z_focus_offset", learning_rates), + } + illumination: dict[str, Any] = { + "tilt": { + "zenith": _convert_calibratable(tf.tilt_angle_zenith, "illumination.tilt.zenith", learning_rates), + "azimuth": _convert_calibratable(tf.tilt_angle_azimuth, "illumination.tilt.azimuth", learning_rates), + }, + } + if contrast == "phase": + illumination["wavelength"] = tf.wavelength_illumination + illumination["numerical_aperture"] = _convert_calibratable( + tf.numerical_aperture_illumination, "illumination.numerical_aperture", learning_rates + ) + illumination["invert_phase_contrast"] = tf.invert_phase_contrast + else: + detection["wavelength"] = tf.wavelength_emission + if tf.confocal_pinhole_diameter is not None: + detection["pinhole"] = {"diameter_um": tf.confocal_pinhole_diameter} + + # The 3.x reconstruction_algorithm splits into (object.prior, noise, + # solver.reconstruction). + noise: dict[str, Any] = {"gaussian": {"sigma": 1.0}} + algorithm = ai.reconstruction_algorithm + if algorithm == "Tikhonov": + prior: Any = "tikhonov" + reconstruction: Any = "closed_form" + elif algorithm == "TV": + prior = "tv" + reconstruction = {"admm": {"iterations": ai.TV_iterations, "rho": ai.TV_rho_strength}} + elif algorithm in ("RL", "RLGC"): + prior = "none" + noise = {"poisson": {}} + rl = ai.rl + reconstruction = { + "richardson_lucy": { + "iterations": rl.iterations, + "background": rl.background, + "stopping_tolerance": rl.stopping_tolerance, + "back_projector": rl.back_projector, + "bp_alpha": rl.bp_alpha, + "bp_beta": rl.bp_beta, + "gradient_consensus": algorithm == "RLGC", + } + } + else: # pragma: no cover - 3.x schema has no other values + raise MigrationError(f"unknown reconstruction_algorithm {algorithm!r}") + + weight: Any = ai.regularization_strength + if ai.auto_regularization is not None: + weight = {"auto": ai.auto_regularization.model_dump()} + + calibration = None + if settings.optimization is not None or learning_rates: + opt = settings.optimization + advanced: dict[str, Any] = {"learning_rates": learning_rates} + search = [{"gradient": {"max_iter": opt.max_iterations if opt else 200}}] + if opt is not None: + advanced.update( + { + "optimizer": opt.method, + "convergence_tol": opt.convergence_tol, + "convergence_patience": opt.convergence_patience, + "use_gradients": opt.use_gradients, + "grid_points": opt.grid_points, + "loss": opt.loss.model_dump(), + "log_dir": opt.log_dir, + } + ) + calibration = { + "score": "sharpness", + "search": search, + "advanced": advanced, + } + + return schema.V4Config( + version=4, + data={ + "input_channel_names": settings.input_channel_names, + "time_indices": settings.time_indices, + "pixel_size": {"yx": tf.yx_pixel_size, "z": tf.z_pixel_size}, + }, + imaging_system={ + "contrast": contrast, + "index_of_refraction_media": tf.index_of_refraction_media, + "detection": detection, + "illumination": illumination, + }, + object={ + "dimension": settings.reconstruction_dimension, + "prior": prior, + "weight": weight, + }, + noise=noise, + solver={ + "padding": {"z": tf.z_padding}, + "reconstruction": reconstruction, + "calibration": calibration, + "apodization_rolloff": ai.apodization_rolloff, + "device": settings.device, + }, + ) + + +# --- 4.0 -> 3.x ------------------------------------------------------------ + + +def _restore_calibratable(value: Any, dotted_name: str, learning_rates: dict[str, float]) -> Any: + """Inverse of :func:`_convert_calibratable`.""" + if isinstance(value, schema.FreeParameter): + if value.tolerance != "free": + raise NotRunnableInLegacyError( + f"parameter {dotted_name!r} has a numeric tolerance; " + "numeric tolerances are not supported by the 3.x execution path yet" + ) + return { + "initial_value": value.nominal, + "learning_rate": learning_rates.get(dotted_name, _DEFAULT_LEARNING_RATE), + } + return value + + +def _require_legacy_supported(cfg: schema.V4Config) -> None: + """Reject 4.0-only features that have no 3.x execution equivalent yet.""" + + def unsupported(reason: str) -> NotRunnableInLegacyError: + return NotRunnableInLegacyError(f"not runnable on the 3.x execution path yet: {reason}") + + if cfg.imaging_system.field_dependence != "shift_invariant": + raise unsupported("imaging_system.field_dependence: tiled") + # Free parameters with no calibration block would silently fall through + # to the 3.x sharpness optimizer, although the design's default score is + # evidence. Fail loudly instead of guessing the user's intent. + if schema.free_parameters(cfg.imaging_system) and cfg.solver.calibration is None: + raise NotRunnableInLegacyError( + "imaging_system has free parameters (tolerances) but no solver.calibration " + "block. The 4.0 default score (evidence) runs through 'wo calibrate-tf'; to " + "use the 3.x sharpness optimizer through 'wo reconstruct', add " + "solver.calibration: {score: sharpness, search: [{gradient: {max_iter: N}}]}." + ) + detection = cfg.imaging_system.detection + if detection.pupil_aberrations is not None: + raise unsupported("imaging_system.detection.pupil_aberrations (estimate them with 'wo calibrate-tf')") + if detection.pinhole == "point" or ( + isinstance(detection.pinhole, schema.Pinhole) and detection.pinhole.diameter_au is not None + ): + raise unsupported("detection.pinhole: point / diameter_au (use diameter_um)") + if not (isinstance(cfg.object.prior, str) and cfg.object.prior in ("tikhonov", "tv", "none")): + raise unsupported(f"object.prior: {cfg.object.prior}") + if cfg.object.weight == "auto": + raise unsupported("object.weight: auto (bare); the migrated 3.x sweep is weight: {auto: {...}}") + if cfg.object.constraints: + raise unsupported(f"object.constraints: {cfg.object.constraints}") + if cfg.reconstruction_kind == "fista": + raise unsupported("solver.reconstruction: fista") + if cfg.solver.padding.yx not in (0, (0, 0)): + raise unsupported("solver.padding.yx != 0") + calibration = cfg.solver.calibration + if calibration is not None: + if calibration.score != "sharpness": + raise unsupported(f"solver.calibration.score: {calibration.score}") + if calibration.final_estimate != "best_H": + raise unsupported("solver.calibration.final_estimate: average_near_best_H") + if len(calibration.search) != 1 or not isinstance(calibration.search[0], schema.SearchGradient): + raise unsupported("solver.calibration.search stages other than a single gradient stage") + + +def v4_to_legacy(cfg: schema.V4Config) -> ReconstructionSettings: + """Convert a 4.0 configuration to the equivalent 3.x settings. + + Parameters + ---------- + cfg : V4Config + A 4.0 configuration limited to features with a 3.x execution path. + + Returns + ------- + ReconstructionSettings + 3.x settings selecting the same code path with the same numbers. + + Raises + ------ + NotRunnableInLegacyError + When the configuration uses 4.0-only features (tiled field + dependence, pupil aberrations, new priors/scores/stages, ...). + """ + _require_legacy_supported(cfg) + + system = cfg.imaging_system + detection = system.detection + illumination = system.illumination + calibration = cfg.solver.calibration + learning_rates = ( + calibration.advanced.learning_rates if calibration is not None and calibration.advanced is not None else {} + ) + + tf: dict[str, Any] = { + "yx_pixel_size": cfg.data.pixel_size.yx, + "z_pixel_size": cfg.data.pixel_size.z, + "z_padding": cfg.solver.padding.z, + "z_focus_offset": _restore_calibratable(detection.z_focus_offset, "detection.z_focus_offset", learning_rates), + "index_of_refraction_media": system.index_of_refraction_media, + "numerical_aperture_detection": _restore_calibratable( + detection.numerical_aperture, "detection.numerical_aperture", learning_rates + ), + "tilt_angle_zenith": _restore_calibratable( + illumination.tilt.zenith, "illumination.tilt.zenith", learning_rates + ), + "tilt_angle_azimuth": _restore_calibratable( + illumination.tilt.azimuth, "illumination.tilt.azimuth", learning_rates + ), + } + if system.contrast == "phase": + tf["wavelength_illumination"] = illumination.wavelength if illumination.wavelength is not None else 0.532 + tf["numerical_aperture_illumination"] = _restore_calibratable( + illumination.numerical_aperture if illumination.numerical_aperture is not None else 0.9, + "illumination.numerical_aperture", + learning_rates, + ) + tf["invert_phase_contrast"] = illumination.invert_phase_contrast + else: + tf["wavelength_emission"] = detection.wavelength if detection.wavelength is not None else 0.532 + if isinstance(detection.pinhole, schema.Pinhole): + diameter = detection.pinhole.diameter_um + if isinstance(diameter, schema.FreeParameter): + raise NotRunnableInLegacyError("a calibratable pinhole diameter is not supported by 3.x") + tf["confocal_pinhole_diameter"] = diameter + + apply_inverse: dict[str, Any] = { + "apodization_rolloff": cfg.solver.apodization_rolloff, + } + prior = cfg.object.prior + recon = cfg.solver.reconstruction + if prior == "tikhonov": + if cfg.reconstruction_kind != "closed_form": + raise NotRunnableInLegacyError( + f"object.prior: tikhonov runs with solver.reconstruction: closed_form in 3.x, " + f"got {cfg.reconstruction_kind}" + ) + apply_inverse["reconstruction_algorithm"] = "Tikhonov" + elif prior == "tv": + if not isinstance(recon, schema.ReconstructionAdmm): + raise NotRunnableInLegacyError("object.prior: tv requires solver.reconstruction: admm in 3.x") + apply_inverse["reconstruction_algorithm"] = "TV" + apply_inverse["TV_rho_strength"] = recon.admm.rho + apply_inverse["TV_iterations"] = recon.admm.iterations + elif prior == "none": + if not isinstance(recon, schema.ReconstructionRichardsonLucy): + raise NotRunnableInLegacyError("object.prior: none requires solver.reconstruction: richardson_lucy") + if system.contrast != "fluorescence": + raise NotRunnableInLegacyError("richardson_lucy is fluorescence-only in 3.x") + rl = recon.richardson_lucy + apply_inverse["reconstruction_algorithm"] = "RLGC" if rl.gradient_consensus else "RL" + apply_inverse["rl"] = { + "iterations": rl.iterations, + "background": rl.background, + "stopping_tolerance": rl.stopping_tolerance, + "back_projector": rl.back_projector, + "bp_alpha": rl.bp_alpha, + "bp_beta": rl.bp_beta, + } + + weight = cfg.object.weight + if isinstance(weight, schema.WeightAuto): + apply_inverse["auto_regularization"] = weight.auto.model_dump() + else: + apply_inverse["regularization_strength"] = weight + + optimization = None + if calibration is not None and calibration.advanced is not None and calibration.advanced.optimizer is not None: + advanced = calibration.advanced + gradient_stage = calibration.search[0] + optimization = OptimizationSettings( + max_iterations=gradient_stage.gradient.max_iter, + method=advanced.optimizer, + convergence_tol=advanced.convergence_tol, + convergence_patience=advanced.convergence_patience, + use_gradients=advanced.use_gradients, + grid_points=advanced.grid_points, + loss=advanced.loss if advanced.loss is not None else {"type": "midband_power"}, + log_dir=advanced.log_dir, + ) + + contrast_settings = {"transfer_function": tf, "apply_inverse": apply_inverse} + return ReconstructionSettings( + input_channel_names=cfg.data.input_channel_names, + time_indices=cfg.data.time_indices, + reconstruction_dimension=cfg.object.dimension, + device=cfg.solver.device, + optimization=optimization, + **{system.contrast: contrast_settings}, + ) + + +# --- YAML helpers ----------------------------------------------------------- + + +def v4_to_yaml_text(cfg: schema.V4Config) -> str: + """Render a 4.0 configuration as YAML, dropping null-valued fields.""" + + def prune(value: Any) -> Any: + if isinstance(value, dict): + return {k: prune(v) for k, v in value.items() if v is not None} + if isinstance(value, list): + return [prune(v) for v in value] + return value + + return yaml.safe_dump(prune(cfg.model_dump()), default_flow_style=False, sort_keys=False) diff --git a/waveorder/config/schema.py b/waveorder/config/schema.py new file mode 100644 index 00000000..db23b58b --- /dev/null +++ b/waveorder/config/schema.py @@ -0,0 +1,750 @@ +"""Pydantic models for the WaveOrder 4.0 configuration schema. + +The schema follows the WaveOrder 4.0 design document: five blocks (``data``, +``imaging_system``, ``object``, ``noise``, ``solver``) plus ``outputs`` and a +required ``version`` field. + +Conventions +----------- +* Any number under ``imaging_system`` may be written as + ``{nominal: x, tolerance: t}`` (a :class:`FreeParameter`), which makes it a + free parameter of the imaging system, estimated from the data. ``tolerance`` + is the standard deviation of a Gaussian prior in the parameter's own units, + or ``"free"`` for a flat prior. Optional ``bounds: [min, max]`` truncate the + prior. A plain number is fixed. +* Where a field offers several options, its value is a bare name for an + option with no parameters, or a mapping with one key (the option's name) + holding its parameters, e.g. ``reconstruction: closed_form`` or + ``reconstruction: {fista: {iterations: 50}}``. +* Zernike coefficients under ``pupil_aberrations`` are named (``defocus``, + ``astigmatism_0``, ``astigmatism_45``, ``coma_x``, ``coma_y``, + ``trefoil_x``, ``trefoil_y``, ``spherical``, then ``z`` by Noll index) + and are in waves of RMS wavefront error (orthonormal Noll normalization). +""" + +from __future__ import annotations + +import warnings +from typing import Any, Literal, Optional, Union + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + NonNegativeFloat, + NonNegativeInt, + PositiveFloat, + PositiveInt, + TypeAdapter, + field_validator, + model_validator, +) + +from waveorder._pixel_size import YXPixelSize +from waveorder.calibrate.zernike_names import noll_indices_up_to_order, validate_zernike_name + + +class V4BaseModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +# --- Free parameters ------------------------------------------------------- + + +class FreeParameter(V4BaseModel): + """A calibratable imaging-system number: ``{nominal, tolerance, bounds}``. + + Parameters + ---------- + nominal : float + Center of the prior and the fixed value when nothing is estimated. + tolerance : float or "free" + Standard deviation of a Gaussian prior centred on ``nominal``, in the + parameter's own units, or ``"free"`` for a flat prior. + bounds : tuple[float, float], optional + Truncates the prior; the solver clamps at these bounds. + """ + + nominal: float + tolerance: Union[PositiveFloat, Literal["free"]] = "free" + bounds: Optional[tuple[float, float]] = None + + @model_validator(mode="after") + def _nominal_inside_bounds(self): + if self.bounds is not None: + lo, hi = self.bounds + if lo >= hi: + raise ValueError(f"bounds must satisfy min < max, got {self.bounds}") + if not (lo <= self.nominal <= hi): + raise ValueError(f"nominal = {self.nominal} lies outside bounds = {self.bounds}") + return self + + +CalibratableFloat = Union[float, FreeParameter] + + +def nominal_of(value: CalibratableFloat | None) -> float | None: + """Return the nominal value of a plain number or a :class:`FreeParameter`.""" + if value is None: + return None + if isinstance(value, FreeParameter): + return value.nominal + return float(value) + + +def is_free(value: Any) -> bool: + """Whether a value is a :class:`FreeParameter` (i.e. carries a tolerance).""" + return isinstance(value, FreeParameter) + + +# --- data ------------------------------------------------------------------ + + +class PixelSize(V4BaseModel): + """Voxel sizes in micrometers.""" + + yx: Union[PositiveFloat, YXPixelSize] = Field( + description="lateral pixel size in um; scalar for isotropic, or {y, x} mapping", + ) + z: PositiveFloat = Field(description="axial pixel size in um") + + @field_validator("yx", mode="before") + @classmethod + def _normalize_yx(cls, v): + return YXPixelSize.from_value(v) + + +class DataSettings(V4BaseModel): + """The ``data`` block: which data *g* is reconstructed and its sampling.""" + + input_channel_names: list[str] = Field( + default_factory=lambda: ["simulated"], + description="channel names in the input zarr (defaults only make sense for simulate)", + ) + time_indices: Union[NonNegativeInt, list[NonNegativeInt], Literal["all"]] = Field( + default="all", description="time points to reconstruct" + ) + pixel_size: PixelSize + + +# --- imaging_system -------------------------------------------------------- + + +class ZernikeUpToOrder(V4BaseModel): + """Free every Zernike term through radial order ``order`` with nominal 0. + + Piston, tip and tilt (Noll 1-3) are excluded: piston does not affect an + intensity PSF and tip/tilt are degenerate with an object translation. + Named entries in :class:`PupilAberrations` override this shorthand for + their term. + """ + + order: PositiveInt = Field(ge=2, description="highest radial order to free") + tolerance: Union[PositiveFloat, Literal["free"]] = Field( + default="free", description="Gaussian prior std in waves, or 'free'" + ) + bounds: Optional[tuple[float, float]] = None + + +_CALIBRATABLE_ADAPTER = TypeAdapter(CalibratableFloat) + + +class PupilAberrations(BaseModel): + """Zernike pupil aberrations by name; unlisted terms are 0 and fixed. + + Accepts a ``zernike_up_to_order`` shorthand plus any number of named + terms (``defocus``, ``astigmatism_0``, ``astigmatism_45``, ``coma_x``, + ``coma_y``, ``trefoil_x``, ``trefoil_y``, ``spherical``, ``z`` by + Noll index). Coefficients are in waves (RMS, orthonormal Noll). + """ + + model_config = ConfigDict(extra="allow") + + zernike_up_to_order: Optional[ZernikeUpToOrder] = None + + @model_validator(mode="after") + def _validate_terms(self): + extra = self.__pydantic_extra__ or {} + for name, value in extra.items(): + validate_zernike_name(name) + extra[name] = _CALIBRATABLE_ADAPTER.validate_python(value) + return self + + @property + def terms(self) -> dict[str, CalibratableFloat]: + """Named Zernike terms, validated, excluding the shorthand.""" + return dict(self.__pydantic_extra__ or {}) + + +class Pinhole(V4BaseModel): + """Confocal pinhole with a calibratable diameter. + + Exactly one of ``diameter_au`` (Airy units) or ``diameter_um`` + (micrometers in demagnified image space, the 3.x convention) must be set. + """ + + diameter_au: Optional[CalibratableFloat] = None + diameter_um: Optional[CalibratableFloat] = None + + @model_validator(mode="after") + def _exactly_one(self): + if (self.diameter_au is None) == (self.diameter_um is None): + raise ValueError("pinhole requires exactly one of 'diameter_au' or 'diameter_um'") + return self + + +class Detection(V4BaseModel): + """Detection arm of the imaging system.""" + + wavelength: Optional[PositiveFloat] = Field( + default=None, description="emission wavelength in um (fluorescence only)" + ) + numerical_aperture: CalibratableFloat = Field(default=1.2, description="detection objective numerical aperture") + pinhole: Union[Literal["none", "point"], Pinhole] = Field( + default="none", description="none (widefield) | point | {diameter_au | diameter_um}" + ) + pupil_aberrations: Optional[PupilAberrations] = None + z_focus_offset: CalibratableFloat = Field(default=0.0, description="offset from center slice in slice units") + + @model_validator(mode="after") + def _positive_na(self): + na = nominal_of(self.numerical_aperture) + if na is not None and na <= 0: + raise ValueError(f"detection.numerical_aperture must be positive, got {na}") + return self + + +class Tilt(V4BaseModel): + """Illumination tilt in radians.""" + + zenith: CalibratableFloat = 0.0 + azimuth: CalibratableFloat = 0.0 + + +class Illumination(V4BaseModel): + """Illumination arm: transillumination for phase, excitation for confocal.""" + + wavelength: Optional[PositiveFloat] = Field(default=None, description="illumination wavelength in um") + numerical_aperture: Optional[CalibratableFloat] = Field( + default=None, description="condenser / excitation numerical aperture" + ) + tilt: Tilt = Field(default_factory=Tilt) + invert_phase_contrast: bool = Field(default=False, description="phase only") + + +class TiledFieldDependence(V4BaseModel): + """Tiled shift-variance: one parameter set per tile, blended in overlaps.""" + + size_yx: tuple[PositiveInt, PositiveInt] = Field(description="tile size in pixels (y, x)") + overlap_yx: tuple[NonNegativeInt, NonNegativeInt] = Field(default=(32, 32), description="tile overlap in pixels") + tile_blend: Literal["linear", "cosine"] = "linear" + + +class FieldDependenceTiled(V4BaseModel): + tiled: TiledFieldDependence + + +class ImagingSystem(V4BaseModel): + """The ``imaging_system`` block: H and its prior (via tolerances).""" + + contrast: Literal["phase", "fluorescence"] + field_dependence: Union[Literal["shift_invariant"], FieldDependenceTiled] = "shift_invariant" + index_of_refraction_media: PositiveFloat = 1.3 + detection: Detection = Field(default_factory=Detection) + illumination: Illumination = Field(default_factory=Illumination) + + @model_validator(mode="after") + def _na_below_index(self): + na_det = nominal_of(self.detection.numerical_aperture) + if na_det is not None and na_det > self.index_of_refraction_media: + raise ValueError( + f"detection.numerical_aperture = {na_det} must be less than or equal to " + f"index_of_refraction_media = {self.index_of_refraction_media}" + ) + na_ill = nominal_of(self.illumination.numerical_aperture) + if na_ill is not None and na_ill > self.index_of_refraction_media: + raise ValueError( + f"illumination.numerical_aperture = {na_ill} must be less than or equal to " + f"index_of_refraction_media = {self.index_of_refraction_media}" + ) + return self + + +# --- object ---------------------------------------------------------------- + + +class BeadsPrior(V4BaseModel): + diameter: PositiveFloat = Field(description="bead diameter in um; sparse spheres at unknown positions") + snr_threshold: PositiveFloat = Field( + default=6.0, + description="initial detection threshold in robust noise units (median + snr_threshold * 1.4826 * MAD); " + "detection relaxes adaptively down to half this value before failing", + ) + max_beads: PositiveInt = Field(default=50, description="fit at most this many beads, brightest first") + min_separation: PositiveInt = Field(default=5, description="non-maximum suppression window in voxels") + + +class DenoiserPrior(V4BaseModel): + checkpoint: str + + +class PriorBeads(V4BaseModel): + beads: BeadsPrior + + +class PriorDenoiser(V4BaseModel): + denoiser: DenoiserPrior + + +Prior = Union[ + Literal["tikhonov", "tv", "none"], + PriorBeads, + PriorDenoiser, +] + + +class LegacyAutoWeight(V4BaseModel): + """The 3.x ``auto_regularization`` sweep, carried losslessly through migration. + + The bare ``weight: auto`` of the 4.0 design (score-based selection of + lambda) is a separate option; this block keeps the 3.x rule-based sweep + (``otsu_cnr`` / ``l_curve``) available to migrated configurations. + """ + + rule: Literal["otsu_cnr", "l_curve"] = "otsu_cnr" + search_min: float = -6.0 + search_max: float = 2.0 + num_samples: PositiveInt = 25 + search_scale: Literal["relative", "absolute"] = "relative" + report_path: Optional[str] = None + plot_path: Optional[str] = None + + +class WeightAuto(V4BaseModel): + auto: LegacyAutoWeight + + +class GaussianSpectrumSource(V4BaseModel): + """Gaussian random field with a power-law spectrum (simulation only).""" + + exponent: float = Field(default=2.0, description="power-spectrum decay exponent") + + +class PhantomSource(V4BaseModel): + kind: str + + +class FromReconSource(V4BaseModel): + path: str + + +class SourceGaussianSpectrum(V4BaseModel): + gaussian_spectrum: GaussianSpectrumSource + + +class SourcePhantom(V4BaseModel): + phantom: PhantomSource + + +class SourceFromRecon(V4BaseModel): + from_recon: FromReconSource + + +Source = Union[SourceGaussianSpectrum, SourcePhantom, SourceFromRecon] + + +class ObjectSettings(V4BaseModel): + """The ``object`` block: the prior on *f* (or its ``source`` in a simulation).""" + + dimension: Literal[2, 3] = Field(default=3, description="3 | 2 (thin object, a single plane)") + prior: Prior = "tikhonov" + weight: Union[NonNegativeFloat, Literal["auto"], WeightAuto] = Field( + default=1e-3, description="regularization weight lambda, or auto" + ) + constraints: list[Literal["nonneg"]] = Field(default_factory=list) + source: Optional[Source] = Field( + default=None, description="the true object of a simulation configuration (design section 5)" + ) + + +# --- noise ----------------------------------------------------------------- + + +class DarkFrames(V4BaseModel): + dark_frames: str + + +class GaussianNoise(V4BaseModel): + sigma: Union[PositiveFloat, Literal["estimate"], DarkFrames] = 1.0 + snr_db: Optional[float] = Field( + default=None, description="simulation only: sets sigma from the clean signal's power" + ) + + +class PoissonNoise(V4BaseModel): + gain: PositiveFloat = 1.0 + offset: float = 0.0 + photons: Optional[PositiveFloat] = Field( + default=None, description="simulation only: mean photons per voxel of the clean signal" + ) + + +class PoissonGaussianNoise(V4BaseModel): + gain: PositiveFloat = 1.0 + offset: float = 0.0 + sigma: Union[PositiveFloat, Literal["estimate"], DarkFrames] = 1.0 + + +class NoiseModel(V4BaseModel): + """The ``noise`` block: exactly one of gaussian / poisson / poisson_gaussian.""" + + gaussian: Optional[GaussianNoise] = None + poisson: Optional[PoissonNoise] = None + poisson_gaussian: Optional[PoissonGaussianNoise] = None + + @model_validator(mode="after") + def _exactly_one(self): + set_options = [name for name in ("gaussian", "poisson", "poisson_gaussian") if getattr(self, name) is not None] + if len(set_options) == 0: + self.gaussian = GaussianNoise() + elif len(set_options) > 1: + raise ValueError(f"noise accepts exactly one option, got {set_options}") + return self + + @property + def kind(self) -> str: + for name in ("gaussian", "poisson", "poisson_gaussian"): + if getattr(self, name) is not None: + return name + return "gaussian" + + +# --- solver ---------------------------------------------------------------- + + +class Padding(V4BaseModel): + # Default 0 preserves 3.x behavior exactly (3.x z_padding defaults to 0). + # The design document's section 2 text says "padding: {z: 5, yx: 0}", + # which contradicts its own defaults-reproduce-3.x requirement; the + # deviation is documented in PR #592. + z: NonNegativeInt = Field(default=0, description="z slices to pad for axial boundary effects") + yx: Union[NonNegativeInt, tuple[NonNegativeInt, NonNegativeInt]] = Field( + default=0, description="transverse padding in pixels; number or [y, x]" + ) + + +class FistaReconstruction(V4BaseModel): + iterations: PositiveInt = 50 + + +class AdmmReconstruction(V4BaseModel): + iterations: PositiveInt = 1 + rho: PositiveFloat = 1e-3 + + +class RichardsonLucyReconstruction(V4BaseModel): + """Richardson-Lucy options, matching the 3.x ``rl`` block field for field.""" + + iterations: PositiveInt = 25 + background: NonNegativeFloat = 0.0 + stopping_tolerance: Optional[NonNegativeFloat] = None + back_projector: Literal["matched", "gaussian", "butterworth", "wiener", "wiener_butterworth"] = "matched" + bp_alpha: Optional[PositiveFloat] = None + bp_beta: Optional[PositiveFloat] = None + gradient_consensus: bool = Field( + default=False, description="use the RLGC gradient-consensus stopping variant (3.x 'RLGC')" + ) + + +class ReconstructionFista(V4BaseModel): + fista: FistaReconstruction + + +class ReconstructionAdmm(V4BaseModel): + admm: AdmmReconstruction + + +class ReconstructionRichardsonLucy(V4BaseModel): + richardson_lucy: RichardsonLucyReconstruction + + +Reconstruction = Union[ + Literal["closed_form"], + ReconstructionFista, + ReconstructionAdmm, + ReconstructionRichardsonLucy, +] + + +class GridStage(V4BaseModel): + points: Union[PositiveInt, dict[str, PositiveInt]] = 7 + + +class RandomStage(V4BaseModel): + samples: PositiveInt = 16 + seed: int = 0 + + +class GradientStage(V4BaseModel): + max_iter: PositiveInt = 200 + method: Literal["split", "alternating", "joint"] = "split" + + +class SearchGrid(V4BaseModel): + grid: GridStage + + +class SearchRandom(V4BaseModel): + random: RandomStage + + +class SearchGradient(V4BaseModel): + gradient: GradientStage + + +SearchStage = Union[SearchGrid, SearchRandom, SearchGradient, Literal["grid", "random", "gradient"]] + + +def _normalize_search_stage(stage: SearchStage) -> Union[SearchGrid, SearchRandom, SearchGradient]: + """Expand a bare stage name into its option model with defaults.""" + if stage == "grid": + return SearchGrid(grid=GridStage()) + if stage == "random": + return SearchRandom(random=RandomStage()) + if stage == "gradient": + return SearchGradient(gradient=GradientStage()) + return stage + + +class CalibrationAdvanced(V4BaseModel): + """Step sizes and other knobs of the search; carries 3.x optimizer settings.""" + + optimizer: Optional[Literal["adam", "nadam", "lbfgs", "nelder_mead", "grid_search"]] = Field( + default=None, description="3.x optimizer backend; null = not carried from a 3.x optimization block" + ) + learning_rates: dict[str, float] = Field( + default_factory=dict, + description="per-parameter gradient step sizes, keyed by dotted parameter name", + ) + convergence_tol: Optional[float] = None + convergence_patience: Optional[PositiveInt] = 5 + use_gradients: Optional[bool] = None + grid_points: PositiveInt = 7 + loss: Optional[dict[str, Any]] = Field(default=None, description="3.x sharpness loss settings (type + parameters)") + log_dir: Optional[str] = None + + +class Calibration(V4BaseModel): + """How H is estimated; read only when something has a tolerance.""" + + score: Literal["evidence", "gcv", "sharpness"] = "evidence" + search: list[SearchStage] = Field(default_factory=lambda: [SearchGradient(gradient=GradientStage())]) + final_estimate: Literal["best_H", "average_near_best_H"] = "best_H" + advanced: Optional[CalibrationAdvanced] = None + + @field_validator("search", mode="after") + @classmethod + def _normalize_stages(cls, v): + return [_normalize_search_stage(stage) for stage in v] + + +class SolverSettings(V4BaseModel): + """The ``solver`` block: numerics of the inner and outer problems.""" + + padding: Padding = Field(default_factory=Padding) + reconstruction: Reconstruction = "closed_form" + calibration: Optional[Calibration] = None + apodization_rolloff: float = Field( + default=0.0, + ge=0.0, + le=1.0, + description="raised-cosine roll-off of the inverse filter at the transverse Nyquist edge", + ) + device: Optional[str] = Field(default=None, description="compute device: null=cpu, auto, cuda:0, mps, ...") + + +# --- outputs --------------------------------------------------------------- + + +OutputName = Literal[ + "reconstruction", + "transfer_function", + "psf", + "parameters", + "parameters_covariance", + "score_trace", + "score_grid", + "weight_trace", + "residual", + "tile_map", +] + +_DEFAULT_OUTPUTS: tuple[str, ...] = ("reconstruction", "transfer_function") + + +# --- top level -------------------------------------------------------------- + + +class SimulateSettings(V4BaseModel): + """The ``simulate`` block of a simulation configuration (design section 5).""" + + seed: int = 0 + shape: tuple[PositiveInt, PositiveInt, PositiveInt] + + +class V4Config(V4BaseModel): + """A complete WaveOrder 4.0 configuration.""" + + version: Literal[4] + data: DataSettings + imaging_system: ImagingSystem + object: ObjectSettings = Field(default_factory=ObjectSettings) + noise: NoiseModel = Field(default_factory=NoiseModel) + solver: SolverSettings = Field(default_factory=SolverSettings) + outputs: Union[Literal["all"], list[OutputName]] = Field(default_factory=lambda: list(_DEFAULT_OUTPUTS)) + simulate: Optional[SimulateSettings] = None + + @model_validator(mode="after") + def _validation_rules(self): + """The cross-block validation rules of design section 4.5.""" + params = free_parameters(self.imaging_system) + anything_estimated = len(params) > 0 + calibration = self.solver.calibration + + if isinstance(self.object.prior, PriorBeads) and not anything_estimated: + raise ValueError("object.prior: beads requires a free imaging-system parameter (beads are for calibration)") + + if calibration is not None and anything_estimated: + if ( + calibration.score == "sharpness" + and self.imaging_system.field_dependence != "shift_invariant" + and all(p.tolerance == "free" for p in params.values()) + ): + warnings.warn( + "tiled field dependence with every tolerance 'free' and score: sharpness " + "estimates each tile independently with no prior; expect overfitting on " + "low-signal tiles.", + UserWarning, + ) + if calibration.score == "evidence": + prior_ok = self.object.prior in ("tikhonov",) or isinstance(self.object.prior, PriorBeads) + if not prior_ok: + raise ValueError("score: evidence requires object.prior in {tikhonov, beads}") + if self.noise.kind == "poisson": + raise ValueError("score: evidence requires Gaussian-family noise") + for stage in calibration.search: + if isinstance(stage, (SearchGrid, SearchRandom)): + unbounded = [name for name, p in params.items() if p.tolerance == "free" and p.bounds is None] + if unbounded: + raise ValueError( + f"a grid or random search stage requires bounds for parameters with a " + f"'free' tolerance; missing bounds on: {unbounded}" + ) + if isinstance(stage, SearchGrid): + points = stage.grid.points + count = 1 + for name in params: + n = points if isinstance(points, int) else points.get(name, 7) + if name.endswith(".zernike_up_to_order"): + shorthand = self.imaging_system.detection.pupil_aberrations.zernike_up_to_order + count *= n ** len(noll_indices_up_to_order(shorthand.order)) + else: + count *= n + if count > 10_000: + warnings.warn( + f"the grid search stage evaluates {count} points over " + f"{len(params)} parameters; consider fewer points per " + f"parameter or a random stage.", + UserWarning, + ) + if ( + isinstance(stage, SearchGradient) + and stage.gradient.method == "joint" + and calibration.score == "sharpness" + ): + warnings.warn( + "gradient method 'joint' with score 'sharpness' collapses to the trivial " + "solution; use 'split' or 'alternating', or an evidence-family score.", + UserWarning, + ) + + if self.object.weight == "auto": + score = calibration.score if calibration is not None else None + if anything_estimated and score not in ("gcv", "evidence"): + raise ValueError( + "weight: auto requires solver.calibration.score in {gcv, evidence} " + "when imaging-system parameters are estimated" + ) + + if self.object.prior == "none" and self.reconstruction_kind == "closed_form": + raise ValueError("object.prior: none with solver.reconstruction: closed_form is an unregularized inverse") + if ( + self.object.prior == "tv" or isinstance(self.object.prior, PriorDenoiser) or self.object.constraints + ) and self.reconstruction_kind == "closed_form": + raise ValueError( + "object.prior tv/denoiser and constraints: [nonneg] have no closed form; " + "choose fista, admm, or richardson_lucy under solver.reconstruction" + ) + if self.reconstruction_kind == "richardson_lucy": + if self.noise.kind != "poisson": + raise ValueError("richardson_lucy requires noise: poisson") + if self.object.prior != "none": + raise ValueError("richardson_lucy requires object.prior: none") + return self + + @property + def reconstruction_kind(self) -> str: + """Bare name of the chosen reconstruction option.""" + recon = self.solver.reconstruction + if isinstance(recon, str): + return recon + return next(iter(recon.__class__.model_fields)) + + @property + def contrast(self) -> str: + return self.imaging_system.contrast + + +def free_parameters(system: ImagingSystem) -> dict[str, FreeParameter]: + """Collect every :class:`FreeParameter` of an imaging system. + + Parameters + ---------- + system : ImagingSystem + The ``imaging_system`` block. + + Returns + ------- + dict[str, FreeParameter] + Mapping from dotted parameter name (e.g. + ``"detection.numerical_aperture"``, + ``"detection.pupil_aberrations.spherical"``) to its prior. The + ``zernike_up_to_order`` shorthand is expanded into named terms by + :mod:`waveorder.calibrate.zernike_names` consumers, not here; it + appears under its own key when present with a tolerance. + """ + found: dict[str, FreeParameter] = {} + + def visit(obj: Any, prefix: str) -> None: + if isinstance(obj, FreeParameter): + found[prefix] = obj + return + if isinstance(obj, PupilAberrations): + shorthand = obj.zernike_up_to_order + if shorthand is not None: + found[f"{prefix}.zernike_up_to_order"] = FreeParameter( + nominal=0.0, tolerance=shorthand.tolerance, bounds=shorthand.bounds + ) + for name, value in obj.terms.items(): + visit(value, f"{prefix}.{name}") + return + if isinstance(obj, BaseModel): + for name in obj.__class__.model_fields: + value = getattr(obj, name) + if isinstance(value, (BaseModel,)): + visit(value, f"{prefix}.{name}" if prefix else name) + elif isinstance(value, FreeParameter): + visit(value, f"{prefix}.{name}" if prefix else name) + + visit(system, "") + return found diff --git a/waveorder/config/simulate.py b/waveorder/config/simulate.py new file mode 100644 index 00000000..ffe82c96 --- /dev/null +++ b/waveorder/config/simulate.py @@ -0,0 +1,204 @@ +"""The 4.0 ``simulate`` verb: write ``data.zarr``, ``object.zarr``, ``tf.zarr``. + +Minimal implementation of design document section 5: fluorescence +contrast, shift-invariant imaging system, ``gaussian_spectrum`` object +source, and Gaussian (``sigma`` or ``snr_db``) or Poisson (``photons``) +noise. A ``simulate`` configuration describes the *true* system: any +imaging-system number written with a numeric ``tolerance`` draws the truth +from that prior (clamped to its bounds) under the configuration's seed. + +Outputs land in the standard formats: ``data.zarr`` and ``object.zarr`` +are single-position HCS OME-Zarr stores (consumable by every verb), and +``tf.zarr`` holds the drawn free parameters, a zero covariance, the true +transfer function, and ``provenance: {kind: simulated, ...}``. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +from iohub.ngff import open_ome_zarr +from iohub.ngff.models import TransformationMeta + +from waveorder.calibrate.synthesis import fluorescence_otf +from waveorder.calibrate.tfzarr import write_tf_zarr +from waveorder.calibrate.zernike_names import name_to_noll, noll_indices_up_to_order, noll_to_name +from waveorder.config.schema import ( + FreeParameter, + SourceGaussianSpectrum, + V4Config, +) + + +def _draw_truth_coefficients(cfg: V4Config, generator: torch.Generator) -> dict[str, float]: + """True Zernike coefficients: fixed values, and draws from numeric tolerances.""" + aberrations = cfg.imaging_system.detection.pupil_aberrations + truth: dict[str, float] = {} + if aberrations is None: + return truth + + def draw(nominal: float, tolerance, bounds) -> float: + if tolerance == "free" or tolerance is None: + return float(nominal) + value = nominal + float(tolerance) * float(torch.randn((), generator=generator)) + if bounds is not None: + value = min(max(value, bounds[0]), bounds[1]) + return float(value) + + shorthand = aberrations.zernike_up_to_order + if shorthand is not None: + for j in noll_indices_up_to_order(shorthand.order): + truth[noll_to_name(j)] = draw(0.0, shorthand.tolerance, shorthand.bounds) + for name, value in aberrations.terms.items(): + canonical = noll_to_name(name_to_noll(name)) + if isinstance(value, FreeParameter): + truth[canonical] = draw(value.nominal, value.tolerance, value.bounds) + else: + truth[canonical] = float(value) + return truth + + +def _gaussian_spectrum_object( + zyx_shape: tuple[int, int, int], + exponent: float, + generator: torch.Generator, +) -> torch.Tensor: + """Zero-mean, unit-std Gaussian random field with power spectrum ``(|k|^2 + k0^2)^(-exponent/2)``.""" + freqs = [torch.fft.fftfreq(n) for n in zyx_shape] + kz, ky, kx = torch.meshgrid(*freqs, indexing="ij") + k2 = kz**2 + ky**2 + kx**2 + k0 = 1.0 / max(zyx_shape) + amplitude = (k2 + k0**2) ** (-exponent / 4.0) + white = torch.randn(zyx_shape, generator=generator) + field = torch.real(torch.fft.ifftn(torch.fft.fftn(white) * amplitude)) + return (field - field.mean()) / field.std().clamp(min=1e-12) + + +def _write_position(path: Path, zyx: np.ndarray, channel_name: str, scale_zyx: tuple[float, float, float]) -> None: + """Write one TCZYX = (1, 1, Z, Y, X) volume as a single-position HCS store.""" + dataset = open_ome_zarr(path, layout="hcs", mode="w", channel_names=[channel_name]) + position = dataset.create_position("0", "0", "0") + position.create_image( + "0", + zyx[None, None].astype(np.float32), + transform=[TransformationMeta(type="scale", scale=[1.0, 1.0, *scale_zyx])], + ) + dataset.close() + + +def run_simulation(cfg: V4Config, output_dirpath: Path) -> dict[str, float]: + """Simulate data from the true imaging system of a 4.0 configuration. + + Parameters + ---------- + cfg : V4Config + Configuration with a ``simulate`` block (seed and shape), a + fluorescence ``imaging_system``, an ``object.source``, and a noise + model. + output_dirpath : Path + Directory receiving ``data.zarr``, ``object.zarr``, ``tf.zarr``. + + Returns + ------- + dict[str, float] + The drawn true Zernike coefficients (waves). + """ + if cfg.simulate is None: + raise ValueError("simulate requires a 'simulate: {seed, shape}' block (design section 5)") + if cfg.imaging_system.contrast != "fluorescence": + raise NotImplementedError("simulate supports fluorescence contrast for now") + if cfg.imaging_system.field_dependence != "shift_invariant": + raise NotImplementedError("simulate is shift-invariant for now") + if not isinstance(cfg.object.source, SourceGaussianSpectrum): + raise NotImplementedError("simulate supports object.source: {gaussian_spectrum: {exponent}} for now") + + detection = cfg.imaging_system.detection + if detection.wavelength is None: + raise ValueError("imaging_system.detection.wavelength is required") + yx = cfg.data.pixel_size.yx + yx_pixel_size = float(yx) if isinstance(yx, (int, float)) else float(yx.y) + z_pixel_size = float(cfg.data.pixel_size.z) + + zyx_shape = tuple(cfg.simulate.shape) + generator = torch.Generator().manual_seed(cfg.simulate.seed) + + truth = _draw_truth_coefficients(cfg, generator) + numerical_aperture = detection.numerical_aperture + if isinstance(numerical_aperture, FreeParameter): + raise NotImplementedError("drawing a true numerical aperture is not supported yet; use a plain number") + + otf = fluorescence_otf( + zyx_shape=zyx_shape, + yx_pixel_size=yx_pixel_size, + z_pixel_size=z_pixel_size, + wavelength_emission=detection.wavelength, + index_of_refraction_media=cfg.imaging_system.index_of_refraction_media, + numerical_aperture_detection=numerical_aperture, + zernike_coefficients=truth, + ) + + exponent = cfg.object.source.gaussian_spectrum.exponent + zyx_object = _gaussian_spectrum_object(zyx_shape, exponent, generator) + clean = torch.real(torch.fft.ifftn(torch.fft.fftn(zyx_object) * otf)) + + if cfg.noise.kind == "gaussian": + noise = cfg.noise.gaussian + if noise.snr_db is not None: + sigma = float(clean.std()) / (10.0 ** (noise.snr_db / 20.0)) + elif isinstance(noise.sigma, (int, float)): + sigma = float(noise.sigma) + else: + raise ValueError("simulation needs noise.gaussian.sigma as a number, or snr_db") + data = clean + sigma * torch.randn(zyx_shape, generator=generator) + noise_description = {"kind": "gaussian", "sigma": sigma} + elif cfg.noise.kind == "poisson": + noise = cfg.noise.poisson + if noise.photons is None: + raise ValueError("simulation needs noise.poisson.photons") + nonneg = clean - clean.min() + scale = noise.photons / float(nonneg.mean().clamp(min=1e-12)) + data = torch.poisson(nonneg * scale, generator=generator) + noise_description = {"kind": "poisson", "photons": float(noise.photons), "scale": scale} + else: + raise NotImplementedError(f"simulate does not support noise: {cfg.noise.kind} yet") + + output_dirpath = Path(output_dirpath) + output_dirpath.mkdir(parents=True, exist_ok=True) + channel_name = cfg.data.input_channel_names[0] + scale_zyx = (z_pixel_size, yx_pixel_size, yx_pixel_size) + _write_position(output_dirpath / "data.zarr", data.numpy(), channel_name, scale_zyx) + _write_position(output_dirpath / "object.zarr", zyx_object.numpy(), "object", scale_zyx) + + # tf.zarr: P = the parameters with a tolerance; fixed terms live in provenance. + aberrations = cfg.imaging_system.detection.pupil_aberrations + free_names: list[str] = [] + if aberrations is not None: + if aberrations.zernike_up_to_order is not None: + free_names.extend(noll_to_name(j) for j in noll_indices_up_to_order(aberrations.zernike_up_to_order.order)) + for name, value in aberrations.terms.items(): + canonical = noll_to_name(name_to_noll(name)) + if isinstance(value, FreeParameter) and canonical not in free_names: + free_names.append(canonical) + free_names.sort(key=name_to_noll) + parameters = np.array([[truth[name] for name in free_names]], dtype=np.float64) + + write_tf_zarr( + output_dirpath / "tf.zarr", + parameters=parameters, + parameter_names=free_names, + tile_grid={ + "size_yx": [int(zyx_shape[1]), int(zyx_shape[2])], + "overlap_yx": [0, 0], + "centers_yx": [[0.0, 0.0]], + }, + provenance={ + "kind": "simulated", + "seed": cfg.simulate.seed, + "true_coefficients_waves": truth, + "noise": noise_description, + }, + transfer_function=otf.numpy().astype(np.complex64), + ) + return truth