From 30da147bb32e0df49ae59fe0b49c0ab0f4a88e93 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:05:47 -0700 Subject: [PATCH 01/24] feat(calibrate): Zernike-named fluorescence PSF/OTF synthesis and evidence score Add waveorder/calibrate: named Zernike terms (Noll indexing, orthonormal, coefficients in waves), differentiable widefield fluorescence PSF/OTF synthesis from a parameter mapping (matches the shift-variant model's numerics at the field center), the exact Whittle/Gaussian negative log evidence for Gaussian noise with a Tikhonov or spectral object prior, and the calibration penalty from parameter tolerances. --- tests/calibrate/__init__.py | 0 tests/calibrate/test_synthesis_scores.py | 143 +++++++++++++++ tests/calibrate/test_zernike_names.py | 78 ++++++++ waveorder/calibrate/__init__.py | 39 ++++ waveorder/calibrate/scores.py | 132 ++++++++++++++ waveorder/calibrate/synthesis.py | 217 +++++++++++++++++++++++ waveorder/calibrate/zernike_names.py | 117 ++++++++++++ 7 files changed, 726 insertions(+) create mode 100644 tests/calibrate/__init__.py create mode 100644 tests/calibrate/test_synthesis_scores.py create mode 100644 tests/calibrate/test_zernike_names.py create mode 100644 waveorder/calibrate/__init__.py create mode 100644 waveorder/calibrate/scores.py create mode 100644 waveorder/calibrate/synthesis.py create mode 100644 waveorder/calibrate/zernike_names.py 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_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/waveorder/calibrate/__init__.py b/waveorder/calibrate/__init__.py new file mode 100644 index 00000000..38cce3c4 --- /dev/null +++ b/waveorder/calibrate/__init__.py @@ -0,0 +1,39 @@ +"""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. +""" + +from waveorder.calibrate.scores import ( + calibration_penalty, + whittle_evidence, +) +from waveorder.calibrate.synthesis import ( + fluorescence_otf, + fluorescence_psf, +) +from waveorder.calibrate.zernike_names import ( + NAMED_TERMS, + name_to_noll, + noll_indices_up_to_order, + noll_to_name, + validate_zernike_name, +) + +__all__ = [ + "NAMED_TERMS", + "calibration_penalty", + "fluorescence_otf", + "fluorescence_psf", + "name_to_noll", + "noll_indices_up_to_order", + "noll_to_name", + "validate_zernike_name", + "whittle_evidence", +] diff --git a/waveorder/calibrate/scores.py b/waveorder/calibrate/scores.py new file mode 100644 index 00000000..6c924c1b --- /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``), + for ``tikhonov_spectral`` priors. + 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..5021ddf1 --- /dev/null +++ b/waveorder/calibrate/synthesis.py @@ -0,0 +1,217 @@ +"""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`). 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, with the focal + plane at index ``Z // 2`` after ``fftshift`` along z (the array is + returned in FFT order along z, matching the shift-variant model: + the in-focus slice is at index 0). + """ + 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/zernike_names.py b/waveorder/calibrate/zernike_names.py new file mode 100644 index 00000000..40f53f1a --- /dev/null +++ b/waveorder/calibrate/zernike_names.py @@ -0,0 +1,117 @@ +"""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)``. +* 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 From 822e3b8036adad550614324d4932f3c999be27f9 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:05:55 -0700 Subject: [PATCH 02/24] feat(config): WaveOrder 4.0 configuration schema and 3.x migration Add waveorder/config: the five-block 4.0 pydantic schema (data, imaging_system, object, noise, solver + outputs and version) per the 4.0 design document, version detection, a deterministic lossless 3.x -> 4.0 migration for phase and fluorescence configurations, the inverse 4.0 -> 3.x settings conversion that selects the same code path, and a version-aware loader. The migrate/convert round trip reproduces the directly-parsed 3.x settings field for field for every configuration in the repository (unit-tested). Birefringence configurations stay 3.x. --- tests/config/__init__.py | 0 tests/config/test_migrate.py | 220 +++++++++++ tests/config/test_v4_schema.py | 205 ++++++++++ waveorder/config/__init__.py | 63 +++ waveorder/config/loader.py | 86 +++++ waveorder/config/migrate.py | 403 +++++++++++++++++++ waveorder/config/schema.py | 679 +++++++++++++++++++++++++++++++++ 7 files changed, 1656 insertions(+) create mode 100644 tests/config/__init__.py create mode 100644 tests/config/test_migrate.py create mode 100644 tests/config/test_v4_schema.py create mode 100644 waveorder/config/__init__.py create mode 100644 waveorder/config/loader.py create mode 100644 waveorder/config/migrate.py create mode 100644 waveorder/config/schema.py diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/config/test_migrate.py b/tests/config/test_migrate.py new file mode 100644 index 00000000..3e3a9bc1 --- /dev/null +++ b/tests/config/test_migrate.py @@ -0,0 +1,220 @@ +"""3.x -> 4.0 migration: deterministic, lossless, and round-trippable.""" + +from pathlib import Path + +import pytest +import yaml + +from waveorder.cli.settings import ReconstructionSettings +from waveorder.config import ( + MigrationError, + detect_version, + migrate_3x_dict, + v4_to_legacy, + v4_to_yaml_text, +) +from waveorder.config.schema import V4Config + +REPO_ROOT = Path(__file__).parent.parent.parent + +# Every 3.x phase / fluorescence configuration in the repository. +MIGRATABLE_CONFIG_PATHS = sorted( + [ + *(REPO_ROOT / "docs/examples/cli/configs").glob("*.yml"), + *(REPO_ROOT / "benchmarks/configs").glob("*.yml"), + *(REPO_ROOT / "docs/examples/optimization").glob("*.yml"), + ] +) +MIGRATABLE_CONFIG_PATHS = [ + p for p in MIGRATABLE_CONFIG_PATHS if "birefringence" not in p.name and "tile-stitch" not in p.name +] + +BIREFRINGENCE_CONFIG_PATHS = sorted((REPO_ROOT / "docs/examples/cli/configs").glob("birefringence*.yml")) + + +def _load(path: Path) -> dict: + with open(path) as f: + return yaml.safe_load(f) + + +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)) + + +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() diff --git a/tests/config/test_v4_schema.py b/tests/config/test_v4_schema.py new file mode 100644 index 00000000..571f60f6 --- /dev/null +++ b/tests/config/test_v4_schema.py @@ -0,0 +1,205 @@ +"""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) 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..9528d284 --- /dev/null +++ b/waveorder/config/loader.py @@ -0,0 +1,86 @@ +"""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. +""" + +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 will be rejected in a " + f"future release. Convert it with: 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..fff91060 --- /dev/null +++ b/waveorder/config/migrate.py @@ -0,0 +1,403 @@ +"""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.""" + 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 configurations, which stay 3.x in this release. + """ + 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") + detection = cfg.imaging_system.detection + if detection.pupil_aberrations is not None: + raise unsupported("imaging_system.detection.pupil_aberrations") + 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..ad278c1d --- /dev/null +++ b/waveorder/config/schema.py @@ -0,0 +1,679 @@ +"""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 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(description="channel names in the input zarr") + 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 TikhonovSpectralPrior(V4BaseModel): + exponent: Optional[float] = Field(default=None, description="power-law exponent of the object spectrum") + fit_from: Optional[str] = Field(default=None, description="path to fit the spectrum from") + + +class BeadsPrior(V4BaseModel): + diameter: PositiveFloat = Field(description="bead diameter in um; sparse spheres at unknown positions") + + +class DenoiserPrior(V4BaseModel): + checkpoint: str + + +class PriorTikhonovSpectral(V4BaseModel): + tikhonov_spectral: TikhonovSpectralPrior + + +class PriorBeads(V4BaseModel): + beads: BeadsPrior + + +class PriorDenoiser(V4BaseModel): + denoiser: DenoiserPrior + + +Prior = Union[ + Literal["tikhonov", "tv", "none"], + PriorTikhonovSpectral, + 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 ObjectSettings(V4BaseModel): + """The ``object`` block: the prior on *f*.""" + + 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) + + +# --- noise ----------------------------------------------------------------- + + +class DarkFrames(V4BaseModel): + dark_frames: str + + +class GaussianNoise(V4BaseModel): + sigma: Union[PositiveFloat, Literal["estimate"], DarkFrames] = 1.0 + + +class PoissonNoise(V4BaseModel): + gain: PositiveFloat = 1.0 + offset: float = 0.0 + + +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): + 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", "sure", "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 calibration is not None and anything_estimated: + if calibration.score == "evidence": + prior_ok = self.object.prior in ("tikhonov",) or isinstance( + self.object.prior, (PriorTikhonovSpectral, PriorBeads) + ) + if not prior_ok: + raise ValueError("score: evidence requires object.prior in {tikhonov, tikhonov_spectral, 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, 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", "sure"): + raise ValueError( + "weight: auto requires solver.calibration.score in {gcv, evidence, sure} " + "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 From 0358e21dc1bf54af40a561df489399bedd73c327 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:06:02 -0700 Subject: [PATCH 03/24] feat(cli): migrate-config verb and version-aware config loading Add 'waveorder migrate-config' (alias mc) to convert 3.x configurations to 4.0, and route every verb's config parsing through the version-aware loader: version: 4 files run through the 4.0 schema, 3.x files run as before while printing their 4.0 equivalent once with a deprecation warning. End-to-end test asserts a migrated configuration reproduces the 3.x reconstruction bit for bit. bench.py forwards the expanded suite's per-case CLI selection when present. --- tests/config/test_migrate_cli.py | 110 ++++++++++++++++++ .../cli/apply_inverse_transfer_function.py | 14 +-- waveorder/cli/bench.py | 5 + waveorder/cli/compute_transfer_function.py | 5 +- waveorder/cli/main.py | 2 + waveorder/cli/migrate_config.py | 66 +++++++++++ waveorder/cli/reconstruct.py | 5 +- waveorder/cli/simulate.py | 5 +- 8 files changed, 195 insertions(+), 17 deletions(-) create mode 100644 tests/config/test_migrate_cli.py create mode 100644 waveorder/cli/migrate_config.py diff --git a/tests/config/test_migrate_cli.py b/tests/config/test_migrate_cli.py new file mode 100644 index 00000000..b375b1ad --- /dev/null +++ b/tests/config/test_migrate_cli.py @@ -0,0 +1,110 @@ +"""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 waveorder.cli.main import app +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, + }, + }, +} + + +@pytest.fixture +def input_zarr(tmp_path): + """A tiny single-position fluorescence dataset.""" + input_path = tmp_path / "input.zarr" + dataset = open_ome_zarr(input_path, layout="hcs", mode="w", channel_names=["GFP"]) + position = dataset.create_position("0", "0", "0") + rng = np.random.default_rng(0) + position.create_image( + "0", + rng.random((2, 1, 8, 16, 16)).astype(np.float32), + transform=[TransformationMeta(type="scale", scale=[1, 1, 0.4, 0.2, 0.2])], + ) + dataset.close() + return input_path + + +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) + + +def test_3x_config_warns_and_prints_v4(tmp_path, input_zarr): + from waveorder.config import loader + + 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/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 03a541c2..837fa6fe 100644 --- a/waveorder/cli/bench.py +++ b/waveorder/cli/bench.py @@ -157,6 +157,11 @@ def run( save_all=save_all, reference=case.reference, crop=case.crop, + # Forward-compatible with the expanded suite + # (bench/expand-hpc-cases), whose cases may select the + # tile-stitch CLI; absent there, run_hpc_case defaults + # to "rec". + **({"cli": case.cli} if getattr(case, "cli", None) is not None else {}), ) results[case_name] = metrics 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..87ada4a8 100644 --- a/waveorder/cli/main.py +++ b/waveorder/cli/main.py @@ -21,6 +21,7 @@ from waveorder.cli.bench import benchmark 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 @@ -94,6 +95,7 @@ def _main( "apply-inverse-transfer-function": (_apply_inverse_transfer_function_cli, "apply-inv-tf"), "tile-stitch": (_tile_stitch_cli, "ts"), "simulate": (_simulate_cli, "sim"), + "migrate-config": (_migrate_config_cli, "mc"), "view": (_view_cli, "v"), "interactive": (gui, "gui"), } 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..e048fc75 100644 --- a/waveorder/cli/simulate.py +++ b/waveorder/cli/simulate.py @@ -31,8 +31,7 @@ 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 def _write_czyx(czyx: xr.DataArray, path: Path): """Write a CZYX xr.DataArray to an HCS OME-Zarr.""" @@ -59,7 +58,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 From 4ac63eda89e3b6a898171f1cdd296bb09b9c6c65 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:13:09 -0700 Subject: [PATCH 04/24] feat(cli): calibrate-tf verb driving the per-tile Zernike recovery Add 'wo calibrate-tf' (design section 3): free pupil_aberrations terms in a 4.0 configuration are estimated with the existing patchwise Zernike recovery (shift_invariant = one full-FOV tile, design case 5; tiled uses the configured grid, case 17; score: sharpness with a single gradient stage). Results land in tf.zarr with parameters, parameters_covariance, parameter_names, tile_grid, and provenance; shift-invariant stores also carry the synthesized transfer function. Scores other than sharpness, numeric tolerances, and grid/random stages raise NotImplementedError pending design step 4. --- tests/calibrate/test_run_tfzarr.py | 186 ++++++++++++++++++++++ waveorder/calibrate/__init__.py | 18 +++ waveorder/calibrate/run.py | 244 +++++++++++++++++++++++++++++ waveorder/calibrate/tfzarr.py | 140 +++++++++++++++++ waveorder/cli/calibrate_tf.py | 158 +++++++++++++++++++ waveorder/cli/main.py | 3 + 6 files changed, 749 insertions(+) create mode 100644 tests/calibrate/test_run_tfzarr.py create mode 100644 waveorder/calibrate/run.py create mode 100644 waveorder/calibrate/tfzarr.py create mode 100644 waveorder/cli/calibrate_tf.py 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/waveorder/calibrate/__init__.py b/waveorder/calibrate/__init__.py index 38cce3c4..f470c131 100644 --- a/waveorder/calibrate/__init__.py +++ b/waveorder/calibrate/__init__.py @@ -10,6 +10,12 @@ related calibration scores. """ +from waveorder.calibrate.run import ( + CalibrationResult, + build_recovery_settings, + run_zernike_calibration, + zernike_calibration_requested, +) from waveorder.calibrate.scores import ( calibration_penalty, whittle_evidence, @@ -18,6 +24,11 @@ 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, @@ -28,7 +39,14 @@ __all__ = [ "NAMED_TERMS", + "CalibrationResult", + "TransferFunctionStore", + "build_recovery_settings", "calibration_penalty", + "read_tf_zarr", + "run_zernike_calibration", + "write_tf_zarr", + "zernike_calibration_requested", "fluorescence_otf", "fluorescence_psf", "name_to_noll", diff --git a/waveorder/calibrate/run.py b/waveorder/calibrate/run.py new file mode 100644 index 00000000..ce528239 --- /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} (evidence/gcv/sure land with design step 4; use sharpness)") + + 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/tfzarr.py b/waveorder/calibrate/tfzarr.py new file mode 100644 index 00000000..bfa87278 --- /dev/null +++ b/waveorder/calibrate/tfzarr.py @@ -0,0 +1,140 @@ +"""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, +) -> 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). + """ + 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 + + 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/cli/calibrate_tf.py b/waveorder/cli/calibrate_tf.py new file mode 100644 index 00000000..8aba04b9 --- /dev/null +++ b/waveorder/cli/calibrate_tf.py @@ -0,0 +1,158 @@ +from pathlib import Path + +from waveorder.cli.parsing import ( + ConfigFilepath, + InputPositionDirpaths, + OutputDirpath, +) + + +def calibrate_transfer_function_cli( + input_position_dirpath: Path, + config_filepath: Path, + output_dirpath: Path, +) -> 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.run import ( + run_zernike_calibration, + zernike_calibration_requested, + ) + from waveorder.calibrate.synthesis import fluorescence_otf + from waveorder.calibrate.tfzarr import 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 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) + finally: + input_dataset.close() + + result = run_zernike_calibration(cfg, zyx_data) + + # 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 = result.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(result.parameter_names, result.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": "refined", + "sources": [str(input_position_dirpath)], + "config_hash": hashlib.sha256(config_text.encode()).hexdigest(), + "score": "sharpness", + "time_index": int(t_index), + } + write_tf_zarr( + output_dirpath, + parameters=result.parameters, + parameter_names=result.parameter_names, + tile_grid=result.tile_grid, + provenance=provenance, + transfer_function=transfer_function, + ) + + echo_headline( + f"Estimated {len(result.parameter_names)} parameters on {n_tiles} tile(s); " + f"final loss {result.loss_history[-1]:.4g}" + ) + if n_tiles == 1: + for name, value in zip(result.parameter_names, result.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, +) -> 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`. + + \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) diff --git a/waveorder/cli/main.py b/waveorder/cli/main.py index 87ada4a8..66431c2c 100644 --- a/waveorder/cli/main.py +++ b/waveorder/cli/main.py @@ -19,6 +19,7 @@ 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 @@ -93,6 +94,7 @@ 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"), @@ -106,6 +108,7 @@ def _main( "reconstruct", "compute-transfer-function", "apply-inverse-transfer-function", + "calibrate-transfer-function", "tile-stitch", }: command_options["cls"] = _InputPathsCommand From 030acd94f8fa3661a1ef7d57ca91633deb5d1109 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:13:32 -0700 Subject: [PATCH 05/24] docs(config): point pupil-aberration configs at calibrate-tf --- waveorder/config/migrate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/waveorder/config/migrate.py b/waveorder/config/migrate.py index fff91060..63824a65 100644 --- a/waveorder/config/migrate.py +++ b/waveorder/config/migrate.py @@ -235,7 +235,7 @@ def unsupported(reason: str) -> NotRunnableInLegacyError: raise unsupported("imaging_system.field_dependence: tiled") detection = cfg.imaging_system.detection if detection.pupil_aberrations is not None: - raise unsupported("imaging_system.detection.pupil_aberrations") + 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 ): From 05e37b7cb3f4eff0cd53e5d12ff7622a8c611c90 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:21:33 -0700 Subject: [PATCH 06/24] feat(calibrate): evidence and gcv scores with grid/random/gradient search Add the MAPk search of design step 4 for shift-invariant fluorescence: score: evidence (exact Whittle marginal likelihood) or gcv (closed form for the Wiener filter), the calibration penalty from numeric tolerances, grid / random / gradient stages chained per the config, bounds clamping, Laplace parameters_covariance, score_trace and score_grid outputs, and noise.gaussian.sigma: estimate / object.weight: auto via joint maximization with a coarse (sigma, lambda) profile before each gradient stage. calibrate-tf dispatches on score; sharpness keeps the existing recovery path. --- tests/calibrate/test_evidence_search.py | 135 ++++++++ waveorder/calibrate/__init__.py | 8 + waveorder/calibrate/evidence_search.py | 400 ++++++++++++++++++++++++ waveorder/calibrate/tfzarr.py | 9 + waveorder/cli/calibrate_tf.py | 56 +++- 5 files changed, 598 insertions(+), 10 deletions(-) create mode 100644 tests/calibrate/test_evidence_search.py create mode 100644 waveorder/calibrate/evidence_search.py diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py new file mode 100644 index 00000000..7325d918 --- /dev/null +++ b/tests/calibrate/test_evidence_search.py @@ -0,0 +1,135 @@ +"""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.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 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) diff --git a/waveorder/calibrate/__init__.py b/waveorder/calibrate/__init__.py index f470c131..5f6227fd 100644 --- a/waveorder/calibrate/__init__.py +++ b/waveorder/calibrate/__init__.py @@ -10,6 +10,11 @@ related calibration scores. """ +from waveorder.calibrate.evidence_search import ( + EvidenceCalibrationResult, + gcv_score, + run_evidence_calibration, +) from waveorder.calibrate.run import ( CalibrationResult, build_recovery_settings, @@ -40,7 +45,10 @@ __all__ = [ "NAMED_TERMS", "CalibrationResult", + "EvidenceCalibrationResult", "TransferFunctionStore", + "gcv_score", + "run_evidence_calibration", "build_recovery_settings", "calibration_penalty", "read_tf_zarr", diff --git a/waveorder/calibrate/evidence_search.py b/waveorder/calibrate/evidence_search.py new file mode 100644 index 00000000..7dc82f3b --- /dev/null +++ b/waveorder/calibrate/evidence_search.py @@ -0,0 +1,400 @@ +"""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.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]) -> list[tuple[float, float]]: + """Grid range per parameter: bounds, else nominal +/- 2 tolerance.""" + ranges = [] + for spec in specs: + if spec.bounds is not None: + ranges.append(spec.bounds) + elif spec.tolerance != "free": + ranges.append((spec.nominal - 2.0 * spec.tolerance, spec.nominal + 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, +) -> 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. + + Returns + ------- + EvidenceCalibrationResult + """ + 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) + + 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 spec.tolerance != "free": + penalty = penalty + (values[i] - spec.nominal) ** 2 / (2.0 * spec.tolerance**2) + 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 = nominals.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) + 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) + 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. + 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/tfzarr.py b/waveorder/calibrate/tfzarr.py index bfa87278..125b816c 100644 --- a/waveorder/calibrate/tfzarr.py +++ b/waveorder/calibrate/tfzarr.py @@ -50,6 +50,7 @@ def write_tf_zarr( 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. @@ -72,6 +73,10 @@ def write_tf_zarr( 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: @@ -109,6 +114,10 @@ def write_tf_zarr( ) 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 {} diff --git a/waveorder/cli/calibrate_tf.py b/waveorder/cli/calibrate_tf.py index 8aba04b9..3266afdc 100644 --- a/waveorder/cli/calibrate_tf.py +++ b/waveorder/cli/calibrate_tf.py @@ -27,6 +27,7 @@ def calibrate_transfer_function_cli( import yaml from iohub.ngff import open_ome_zarr + from waveorder.calibrate.evidence_search import run_evidence_calibration from waveorder.calibrate.run import ( run_zernike_calibration, zernike_calibration_requested, @@ -80,17 +81,50 @@ def calibrate_transfer_function_cli( finally: input_dataset.close() - result = run_zernike_calibration(cfg, zyx_data) + calibration = cfg.solver.calibration + score = calibration.score if calibration is not None else "evidence" + extra_provenance = {} + score_arrays = {} + covariance = None + if score in ("evidence", "gcv"): + evidence_result = run_evidence_calibration(cfg, zyx_data) + 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 = result.parameters.shape[0] + 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(result.parameter_names, result.parameters[0])} + coefficients = {name_to_noll(name): value for name, value in zip(parameter_names, parameters[0])} transfer_function = ( fluorescence_otf( zyx_shape=zyx_data.shape, @@ -113,24 +147,26 @@ def calibrate_transfer_function_cli( "kind": "refined", "sources": [str(input_position_dirpath)], "config_hash": hashlib.sha256(config_text.encode()).hexdigest(), - "score": "sharpness", + "score": score, "time_index": int(t_index), + **extra_provenance, } write_tf_zarr( output_dirpath, - parameters=result.parameters, - parameter_names=result.parameter_names, - tile_grid=result.tile_grid, + 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(result.parameter_names)} parameters on {n_tiles} tile(s); " - f"final loss {result.loss_history[-1]:.4g}" + 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(result.parameter_names, result.parameters[0]): + for name, value in zip(parameter_names, parameters[0]): print(f" {name:16s} {value:+.4f} waves") echo_headline(f"Wrote {output_dirpath}") From 7626297368037f16d18fb319cc8bd5488f79639a Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:23:23 -0700 Subject: [PATCH 07/24] fix(calibrate): break the config <-> calibrate import cycle waveorder.config.schema imports the Zernike names from waveorder.calibrate, whose package init eagerly imported the config-dependent search modules, so a cold process importing waveorder.config first crashed with a circular import (pytest's import order hid it; every regression case failed). Load run/evidence_search lazily via PEP 562 and add fresh-interpreter import tests for both entry orders. --- tests/config/test_fresh_imports.py | 32 ++++++++++++++++++++ waveorder/calibrate/__init__.py | 48 +++++++++++++++++++----------- 2 files changed, 63 insertions(+), 17 deletions(-) create mode 100644 tests/config/test_fresh_imports.py 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/waveorder/calibrate/__init__.py b/waveorder/calibrate/__init__.py index 5f6227fd..e441e114 100644 --- a/waveorder/calibrate/__init__.py +++ b/waveorder/calibrate/__init__.py @@ -8,19 +8,15 @@ 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). """ -from waveorder.calibrate.evidence_search import ( - EvidenceCalibrationResult, - gcv_score, - run_evidence_calibration, -) -from waveorder.calibrate.run import ( - CalibrationResult, - build_recovery_settings, - run_zernike_calibration, - zernike_calibration_requested, -) +import importlib + from waveorder.calibrate.scores import ( calibration_penalty, whittle_evidence, @@ -42,24 +38,42 @@ 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", - "gcv_score", - "run_evidence_calibration", "build_recovery_settings", "calibration_penalty", - "read_tf_zarr", - "run_zernike_calibration", - "write_tf_zarr", - "zernike_calibration_requested", "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) From dd6eb3e8e04f052a968728501fcaf928b20db363 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:34:44 -0700 Subject: [PATCH 08/24] feat(calibrate): beads prior and tf.zarr covariance round-trip Add prior: {beads: {diameter}} (design example 4.3): sparse spheres at unknown subpixel positions fit jointly with the free pupil-aberration terms by MAP in Fourier space (analytic sphere form factor, phase-ramp positions), with local-maximum initialization and Laplace covariance. The known object breaks the global sign-flip twin, so bead calibration recovers signed coefficients. calibrate-tf now accepts -t (a tf.zarr whose parameters move the prior center and whose non-zero covariance replaces the configured tolerances as a correlated Gaussian prior; design section 3) and --init (start point only). Zernike docstrings now carry a prominent waves <-> radians conversion warning (cross-track currency is radians = 2 pi x waves). --- tests/calibrate/test_beads.py | 121 ++++++++++ tests/calibrate/test_evidence_search.py | 66 ++++++ waveorder/calibrate/beads.py | 292 ++++++++++++++++++++++++ waveorder/calibrate/evidence_search.py | 56 ++++- waveorder/calibrate/synthesis.py | 7 +- waveorder/calibrate/zernike_names.py | 6 + waveorder/cli/calibrate_tf.py | 78 ++++++- 7 files changed, 616 insertions(+), 10 deletions(-) create mode 100644 tests/calibrate/test_beads.py create mode 100644 waveorder/calibrate/beads.py diff --git a/tests/calibrate/test_beads.py b/tests/calibrate/test_beads.py new file mode 100644 index 00000000..95d2d723 --- /dev/null +++ b/tests/calibrate/test_beads.py @@ -0,0 +1,121 @@ +"""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 diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py index 7325d918..15f36071 100644 --- a/tests/calibrate/test_evidence_search.py +++ b/tests/calibrate/test_evidence_search.py @@ -123,6 +123,72 @@ def test_sigma_estimate_and_weight_auto(synthetic_data): assert np.isfinite(result.weight) and result.weight > 0 +def _store(values: dict[str, float], variances: dict[str, float] | None = None): + from waveorder.calibrate.tfzarr import TransferFunctionStore + + 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.""" + from waveorder.calibrate.tfzarr import TransferFunctionStore + + 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)) diff --git a/waveorder/calibrate/beads.py b/waveorder/calibrate/beads.py new file mode 100644 index 00000000..0021ad89 --- /dev/null +++ b/waveorder/calibrate/beads.py @@ -0,0 +1,292 @@ +"""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 +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, +) -> 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 * MAD``. + + Returns + ------- + Tensor + ``[N, 3]`` integer ``(z, y, x)`` peak positions. + """ + median = zyx_data.median() + mad = (zyx_data - median).abs().median().clamp(min=1e-12) + threshold = median + snr_threshold * 1.4826 * mad + + kernel = 2 * (min_separation // 2) + 1 + pooled = functional.max_pool3d(zyx_data[None, None], kernel_size=kernel, stride=1, padding=kernel // 2)[0, 0] + peaks = (zyx_data == pooled) & (zyx_data > threshold) + positions = peaks.nonzero() + if positions.shape[0] == 0: + raise ValueError( + f"no beads found above median + {snr_threshold} * MAD; check the channel, or lower the threshold" + ) + 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 = 50, +) -> 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 + At most this many beads are fit. + + 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. + peak_positions = detect_beads(data, max_beads=max_beads) + 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 index 7dc82f3b..837ae0fe 100644 --- a/waveorder/calibrate/evidence_search.py +++ b/waveorder/calibrate/evidence_search.py @@ -26,6 +26,7 @@ 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, @@ -164,6 +165,8 @@ 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. @@ -177,6 +180,16 @@ def run_evidence_calibration( 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 ------- @@ -196,6 +209,40 @@ def run_evidence_calibration( 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): + prior_precision = torch.linalg.inv(torch.as_tensor(block, dtype=torch.float64)) + 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") @@ -254,8 +301,13 @@ def total_score(values: torch.Tensor, log_sigma: torch.Tensor, log_weight: torch 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] - spec.nominal) ** 2 / (2.0 * spec.tolerance**2) + 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: @@ -265,7 +317,7 @@ def clamp(values: torch.Tensor) -> torch.Tensor: clamped[i] = clamped[i].clamp(spec.bounds[0], spec.bounds[1]) return clamped - current = nominals.clone() + 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] = [] diff --git a/waveorder/calibrate/synthesis.py b/waveorder/calibrate/synthesis.py index 5021ddf1..8491c3d3 100644 --- a/waveorder/calibrate/synthesis.py +++ b/waveorder/calibrate/synthesis.py @@ -6,8 +6,11 @@ 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`). This mirrors +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 diff --git a/waveorder/calibrate/zernike_names.py b/waveorder/calibrate/zernike_names.py index 40f53f1a..de8795cb 100644 --- a/waveorder/calibrate/zernike_names.py +++ b/waveorder/calibrate/zernike_names.py @@ -7,6 +7,12 @@ 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 diff --git a/waveorder/cli/calibrate_tf.py b/waveorder/cli/calibrate_tf.py index 3266afdc..5f85eccb 100644 --- a/waveorder/cli/calibrate_tf.py +++ b/waveorder/cli/calibrate_tf.py @@ -1,4 +1,7 @@ from pathlib import Path +from typing import Annotated, Optional + +import typer from waveorder.cli.parsing import ( ConfigFilepath, @@ -6,11 +9,31 @@ 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. @@ -27,19 +50,20 @@ def calibrate_transfer_function_cli( 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 write_tf_zarr + 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 free_parameters + from waveorder.config.schema import PriorBeads, free_parameters config_text = Path(config_filepath).read_text() raw = yaml.safe_load(config_text) @@ -86,8 +110,37 @@ def calibrate_transfer_function_cli( extra_provenance = {} score_arrays = {} covariance = None - if score in ("evidence", "gcv"): - evidence_result = run_evidence_calibration(cfg, zyx_data) + 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 @@ -144,7 +197,7 @@ def calibrate_transfer_function_cli( ) provenance = { - "kind": "refined", + "kind": provenance_kind, "sources": [str(input_position_dirpath)], "config_hash": hashlib.sha256(config_text.encode()).hexdigest(), "score": score, @@ -175,6 +228,8 @@ 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. @@ -183,6 +238,11 @@ def _calibrate_tf_cli( 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 @@ -191,4 +251,10 @@ def _calibrate_tf_cli( 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) + calibrate_transfer_function_cli( + input_position_dirpaths[0], + config_filepath, + output_dirpath, + prior_tf_dirpath=prior_tf_dirpath, + init_tf_dirpath=init_tf_dirpath, + ) From ba87d23ceed48b6cff0a1da6a85a3fe4d6cbc8c8 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 17:40:59 -0700 Subject: [PATCH 09/24] feat(cli): 4.0 simulate verb (design section 5, minimal) A version: 4 configuration with a simulate: {seed, shape} block describes the true imaging system and writes data.zarr, object.zarr and tf.zarr (drawn parameters, zero covariance, true transfer function, provenance: simulated) into the output directory. Implemented: fluorescence contrast, gaussian_spectrum object source, gaussian (sigma | snr_db) and poisson (photons) noise; a numeric tolerance draws the truth from its prior under the seed. 3.x configurations (and 4.0 without a simulate block) keep the existing phantom simulation. Closure test: evidence calibration on simulated data recovers the drawn coefficient. --- tests/config/test_simulate_v4.py | 115 +++++++++++++++++ waveorder/cli/simulate.py | 22 +++- waveorder/config/schema.py | 45 ++++++- waveorder/config/simulate.py | 205 +++++++++++++++++++++++++++++++ 4 files changed, 383 insertions(+), 4 deletions(-) create mode 100644 tests/config/test_simulate_v4.py create mode 100644 waveorder/config/simulate.py diff --git a/tests/config/test_simulate_v4.py b/tests/config/test_simulate_v4.py new file mode 100644 index 00000000..625c93d4 --- /dev/null +++ b/tests/config/test_simulate_v4.py @@ -0,0 +1,115 @@ +"""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.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]) + + from waveorder.calibrate.evidence_search import run_evidence_calibration + + 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/waveorder/cli/simulate.py b/waveorder/cli/simulate.py index e048fc75..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,7 +37,19 @@ def _simulate_cli(config_filepath: ConfigFilepath, output_dirpath: OutputDirpath phase, ) from waveorder.cli.printing import echo_headline - from waveorder.config.loader import load_reconstruction_settings + 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.""" diff --git a/waveorder/config/schema.py b/waveorder/config/schema.py index ad278c1d..66aab331 100644 --- a/waveorder/config/schema.py +++ b/waveorder/config/schema.py @@ -117,7 +117,10 @@ def _normalize_yx(cls, v): class DataSettings(V4BaseModel): """The ``data`` block: which data *g* is reconstructed and its sampling.""" - input_channel_names: list[str] = Field(description="channel names in the input zarr") + 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" ) @@ -324,8 +327,37 @@ 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*.""" + """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" @@ -333,6 +365,9 @@ class ObjectSettings(V4BaseModel): 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 ----------------------------------------------------------------- @@ -344,11 +379,17 @@ class DarkFrames(V4BaseModel): 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): diff --git a/waveorder/config/simulate.py b/waveorder/config/simulate.py new file mode 100644 index 00000000..98ba57ec --- /dev/null +++ b/waveorder/config/simulate.py @@ -0,0 +1,205 @@ +"""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 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.""" + from iohub.ngff import open_ome_zarr + from iohub.ngff.models import TransformationMeta + + 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 From 7c15216d2aa3ee855887e1a71a64dd50ab86b7d7 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Fri, 30 Oct 2026 18:02:46 -0700 Subject: [PATCH 10/24] fix: address verifier review of PR #592 - wo reconstruct on a 4.0 config with free parameters but no solver.calibration block now fails loudly instead of silently running the 3.x sharpness optimizer (the 4.0 default score is evidence, which runs through calibrate-tf); an explicit sharpness block still selects the 3.x path. - Bead detection thresholds are configurable (object.prior.beads {snr_threshold, max_beads, min_separation}) and relax adaptively (with a warning) before failing with actionable advice. - Design 4.5 validations added: beads-requires-free-parameter at schema level, grid point-count warning (shorthand expanded to its Zernike terms), tiled all-free sharpness overfit warning. - -t prior: singular covariance blocks are rejected with a clear error; grid stage ranges center on the -t-moved prior center rather than the config nominal. - Documented: the evidence path drops the DC term (exclude_dc) and parameters_covariance is a profile (sigma, lambda fixed) curvature; the padding z default 0 deviation from the design text; fluorescence_psf focal-plane wording; calibrate-tf prints which time index and channel it uses; detect_version rejects empty YAML files with a clear message. --- tests/calibrate/test_beads.py | 25 +++++++++++++ tests/calibrate/test_evidence_search.py | 25 +++++++++++++ tests/config/test_migrate.py | 27 ++++++++++++++ tests/config/test_v4_schema.py | 25 +++++++++++++ waveorder/calibrate/beads.py | 47 ++++++++++++++++++++----- waveorder/calibrate/evidence_search.py | 44 ++++++++++++++++++----- waveorder/calibrate/synthesis.py | 9 ++--- waveorder/cli/calibrate_tf.py | 3 ++ waveorder/config/migrate.py | 12 +++++++ waveorder/config/schema.py | 44 ++++++++++++++++++++++- 10 files changed, 239 insertions(+), 22 deletions(-) diff --git a/tests/calibrate/test_beads.py b/tests/calibrate/test_beads.py index 95d2d723..c54c8d88 100644 --- a/tests/calibrate/test_beads.py +++ b/tests/calibrate/test_beads.py @@ -119,3 +119,28 @@ def test_bead_calibration_cli(tmp_path, bead_volume): 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 index 15f36071..0dc7a263 100644 --- a/tests/calibrate/test_evidence_search.py +++ b/tests/calibrate/test_evidence_search.py @@ -199,3 +199,28 @@ def test_gcv_score_matches_definition(): 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/config/test_migrate.py b/tests/config/test_migrate.py index 3e3a9bc1..cffe3b47 100644 --- a/tests/config/test_migrate.py +++ b/tests/config/test_migrate.py @@ -218,3 +218,30 @@ def test_auto_regularization_maps_to_weight_auto(): 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).""" + from waveorder.config import NotRunnableInLegacyError + + 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_v4_schema.py b/tests/config/test_v4_schema.py index 571f60f6..db2f7e13 100644 --- a/tests/config/test_v4_schema.py +++ b/tests/config/test_v4_schema.py @@ -203,3 +203,28 @@ def test_extra_fields_rejected(): 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/waveorder/calibrate/beads.py b/waveorder/calibrate/beads.py index 0021ad89..502bd431 100644 --- a/waveorder/calibrate/beads.py +++ b/waveorder/calibrate/beads.py @@ -22,6 +22,7 @@ from __future__ import annotations import math +import warnings from dataclasses import dataclass import numpy as np @@ -94,6 +95,7 @@ def detect_beads( 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. @@ -106,7 +108,10 @@ def detect_beads( min_separation : int Non-maximum suppression window in voxels. snr_threshold : float - Peaks must exceed ``median + snr_threshold * MAD``. + 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 ------- @@ -115,15 +120,32 @@ def detect_beads( """ median = zyx_data.median() mad = (zyx_data - median).abs().median().clamp(min=1e-12) - threshold = median + snr_threshold * 1.4826 * mad kernel = 2 * (min_separation // 2) + 1 pooled = functional.max_pool3d(zyx_data[None, None], kernel_size=kernel, stride=1, padding=kernel // 2)[0, 0] - peaks = (zyx_data == pooled) & (zyx_data > threshold) - positions = peaks.nonzero() + 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 + {snr_threshold} * MAD; check the channel, or lower the threshold" + 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] @@ -134,7 +156,7 @@ def run_bead_calibration( cfg: V4Config, zyx_data: np.ndarray | torch.Tensor, gradient_lr: float = 0.01, - max_beads: int = 50, + max_beads: int | None = None, ) -> BeadCalibrationResult: """Fit free imaging-system parameters to a bead volume. @@ -148,8 +170,9 @@ def run_bead_calibration( 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 - At most this many beads are fit. + max_beads : int, optional + At most this many beads are fit; None reads + ``object.prior.beads.max_beads``. Returns ------- @@ -192,7 +215,13 @@ def run_bead_calibration( form_factor = sphere_form_factor(zyx_shape, yx_pixel_size, z_pixel_size, cfg.object.prior.beads.diameter) # Initialization from the data. - peak_positions = detect_beads(data, max_beads=max_beads) + 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) diff --git a/waveorder/calibrate/evidence_search.py b/waveorder/calibrate/evidence_search.py index 837ae0fe..6ea2187a 100644 --- a/waveorder/calibrate/evidence_search.py +++ b/waveorder/calibrate/evidence_search.py @@ -148,14 +148,19 @@ def _fixed_coefficients(cfg: V4Config, estimated: list[ParameterSpec]) -> dict[i return fixed -def _stage_ranges(specs: list[ParameterSpec]) -> list[tuple[float, float]]: - """Grid range per parameter: bounds, else nominal +/- 2 tolerance.""" +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 in specs: + for spec, center in zip(specs, centers.tolist()): if spec.bounds is not None: ranges.append(spec.bounds) elif spec.tolerance != "free": - ranges.append((spec.nominal - 2.0 * spec.tolerance, spec.nominal + 2.0 * spec.tolerance)) + 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 @@ -194,6 +199,20 @@ def run_evidence_calibration( 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": @@ -230,7 +249,15 @@ def run_evidence_calibration( 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): - prior_precision = torch.linalg.inv(torch.as_tensor(block, dtype=torch.float64)) + 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) @@ -362,7 +389,7 @@ def profile_sigma_weight(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tens 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) + ranges = _stage_ranges(specs, nominals) points = stage.grid.points axes = [] for spec, (lo, hi) in zip(specs, ranges): @@ -381,7 +408,7 @@ def profile_sigma_weight(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tens score_grid = np.asarray(rows) elif isinstance(stage, SearchRandom): generator = torch.Generator().manual_seed(stage.random.seed) - ranges = _stage_ranges(specs) + 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): @@ -429,7 +456,8 @@ def profile_sigma_weight(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tens best_score = float(total_score(current, log_sigma, log_weight)) - # Laplace covariance: inverse Hessian of the total score at the optimum. + # 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) diff --git a/waveorder/calibrate/synthesis.py b/waveorder/calibrate/synthesis.py index 8491c3d3..868d9003 100644 --- a/waveorder/calibrate/synthesis.py +++ b/waveorder/calibrate/synthesis.py @@ -135,10 +135,11 @@ def fluorescence_psf( Returns ------- Tensor - Real PSF, shape ``(Z, Y, X)``, sum-normalized to 1, with the focal - plane at index ``Z // 2`` after ``fftshift`` along z (the array is - returned in FFT order along z, matching the shift-variant model: - the in-focus slice is at index 0). + 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:] diff --git a/waveorder/cli/calibrate_tf.py b/waveorder/cli/calibrate_tf.py index 5f85eccb..400051cd 100644 --- a/waveorder/cli/calibrate_tf.py +++ b/waveorder/cli/calibrate_tf.py @@ -102,6 +102,9 @@ def calibrate_transfer_function_cli( 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() diff --git a/waveorder/config/migrate.py b/waveorder/config/migrate.py index 63824a65..401e66f8 100644 --- a/waveorder/config/migrate.py +++ b/waveorder/config/migrate.py @@ -42,6 +42,8 @@ class NotRunnableInLegacyError(ValueError): 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)) @@ -233,6 +235,16 @@ def unsupported(reason: str) -> NotRunnableInLegacyError: 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')") diff --git a/waveorder/config/schema.py b/waveorder/config/schema.py index 66aab331..9e988454 100644 --- a/waveorder/config/schema.py +++ b/waveorder/config/schema.py @@ -41,7 +41,7 @@ ) from waveorder._pixel_size import YXPixelSize -from waveorder.calibrate.zernike_names import validate_zernike_name +from waveorder.calibrate.zernike_names import noll_indices_up_to_order, validate_zernike_name class V4BaseModel(BaseModel): @@ -280,6 +280,13 @@ class TikhonovSpectralPrior(V4BaseModel): 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): @@ -426,6 +433,10 @@ def kind(self) -> str: 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]" @@ -610,7 +621,21 @@ def _validation_rules(self): 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, (PriorTikhonovSpectral, PriorBeads) @@ -627,6 +652,23 @@ def _validation_rules(self): 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" From a49fac5eed67482d6593324f0c548874863bbc86 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 1 Nov 2026 18:31:48 -0800 Subject: [PATCH 11/24] feat(bench): --migrate-configs regression mode (design section 2) Each case's 3.x config is converted with 'wo migrate-config' and the reconstruction runs from the 4.0 file, so the standard comparison machinery asserts output equality against the stored 3.x results. Configs that intentionally stay 3.x (birefringence) run unmigrated and are flagged via metrics['config_migrated']. --- benchmarks/runner.py | 81 +++++++++++++++++++++++++++++++--- tests/test_benchmark_runner.py | 63 +++++++++++++++++++++++++- waveorder/cli/bench.py | 17 +++++++ 3 files changed, 155 insertions(+), 6 deletions(-) diff --git a/benchmarks/runner.py b/benchmarks/runner.py index f2249d87..1d78b466 100644 --- a/benchmarks/runner.py +++ b/benchmarks/runner.py @@ -41,6 +41,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, @@ -94,6 +95,47 @@ def _wo_rec_argv(position_path: Path, config_path: Path, recon_path: Path) -> li return ["wo", "rec", "-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"} @@ -217,11 +259,15 @@ def _run_and_measure( case_dir: Path, tt: TimingTree, ground_truth: torch.Tensor | None = None, + migrate_config: bool = False, ) -> dict: """Write config, shell out to ``wo rec``, 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_config.get(modality, {}).get("transfer_function", {}) @@ -229,8 +275,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) + argv = _wo_rec_argv(position_path, exec_config_path, recon_path) (case_dir / "cli_command.sh").write_text(f"#!/bin/bash\n{shlex.join(argv)}\n") with tt.time("reconstruct"): @@ -241,13 +295,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: @@ -386,6 +443,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``. @@ -404,6 +462,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 ------- @@ -473,6 +536,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: @@ -637,6 +701,7 @@ def run_hpc_case( save_all: bool = False, reference: dict[str, ReferenceBound] | None = None, crop: CropConfig | None = None, + migrate_config: bool = False, ) -> dict: """Run a single HPC benchmark case on existing data via ``wo rec``. @@ -655,11 +720,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 ------- @@ -681,6 +751,7 @@ def run_hpc_case( recon_config=recon_config, case_dir=case_dir, tt=tt, + migrate_config=migrate_config, ) ref_check = check_reference(case_dir, metrics, reference) if ref_check is not None: 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/cli/bench.py b/waveorder/cli/bench.py index 837fa6fe..4f08e718 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( @@ -162,6 +177,8 @@ def run( # tile-stitch CLI; absent there, run_hpc_case defaults # to "rec". **({"cli": case.cli} if getattr(case, "cli", None) is not None else {}), + # Forward-compatible with runners that predate the flag. + **({"migrate_config": True} if migrate_configs else {}), ) results[case_name] = metrics From f6a1eb916e93b38f261d469583bb50585d8ed5ac Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 1 Nov 2026 18:31:49 -0800 Subject: [PATCH 12/24] test(config): migrate-and-run equivalence for every repo config Parametrized end-to-end test: every 3.x configuration in docs/ and benchmarks/ is migrated and run on a small synthetic input; outputs must match the 3.x run array by array. A guard test fails if a new config directory appears without coverage. Shared path lists move to tests/config/repo_configs.py. --- tests/config/repo_configs.py | 33 ++++++++++++++++++ tests/config/test_migrate.py | 48 ++++++++++++++++---------- tests/config/test_migrate_cli.py | 58 +++++++++++++++++++++++++++----- 3 files changed, 113 insertions(+), 26 deletions(-) create mode 100644 tests/config/repo_configs.py diff --git a/tests/config/repo_configs.py b/tests/config/repo_configs.py new file mode 100644 index 00000000..fef8e577 --- /dev/null +++ b/tests/config/repo_configs.py @@ -0,0 +1,33 @@ +"""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 + +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] + +#: 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 "tile-stitch" in p.name] + +#: 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_migrate.py b/tests/config/test_migrate.py index cffe3b47..450dbb53 100644 --- a/tests/config/test_migrate.py +++ b/tests/config/test_migrate.py @@ -5,9 +5,18 @@ 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, @@ -15,28 +24,33 @@ ) from waveorder.config.schema import V4Config -REPO_ROOT = Path(__file__).parent.parent.parent - -# Every 3.x phase / fluorescence configuration in the repository. -MIGRATABLE_CONFIG_PATHS = sorted( - [ - *(REPO_ROOT / "docs/examples/cli/configs").glob("*.yml"), - *(REPO_ROOT / "benchmarks/configs").glob("*.yml"), - *(REPO_ROOT / "docs/examples/optimization").glob("*.yml"), - ] -) -MIGRATABLE_CONFIG_PATHS = [ - p for p in MIGRATABLE_CONFIG_PATHS if "birefringence" not in p.name and "tile-stitch" not in p.name -] - -BIREFRINGENCE_CONFIG_PATHS = sorted((REPO_ROOT / "docs/examples/cli/configs").glob("birefringence*.yml")) - 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 @@ -228,8 +242,6 @@ def test_detect_version_rejects_empty_file(): 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).""" - from waveorder.config import NotRunnableInLegacyError - raw = { "version": 4, "data": {"input_channel_names": ["BF"], "pixel_size": {"yx": 0.1, "z": 0.25}}, diff --git a/tests/config/test_migrate_cli.py b/tests/config/test_migrate_cli.py index b375b1ad..0dd6e6fb 100644 --- a/tests/config/test_migrate_cli.py +++ b/tests/config/test_migrate_cli.py @@ -10,7 +10,9 @@ 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 = { @@ -32,22 +34,42 @@ } -@pytest.fixture -def input_zarr(tmp_path): - """A tiny single-position fluorescence dataset.""" - input_path = tmp_path / "input.zarr" - dataset = open_ome_zarr(input_path, layout="hcs", mode="w", channel_names=["GFP"]) +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((2, 1, 8, 16, 16)).astype(np.float32), + 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)) @@ -95,9 +117,29 @@ def test_3x_and_migrated_v4_reconstructions_match(tmp_path, input_zarr): np.testing.assert_array_equal(recon_3x, recon_v4) -def test_3x_config_warns_and_prints_v4(tmp_path, input_zarr): - from waveorder.config import loader +@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)) From 796fbb03a9ddf5e0cd4bbaf101b9bb8d4ca72bde Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 1 Nov 2026 18:31:50 -0800 Subject: [PATCH 13/24] docs(config): state the 3.x sunset policy in the loader 3.x configs are accepted with a deprecation warning for two releases (4.0 and 4.1), then rejected with a pointer to migrate-config. --- waveorder/config/loader.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/waveorder/config/loader.py b/waveorder/config/loader.py index 9528d284..7c54d173 100644 --- a/waveorder/config/loader.py +++ b/waveorder/config/loader.py @@ -8,6 +8,11 @@ * 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 @@ -78,8 +83,9 @@ def _warn_once_about_3x(settings: ReconstructionSettings, config_path: str) -> N # 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 will be rejected in a " - f"future release. Convert it with: waveorder migrate-config {config_path}", + 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, ) From 12a9da703af940845c49654ac356d13fa6fd7a14 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 1 Nov 2026 18:31:51 -0800 Subject: [PATCH 14/24] style: move inline imports to module top --- tests/calibrate/test_evidence_search.py | 5 +---- tests/config/test_simulate_v4.py | 3 +-- waveorder/config/simulate.py | 5 ++--- 3 files changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py index 0dc7a263..3b599614 100644 --- a/tests/calibrate/test_evidence_search.py +++ b/tests/calibrate/test_evidence_search.py @@ -10,6 +10,7 @@ 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) @@ -124,8 +125,6 @@ def test_sigma_estimate_and_weight_auto(synthetic_data): def _store(values: dict[str, float], variances: dict[str, float] | None = None): - from waveorder.calibrate.tfzarr import TransferFunctionStore - names = list(values) covariance = np.zeros((1, len(names), len(names))) if variances is not None: @@ -167,8 +166,6 @@ def test_zero_covariance_moves_center_only(synthetic_data): def test_sequential_calibration(synthetic_data): """Design case 12: -t yesterday's tf_hat.zarr as today's prior.""" - from waveorder.calibrate.tfzarr import TransferFunctionStore - first = run_evidence_calibration(_config(), synthetic_data) store = TransferFunctionStore( parameters=first.parameters, diff --git a/tests/config/test_simulate_v4.py b/tests/config/test_simulate_v4.py index 625c93d4..5ff4b8a6 100644 --- a/tests/config/test_simulate_v4.py +++ b/tests/config/test_simulate_v4.py @@ -6,6 +6,7 @@ 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 @@ -97,8 +98,6 @@ def test_calibration_closure_on_simulated_data(tmp_path): with open_ome_zarr(tmp_path / "sim" / "data.zarr" / "0" / "0" / "0") as data: zyx = np.asarray(data.data[0, 0]) - from waveorder.calibrate.evidence_search import run_evidence_calibration - recon_raw = yaml.safe_load(yaml.safe_dump(SIM_RAW)) recon_raw.pop("simulate") recon_raw["object"] = {"prior": "tikhonov", "weight": 1e-3} diff --git a/waveorder/config/simulate.py b/waveorder/config/simulate.py index 98ba57ec..ffe82c96 100644 --- a/waveorder/config/simulate.py +++ b/waveorder/config/simulate.py @@ -19,6 +19,8 @@ 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 @@ -76,9 +78,6 @@ def _gaussian_spectrum_object( 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.""" - from iohub.ngff import open_ome_zarr - from iohub.ngff.models import TransformationMeta - dataset = open_ome_zarr(path, layout="hcs", mode="w", channel_names=[channel_name]) position = dataset.create_position("0", "0", "0") position.create_image( From 761d726a55198234449a309adc22a9d2587a24c2 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 1 Nov 2026 18:40:56 -0800 Subject: [PATCH 15/24] fix(bench): accept config_v4_optimized.yml in parameter reference checks In migrated-config mode the reconstruction runs from config_v4.yml, so the optimizer writes config_v4_optimized.yml; check_reference now reads either name. Caught by the migrated-mode regression run (case ops_pheno_center_opt, slurm 37801387). --- benchmarks/runner.py | 9 ++++++++- tests/test_benchmark_reference.py | 11 +++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/benchmarks/runner.py b/benchmarks/runner.py index 1d78b466..54d5e738 100644 --- a/benchmarks/runner.py +++ b/benchmarks/runner.py @@ -366,9 +366,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}' " 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"): From 150dc4deba1cf68337e303db86f12da084abcd25 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 8 Nov 2026 14:44:27 -0800 Subject: [PATCH 16/24] feat(4.0): nn.Module-per-distribution layer (ImagingSystem/ObjectPrior/Noise/ParameterPrior) Additive OO layer over the existing calibrate numerics for the g = Hf + n model: Sampling, ImagingSystem base + Fluorescence3D (wraps fluorescence_otf) and Phase3D (wraps phase_thick_3d), ParameterPrior (wraps calibration_penalty), ObjectPrior base + Tikhonov + TikhonovSpectral (+ FittedSpectrum TODO stub), Noise base + Gaussian (+ Poisson stub), and free functions evidence/calibrate/ reconstruct plus a build_system_from_config adapter. {nominal,tolerance} values become nn.Parameter, plain numbers become buffers, and dotted named_parameters names double as tf.zarr parameter_names. All numerics delegate to the tested functions, so the regression suite is unchanged. --- tests/calibrate/test_distributions.py | 360 ++++++++++ waveorder/distributions.py | 969 ++++++++++++++++++++++++++ 2 files changed, 1329 insertions(+) create mode 100644 tests/calibrate/test_distributions.py create mode 100644 waveorder/distributions.py diff --git a/tests/calibrate/test_distributions.py b/tests/calibrate/test_distributions.py new file mode 100644 index 00000000..949c9523 --- /dev/null +++ b/tests/calibrate/test_distributions.py @@ -0,0 +1,360 @@ +"""The nn.Module-per-distribution layer wraps the existing numerics exactly. + +These are the equivalence/behavior tests of the WaveOrder 4.0 distribution +layer (``waveorder.distributions``). Each one pins the OO layer to the tested +function API it wraps, so the regression suite and the new layer cannot drift. +""" + +import numpy as np +import torch +import yaml + +from waveorder import distributions as D +from waveorder.calibrate.scores import calibration_penalty, whittle_evidence +from waveorder.calibrate.synthesis import fluorescence_otf +from waveorder.config.schema import V4Config +from waveorder.models import phase_thick_3d +from waveorder.reconstruct import tikhonov_regularized_inverse_filter + +ZYX = (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, +) + + +def _fluorescence_system(coefficients=None, normalize="peak"): + return D.Fluorescence3D( + wavelength_emission=OPTICS["wavelength_emission"], + index_of_refraction_media=OPTICS["index_of_refraction_media"], + numerical_aperture_detection=OPTICS["numerical_aperture_detection"], + zernike_coefficients=coefficients or {}, + normalize=normalize, + ) + + +def _sampling(shape=ZYX, z_padding=0): + return D.Sampling( + zyx_shape=shape, + yx_pixel_size=OPTICS["yx_pixel_size"], + z_pixel_size=OPTICS["z_pixel_size"], + z_padding=z_padding, + ) + + +def _simulate(otf_truth, noise_sigma, regularization_strength, seed=0): + """Draw g = H f + n with f from the Tikhonov prior S_k = sigma^2 / lambda.""" + generator = torch.Generator().manual_seed(seed) + prior_std = noise_sigma / regularization_strength**0.5 + f = torch.randn(otf_truth.shape, generator=generator) * prior_std + g = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf_truth)) + g = g + torch.randn(otf_truth.shape, generator=generator) * noise_sigma + return g.double() + + +# --- 1. evidence == whittle_evidence ---------------------------------------- # + + +def test_evidence_matches_whittle_evidence(): + s = _sampling() + coefficients = {"astigmatism_0": 0.12, "spherical": -0.05} + system = _fluorescence_system(coefficients) + sigma, weight = 0.1, 1e-2 + g = torch.randn(ZYX, dtype=torch.float64) + + value = D.evidence(g, system, D.Tikhonov(weight), D.Gaussian(sigma), s) + reference = whittle_evidence( + data=g, + otf=fluorescence_otf(zyx_shape=ZYX, zernike_coefficients=coefficients, **OPTICS), + noise_sigma=sigma, + regularization_strength=weight, + ) + assert abs(float(value) - float(reference)) < 1e-10 + + +def test_evidence_matches_whittle_evidence_exclude_dc(): + s = _sampling() + system = _fluorescence_system({"coma_x": 0.07}) + sigma, weight = 0.2, 5e-3 + g = torch.randn(ZYX, dtype=torch.float64) + + value = D.evidence(g, system, D.Tikhonov(weight), D.Gaussian(sigma), s, exclude_dc=True) + reference = whittle_evidence( + data=g, + otf=fluorescence_otf(zyx_shape=ZYX, zernike_coefficients={"coma_x": 0.07}, **OPTICS), + noise_sigma=sigma, + regularization_strength=weight, + exclude_dc=True, + ) + assert abs(float(value) - float(reference)) < 1e-10 + + +# --- 2. Fluorescence3D.transfer == fluorescence_otf ------------------------- # + + +def test_fluorescence_transfer_matches_otf(): + s = _sampling() + coefficients = {"astigmatism_0": 0.12, "spherical": -0.05, "defocus": 0.03} + system = _fluorescence_system(coefficients) + transfer = system.transfer(s) + reference = fluorescence_otf(zyx_shape=ZYX, zernike_coefficients=coefficients, **OPTICS) + assert torch.equal(transfer, reference) + + +def test_fluorescence_transfer_with_z_padding(): + s = _sampling(z_padding=2) + system = _fluorescence_system({"spherical": 0.1}) + transfer = system.transfer(s) + reference = fluorescence_otf( + zyx_shape=(ZYX[0] + 4, ZYX[1], ZYX[2]), + zernike_coefficients={"spherical": 0.1}, + **OPTICS, + ) + assert torch.equal(transfer, reference) + + +# --- 3. ParameterPrior.nll == calibration_penalty -------------------------- # + + +def test_parameter_prior_matches_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"} + prior = D.ParameterPrior(nominals, tolerances) + assert torch.equal(prior.nll(values), calibration_penalty(values, nominals, tolerances)) + + +def test_parameter_prior_from_system(): + system = _fluorescence_system({"spherical": D.free(0.0, 0.1), "astigmatism_0": 0.2}) + prior = D.ParameterPrior.from_system(system) + assert prior.nominals == {"pupil_aberrations.spherical": 0.0} + assert prior.tolerances == {"pupil_aberrations.spherical": 0.1} + values = {"pupil_aberrations.spherical": 0.05} + assert torch.equal(prior.nll(values), calibration_penalty(values, prior.nominals, prior.tolerances)) + + +# --- 4. reconstruct (Tikhonov) == closed-form Wiener ----------------------- # + + +def test_reconstruct_matches_tikhonov_closed_form_phase(): + s = _sampling(shape=(12, 24, 24), z_padding=0) + system = D.Phase3D( + wavelength_illumination=0.532, + index_of_refraction_media=1.3, + numerical_aperture_illumination=0.9, + numerical_aperture_detection=1.2, + ) + sigma, weight = 0.1, 1e-2 + noise = D.Gaussian(sigma) + obj = D.Tikhonov(weight) + g = torch.randn(12, 24, 24, dtype=torch.float64) + + estimate = D.reconstruct(g, system, obj, noise, s) + otf = system.transfer(s) + regularization = float(noise.variance() / obj.spectrum(s, noise)) + inverse_filter = tikhonov_regularized_inverse_filter(otf, regularization) + reference = torch.real(torch.fft.ifftn(torch.fft.fftn(g.to(torch.complex64)) * inverse_filter)) + assert torch.allclose(estimate, reference, atol=1e-6) + + +def test_reconstruct_tikhonov_regularization_is_sigma2_over_s_k(): + # S_k = sigma^2 / lambda, so sigma^2 / S_k = lambda exactly. + s = _sampling() + system = _fluorescence_system({"spherical": 0.1}) + sigma, weight = 0.3, 7e-3 + noise = D.Gaussian(sigma) + obj = D.Tikhonov(weight) + g = torch.randn(ZYX, dtype=torch.float64) + + estimate = D.reconstruct(g, system, obj, noise, s) + inverse_filter = tikhonov_regularized_inverse_filter(system.transfer(s), weight) + reference = torch.real(torch.fft.ifftn(torch.fft.fftn(g.to(torch.complex64)) * inverse_filter)) + assert torch.allclose(estimate, reference, atol=1e-6) + + +# --- 5. Phase3D.transfer == phase_thick_3d.calculate_transfer_function ----- # + + +def test_phase_transfer_matches_model(): + s = _sampling(shape=(8, 16, 16), z_padding=1) + system = D.Phase3D( + wavelength_illumination=0.532, + index_of_refraction_media=1.3, + numerical_aperture_illumination=0.9, + numerical_aperture_detection=1.2, + tilt_angle_zenith=0.05, + ) + real_tf, imag_tf = system.transfer_real_imag(s) + reference_real, reference_imag = phase_thick_3d.calculate_transfer_function( + (8, 16, 16), + 0.2, + 0.4, + 0.532, + 1, + 1.3, + 0.9, + 1.2, + tilt_angle_zenith=0.05, + ) + assert torch.equal(real_tf, reference_real) + assert torch.equal(imag_tf, reference_imag) + # transfer() returns the phase (real-potential) WOTF as the single H_k. + assert torch.equal(system.transfer(s), reference_real) + + +# --- 6. going blind flips an argument and calibrate recovers --------------- # + + +def test_going_blind_flips_one_argument_to_parameter(): + fixed = _fluorescence_system({"spherical": 0.0}) + assert fixed.free_parameter_names() == [] + blind = _fluorescence_system({"spherical": D.free(0.0, 0.5)}) + assert blind.free_parameter_names() == ["pupil_aberrations.spherical"] + # The free value is an nn.Parameter; a plain number is a buffer. + names = dict(blind.named_parameters()) + assert isinstance(names["pupil_aberrations.spherical"], torch.nn.Parameter) + assert "pupil_aberrations.spherical" not in dict(fixed.named_parameters()) + + +def test_calibrate_recovers_injected_aberration(): + # Only the magnitude of a lone astigmatism term is identifiable (the + # PSF-flip sign ambiguity), matching test_synthesis_scores.py. + s = _sampling() + sigma, weight = 0.05, 1e-2 + truth = 0.2 + otf_truth = fluorescence_otf(zyx_shape=ZYX, zernike_coefficients={"astigmatism_0": truth}, **OPTICS) + g = _simulate(otf_truth, sigma, weight, seed=3) + + system = _fluorescence_system({"astigmatism_0": D.free(0.0, 0.5)}) + prior = D.ParameterPrior.from_system(system) + result = D.calibrate( + g, system, prior, D.Tikhonov(weight), D.Gaussian(sigma), s, max_iter=200, lr=0.02, exclude_dc=True + ) + recovered = result.parameters["pupil_aberrations.astigmatism_0"] + assert abs(abs(recovered) - truth) < 0.11, result.parameters + assert result.covariance.shape == (1, 1) + assert np.isfinite(result.covariance).all() + + +# --- 7. state_dict round-trip, .to, dotted tf.zarr names ------------------- # + + +def test_state_dict_round_trip(): + s = _sampling() + system = _fluorescence_system({"spherical": D.free(0.17, 0.1), "astigmatism_0": 0.05}) + restored = _fluorescence_system({"spherical": D.free(0.0, 0.1), "astigmatism_0": 0.05}) + restored.load_state_dict(system.state_dict()) + assert torch.equal(system.transfer(s), restored.transfer(s)) + + +def test_to_device_cpu(): + system = _fluorescence_system({"spherical": D.free(0.1, 0.1)}) + moved = system.to("cpu") + assert moved is system + s = _sampling() + assert moved.transfer(s).shape == ZYX + + +def test_dotted_names_usable_as_tf_zarr_parameter_names(): + import tempfile + + from waveorder.calibrate.tfzarr import read_tf_zarr, write_tf_zarr + + system = _fluorescence_system({"astigmatism_0": D.free(0.1, 0.05), "spherical": D.free(-0.02, 0.05)}) + names = system.free_parameter_names() + assert names == ["pupil_aberrations.astigmatism_0", "pupil_aberrations.spherical"] + values = np.array([[float(p.detach()) for _, p in system.named_parameters()]]) + with tempfile.TemporaryDirectory() as tmp: + path = f"{tmp}/tf.zarr" + write_tf_zarr(path, values, names, tile_grid={"size_yx": [32, 32]}) + store = read_tf_zarr(path) + assert store.parameter_names == names + assert np.allclose(store.parameters, values) + + +# --- config adapter -------------------------------------------------------- # + + +def test_build_system_from_config_fluorescence(): + cfg = V4Config( + **yaml.safe_load( + """ +version: 4 +data: + 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: {nominal: 0.0, tolerance: 0.1} + astigmatism_0: 0.05 +object: {prior: tikhonov, weight: 0.01} +noise: {gaussian: {sigma: 0.1}} +""" + ) + ) + system, obj, noise, prior, sampling = D.build_system_from_config(cfg) + assert isinstance(system, D.Fluorescence3D) + assert system.free_parameter_names() == ["pupil_aberrations.spherical"] + assert isinstance(obj, D.Tikhonov) + assert isinstance(noise, D.Gaussian) + assert prior.tolerances == {"pupil_aberrations.spherical": 0.1} + assert float(obj.weight) == 0.01 + assert float(noise.sigma) == 0.1 + + +def test_build_system_from_config_phase_free_na(): + cfg = V4Config( + **yaml.safe_load( + """ +version: 4 +data: + pixel_size: {yx: 0.2, z: 0.4} +imaging_system: + contrast: phase + index_of_refraction_media: 1.3 + detection: {numerical_aperture: 1.2} + illumination: + wavelength: 0.532 + numerical_aperture: {nominal: 0.9, tolerance: 0.05} +object: {prior: tikhonov, weight: 0.001} +noise: {gaussian: {sigma: 1.0}} +""" + ) + ) + system, _, _, prior, _ = D.build_system_from_config(cfg) + assert isinstance(system, D.Phase3D) + assert system.free_parameter_names() == ["numerical_aperture_illumination"] + assert prior.tolerances == {"numerical_aperture_illumination": 0.05} + + +# --- stubs / Poisson / FittedSpectrum -------------------------------------- # + + +def test_fitted_spectrum_is_a_todo_stub(): + import pytest + + s = _sampling() + with pytest.raises(NotImplementedError): + D.FittedSpectrum().spectrum(s, D.Gaussian(0.1)) + + +def test_gaussian_noise_sample_and_variance(): + noise = D.Gaussian(0.5) + assert abs(float(noise.variance()) - 0.25) < 1e-12 + generator = torch.Generator().manual_seed(0) + draw = noise.sample(torch.zeros(5000, dtype=torch.float64), generator=generator) + assert abs(float(draw.std()) - 0.5) < 0.05 + + +def test_poisson_sample_mean(): + noise = D.Poisson(gain=1.0) + generator = torch.Generator().manual_seed(0) + draw = noise.sample(torch.full((5000,), 10.0, dtype=torch.float64), generator=generator) + assert abs(float(draw.mean()) - 10.0) < 0.3 diff --git a/waveorder/distributions.py b/waveorder/distributions.py new file mode 100644 index 00000000..4879adad --- /dev/null +++ b/waveorder/distributions.py @@ -0,0 +1,969 @@ +"""Object-oriented, one-``nn.Module``-per-distribution layer over the calibrate numerics. + +WaveOrder 4.0 writes the forward model as ``g = H f + n`` with three +distributions: + +* ``p(H)`` -- the imaging system and its prior, +* ``p(f)`` -- the object prior, +* ``p(n)`` -- the noise. + +This module gives each distribution a :class:`torch.nn.Module` (or a plain +class where there is nothing to carry), mirroring Sri's design. It is a thin, +*additive* layer: every numeric here calls the existing, tested functions in +:mod:`waveorder.calibrate` and :mod:`waveorder.models.phase_thick_3d`. Nothing +reimplements the math, so the regression suite stays byte-for-byte valid. + +Distribution -> module -> config map +------------------------------------ +====================== ========================= ================================ +distribution module config block +====================== ========================= ================================ +``p(H)`` imaging system :class:`Fluorescence3D` ``imaging_system`` (fluorescence) +``p(H)`` imaging system :class:`Phase3D` ``imaging_system`` (phase) +``p(H)`` parameter prior :class:`ParameterPrior` ``imaging_system`` tolerances +``p(f)`` object prior :class:`Tikhonov` ``object.prior: tikhonov`` +``p(f)`` object prior :class:`TikhonovSpectral` ``object.prior.tikhonov_spectral`` +``p(f)`` object prior :class:`FittedSpectrum` (TODO) ``object.prior`` fitted +``p(n)`` noise :class:`Gaussian` ``noise.gaussian`` +``p(n)`` noise :class:`Poisson` ``noise.poisson`` +====================== ========================= ================================ + +A configuration value written ``{nominal, tolerance}`` becomes an +:class:`torch.nn.Parameter` (a free parameter ``P`` of the design, estimated +from the data); a plain number becomes a registered buffer. The dotted names +from :meth:`torch.nn.Module.named_parameters` are the ``parameter_names`` of a +``tf.zarr`` store (:mod:`waveorder.calibrate.tfzarr`). +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import numpy as np +import torch +from torch import Tensor, nn + +from waveorder import sampling +from waveorder.calibrate.scores import calibration_penalty, whittle_evidence +from waveorder.calibrate.synthesis import fluorescence_otf +from waveorder.calibrate.zernike_names import name_to_noll, noll_to_name +from waveorder.models import phase_thick_3d +from waveorder.reconstruct import tikhonov_regularized_inverse_filter + +__all__ = [ + "Sampling", + "ImagingSystem", + "Fluorescence3D", + "Phase3D", + "ParameterPrior", + "ObjectPrior", + "Tikhonov", + "TikhonovSpectral", + "FittedSpectrum", + "Noise", + "Gaussian", + "Poisson", + "evidence", + "calibrate", + "reconstruct", + "CalibrationResult", +] + + +# --------------------------------------------------------------------------- # +# Sampling: grid / geometry # +# --------------------------------------------------------------------------- # + + +@dataclass +class Sampling: + """Grid geometry shared by an imaging system and an object prior. + + A thin container around the pixel sizes and shape; the Nyquist helpers + delegate to :mod:`waveorder.sampling` so there is a single source of truth + for the sample-spacing formulas. + + Parameters + ---------- + zyx_shape : tuple[int, int, int] + ``(Z, Y, X)`` data shape (before axial padding). + yx_pixel_size : float + Transverse pixel size, um. + z_pixel_size : float + Axial pixel size, um. + z_padding : int + Axial padding added on each side before building the transfer + function. + """ + + zyx_shape: tuple[int, int, int] + yx_pixel_size: float + z_pixel_size: float + z_padding: int = 0 + + @property + def shape(self) -> tuple[int, int, int]: + """The unpadded ``(Z, Y, X)`` shape.""" + return tuple(self.zyx_shape) + + @property + def padded_shape(self) -> tuple[int, int, int]: + """The ``(Z + 2 z_padding, Y, X)`` shape used for transfer functions.""" + z, y, x = self.zyx_shape + return (z + 2 * self.z_padding, y, x) + + @property + def yx(self) -> tuple[int, int]: + """The transverse ``(Y, X)`` shape.""" + return tuple(self.zyx_shape[1:]) + + def transverse_nyquist(self, na_illumination: float, na_detection: float, wavelength: float) -> float: + """Transverse Nyquist spacing (see :func:`waveorder.sampling.transverse_nyquist`).""" + return sampling.transverse_nyquist(wavelength, na_illumination, na_detection) + + def axial_nyquist(self, na_detection: float, wavelength: float, index_of_refraction_media: float) -> float: + """Axial Nyquist spacing (see :func:`waveorder.sampling.axial_nyquist`).""" + return sampling.axial_nyquist(wavelength, na_detection, index_of_refraction_media) + + +# --------------------------------------------------------------------------- # +# Parameter bookkeeping: {nominal, tolerance} -> Parameter, number -> buffer # +# --------------------------------------------------------------------------- # + + +@dataclass +class _Free: + """A ``{nominal, tolerance}`` configuration value: becomes an nn.Parameter.""" + + nominal: float + tolerance: float | str = "free" + + +def free(nominal: float, tolerance: float | str = "free") -> _Free: + """Mark a value free for calibration: ``free(nominal, tolerance)``. + + A :class:`Fluorescence3D` or :class:`Phase3D` built with a :class:`_Free` + value registers it as an :class:`torch.nn.Parameter` (appears in + ``named_parameters``); a plain number becomes a buffer. + + Parameters + ---------- + nominal : float + Prior center. + tolerance : float or "free" + Gaussian prior standard deviation, or ``"free"`` for an improper + (flat) prior. + + Returns + ------- + _Free + """ + return _Free(float(nominal), tolerance) + + +class ImagingSystem(nn.Module): + """Base class for ``p(H)``: an imaging system with free parameters. + + Subclasses own named optical parameters as buffers (fixed numbers) or + :class:`torch.nn.Parameter` (free ``{nominal, tolerance}`` values), build + the complex transfer function ``H_k`` with :meth:`transfer`, and apply it + with :meth:`forward`. The tolerances/nominals of the free parameters are + carried alongside so a :class:`ParameterPrior` can score them. + """ + + #: Optical-parameter dtype. ``float32`` matches the function API + #: (:func:`fluorescence_otf`, :func:`phase_thick_3d.calculate_transfer_function`) + #: exactly, so :meth:`transfer` reproduces it to the bit; the scores still + #: accumulate in ``float64`` internally. + param_dtype = torch.float32 + + def __init__(self) -> None: + super().__init__() + # name -> (nominal, tolerance) for the free parameters only. + self._prior: dict[str, tuple[float, float | str]] = {} + + def _register_value(self, name: str, value: float | _Free | Tensor) -> None: + """Register ``value`` as a Parameter (if free) or a buffer (if a number).""" + if isinstance(value, _Free): + tensor = torch.tensor(float(value.nominal), dtype=self.param_dtype) + self.register_parameter(name, nn.Parameter(tensor)) + self._prior[name] = (float(value.nominal), value.tolerance) + else: + tensor = torch.as_tensor(float(value), dtype=self.param_dtype) + self.register_buffer(name, tensor) + + def free_parameter_names(self) -> list[str]: + """Dotted names of the free parameters (``tf.zarr`` ``parameter_names``).""" + return [name for name, _ in self.named_parameters()] + + def parameter_values(self) -> dict[str, Tensor]: + """Current values of the free parameters by dotted name.""" + return {name: param for name, param in self.named_parameters()} + + def nominals(self) -> dict[str, float]: + """Prior centers of the free parameters by dotted name.""" + return {name: nominal for name, (nominal, _) in self._prior.items()} + + def tolerances(self) -> dict[str, float | str]: + """Prior tolerances of the free parameters by dotted name.""" + return {name: tolerance for name, (_, tolerance) in self._prior.items()} + + def transfer(self, s: Sampling) -> Tensor: # pragma: no cover - abstract + """Complex transfer function ``H_k`` in FFT order.""" + raise NotImplementedError + + def prepare(self, raw: Tensor, s: Sampling) -> Tensor: + """Turn a raw stack into the modeled data ``g`` (identity by default).""" + return torch.as_tensor(raw) + + def forward(self, f: Tensor, s: Sampling) -> Tensor: + """Apply the transfer function: ``g = H f`` (real part of the IFFT).""" + otf = self.transfer(s) + f_fft = torch.fft.fftn(torch.as_tensor(f, dtype=otf.dtype if otf.is_complex() else torch.complex64)) + return torch.real(torch.fft.ifftn(f_fft * otf)) + + +class Fluorescence3D(ImagingSystem): + """Widefield fluorescence ``p(H)``; wraps :func:`fluorescence_otf`. + + Owns the detection NA, named Zernike coefficients, focus offset, emission + wavelength, medium index and pixel sizes. Any value passed as + :func:`free` becomes a calibratable :class:`torch.nn.Parameter`; a plain + number becomes a buffer. :meth:`transfer` calls :func:`fluorescence_otf` + with the current values, so the OTF is identical to the function API. + + Parameters + ---------- + wavelength_emission : float + Emission wavelength, um. + index_of_refraction_media : float + Medium refractive index. + numerical_aperture_detection : float or _Free + Detection NA. + zernike_coefficients : mapping, optional + ``{name_or_noll: value}`` in waves; each value a number (buffer) or a + :func:`free` value (Parameter). The dotted parameter name is + ``pupil_aberrations.`` (e.g. + ``pupil_aberrations.astigmatism_0``). + z_focus_offset : float or _Free + Focus offset in slices. + normalize : {"peak", "none"} + OTF normalization, passed to :func:`fluorescence_otf`. + """ + + def __init__( + self, + wavelength_emission: float, + index_of_refraction_media: float, + numerical_aperture_detection: float | _Free, + zernike_coefficients: Mapping[str | int, float | _Free] | None = None, + z_focus_offset: float | _Free = 0.0, + normalize: str = "peak", + ) -> None: + super().__init__() + self.wavelength_emission = float(wavelength_emission) + self.index_of_refraction_media = float(index_of_refraction_media) + self.normalize = normalize + + self._register_value("numerical_aperture_detection", numerical_aperture_detection) + self._register_value("z_focus_offset", z_focus_offset) + + # Zernike coefficients live under a dotted "pupil_aberrations." prefix + # so their free names round-trip through tf.zarr parameter_names. + self._zernike_nolls: list[int] = [] + self.pupil_aberrations = nn.Module() + aberrations_prior: dict[str, tuple[float, float | str]] = {} + for key, value in (zernike_coefficients or {}).items(): + noll = key if isinstance(key, int) else name_to_noll(key) + canonical = noll_to_name(noll) + self._zernike_nolls.append(noll) + if isinstance(value, _Free): + tensor = torch.tensor(float(value.nominal), dtype=self.param_dtype) + self.pupil_aberrations.register_parameter(canonical, nn.Parameter(tensor)) + aberrations_prior[f"pupil_aberrations.{canonical}"] = (float(value.nominal), value.tolerance) + else: + tensor = torch.as_tensor(float(value), dtype=self.param_dtype) + self.pupil_aberrations.register_buffer(canonical, tensor) + self._prior.update(aberrations_prior) + + def _coefficients(self) -> dict[int, float | Tensor]: + """Current Zernike coefficients keyed by Noll index.""" + out: dict[int, float | Tensor] = {} + for noll in self._zernike_nolls: + canonical = noll_to_name(noll) + value = getattr(self.pupil_aberrations, canonical) + out[noll] = value + return out + + def transfer(self, s: Sampling) -> Tensor: + """Complex OTF from :func:`fluorescence_otf` at the current parameters.""" + return fluorescence_otf( + zyx_shape=s.padded_shape, + yx_pixel_size=s.yx_pixel_size, + z_pixel_size=s.z_pixel_size, + wavelength_emission=self.wavelength_emission, + index_of_refraction_media=self.index_of_refraction_media, + numerical_aperture_detection=self.numerical_aperture_detection, + zernike_coefficients=self._coefficients(), + z_focus_offset=self.z_focus_offset, + normalize=self.normalize, + ) + + +class Phase3D(ImagingSystem): + """Thick-phase ``p(H)``; wraps :func:`phase_thick_3d.calculate_transfer_function`. + + The weak-object model has two transfer functions: the real (phase) WOTF + and the imaginary (absorption) WOTF, both returned by WaveOrder as complex + tensors. :meth:`transfer_real_imag` exposes the pair directly for the + label-free reconstruction path. :meth:`transfer` returns the phase WOTF as + the single ``H_k`` (the usual pure-phase case; combine with the absorption + WOTF via an ``absorption_ratio`` at reconstruction time, as + :func:`phase_thick_3d.apply_inverse_transfer_function` does). + + Parameters + ---------- + wavelength_illumination : float + Illumination wavelength, um. + index_of_refraction_media : float + Medium refractive index. + numerical_aperture_illumination : float or _Free + Condenser NA. + numerical_aperture_detection : float or _Free + Detection NA. + tilt_angle_zenith, tilt_angle_azimuth : float or _Free + Illumination tilt angles, radians. + invert_phase_contrast : bool + Flip the defocus axis (matches 3.x). + pupil_steepness : float + Sigmoid pupil-edge steepness. + """ + + def __init__( + self, + wavelength_illumination: float, + index_of_refraction_media: float, + numerical_aperture_illumination: float | _Free = 0.9, + numerical_aperture_detection: float | _Free = 1.2, + tilt_angle_zenith: float | _Free = 0.0, + tilt_angle_azimuth: float | _Free = 0.0, + invert_phase_contrast: bool = False, + pupil_steepness: float = 1e4, + ) -> None: + super().__init__() + self.wavelength_illumination = float(wavelength_illumination) + self.index_of_refraction_media = float(index_of_refraction_media) + self.invert_phase_contrast = bool(invert_phase_contrast) + self.pupil_steepness = float(pupil_steepness) + + self._register_value("numerical_aperture_illumination", numerical_aperture_illumination) + self._register_value("numerical_aperture_detection", numerical_aperture_detection) + self._register_value("tilt_angle_zenith", tilt_angle_zenith) + self._register_value("tilt_angle_azimuth", tilt_angle_azimuth) + + def transfer_real_imag(self, s: Sampling) -> tuple[Tensor, Tensor]: + """Real and imaginary potential WOTFs from the phase model. + + The optical parameters are ``float32`` (``param_dtype``), matching the + model's internal cast, so the returned WOTFs reproduce the function + API exactly while preserving autograd through the free parameters. + """ + return phase_thick_3d.calculate_transfer_function( + zyx_shape=s.shape, + yx_pixel_size=s.yx_pixel_size, + z_pixel_size=s.z_pixel_size, + wavelength_illumination=self.wavelength_illumination, + z_padding=s.z_padding, + index_of_refraction_media=self.index_of_refraction_media, + numerical_aperture_illumination=self.numerical_aperture_illumination, + numerical_aperture_detection=self.numerical_aperture_detection, + invert_phase_contrast=self.invert_phase_contrast, + tilt_angle_zenith=self.tilt_angle_zenith, + tilt_angle_azimuth=self.tilt_angle_azimuth, + pupil_steepness=self.pupil_steepness, + ) + + def transfer(self, s: Sampling) -> Tensor: + """The phase (real-potential) WOTF ``H_k`` for the pure-phase case.""" + real_tf, _ = self.transfer_real_imag(s) + return real_tf + + +# --------------------------------------------------------------------------- # +# ParameterPrior: C(H) = -log p(H) # +# --------------------------------------------------------------------------- # + + +class ParameterPrior: + """The calibration penalty ``C(H) = -log p(H)``; wraps :func:`calibration_penalty`. + + Parameters + ---------- + nominals : dict[str, float] + Prior centers by parameter name. + tolerances : dict[str, float or "free"] + Gaussian prior standard deviations by parameter name (or ``"free"``). + """ + + def __init__(self, nominals: Mapping[str, float], tolerances: Mapping[str, float | str]) -> None: + self.nominals = dict(nominals) + self.tolerances = dict(tolerances) + + @classmethod + def from_system(cls, system: ImagingSystem) -> "ParameterPrior": + """Build the prior from an imaging system's free-parameter nominals/tolerances.""" + return cls(system.nominals(), system.tolerances()) + + def nll(self, params: Mapping[str, float | Tensor]) -> Tensor: + """Negative log prior ``C = sum_i (h_i - h_i0)^2 / (2 tau_i^2)``.""" + return calibration_penalty(dict(params), self.nominals, self.tolerances) + + def sample(self, generator: torch.Generator | None = None) -> dict[str, float]: + """Draw a parameter dict from the Gaussian prior (``free`` -> nominal).""" + out: dict[str, float] = {} + for name, nominal in self.nominals.items(): + tolerance = self.tolerances.get(name, "free") + if tolerance == "free": + out[name] = float(nominal) + else: + noise = torch.randn((), generator=generator, dtype=torch.float64) + out[name] = float(nominal) + float(tolerance) * float(noise) + return out + + +# --------------------------------------------------------------------------- # +# ObjectPrior: p(f) -> S_k # +# --------------------------------------------------------------------------- # + + +class ObjectPrior(nn.Module): + """Base class for ``p(f)``: returns the object power spectrum ``S_k``.""" + + def spectrum(self, s: Sampling, noise: "Noise") -> Tensor | None: # pragma: no cover - abstract + """Object power spectrum ``S_k`` (same shape as ``H_k``), or ``None``.""" + raise NotImplementedError + + +class Tikhonov(ObjectPrior): + """Flat Tikhonov prior ``S_k = sigma^2 / lambda`` (the current default). + + Parameters + ---------- + weight : float + Regularization weight ``lambda``. + """ + + def __init__(self, weight: float) -> None: + super().__init__() + self.register_buffer("weight", torch.as_tensor(float(weight), dtype=torch.float64)) + + def spectrum(self, s: Sampling, noise: "Noise") -> Tensor: + """``S_k = sigma^2 / lambda`` (a scalar; broadcast against ``H_k``).""" + return noise.variance() / self.weight + + +class TikhonovSpectral(ObjectPrior): + """Explicit per-frequency object spectrum ``S_k``. + + Parameters + ---------- + spectrum : Tensor + The object power spectrum ``S_k`` in FFT order, same shape as ``H_k``. + """ + + def __init__(self, spectrum: Tensor) -> None: + super().__init__() + self.register_buffer("spectrum_buffer", torch.as_tensor(spectrum, dtype=torch.float64)) + + def spectrum(self, s: Sampling, noise: "Noise") -> Tensor: + """Return the stored ``S_k``.""" + return self.spectrum_buffer + + +class FittedSpectrum(ObjectPrior): + """Isotropic shell-fit object spectrum ``S(|k|)`` estimated from the data. + + TODO (next pass): the winning iso-Whittle strategy of the 2026-10-30 + sprint. Fit an isotropic object power spectrum ``S(|k|)`` by a + weighted-least-squares shell fit of the data periodogram with ridge + shrinkage, then broadcast it onto the full ``H_k`` grid. See + ``winning_method_spec.md``. The interface is fixed: + + def spectrum(self, s: Sampling, noise: Noise) -> Tensor: + # 1. periodogram P_k = |FFT(g)|^2 (ortho) of the prepared data, + # 2. bin by |k| shells on the Sampling grid, + # 3. WLS fit S(|k|) = (P_shell - sigma^2) / |H_shell|^2 with + # ridge shrinkage toward a smooth power law, + # 4. scatter S(|k|) back to every frequency -> S_k. + + This flat-prior-beats-nothing placeholder is intentionally inert so that + the interface exists and the config builder can reference it before the + numeric fit lands. + """ + + def __init__(self) -> None: + super().__init__() + + def spectrum(self, s: Sampling, noise: "Noise") -> Tensor | None: + """Not implemented yet; see the class docstring for the planned fit.""" + raise NotImplementedError( + "FittedSpectrum (iso-Whittle shell fit) lands in the next pass; use Tikhonov or TikhonovSpectral for now." + ) + + +# --------------------------------------------------------------------------- # +# Noise: p(n) # +# --------------------------------------------------------------------------- # + + +class Noise(nn.Module): + """Base class for ``p(n)``.""" + + def variance(self) -> Tensor: # pragma: no cover - abstract + """Noise variance ``sigma^2``.""" + raise NotImplementedError + + def nll(self, g: Tensor, mean: Tensor) -> Tensor: # pragma: no cover - abstract + """Negative log likelihood of ``g`` given its mean (the data misfit ``D``).""" + raise NotImplementedError + + def sample(self, mean: Tensor, generator: torch.Generator | None = None) -> Tensor: # pragma: no cover + """Draw a noisy observation around ``mean``.""" + raise NotImplementedError + + +class Gaussian(Noise): + """Additive white Gaussian noise ``n ~ N(0, sigma^2 I)``. + + Parameters + ---------- + sigma : float + Noise standard deviation in data units. + """ + + def __init__(self, sigma: float) -> None: + super().__init__() + self.register_buffer("sigma", torch.as_tensor(float(sigma), dtype=torch.float64)) + + def variance(self) -> Tensor: + """``sigma^2``.""" + return self.sigma**2 + + def nll(self, g: Tensor, mean: Tensor) -> Tensor: + """Gaussian data misfit ``||g - mean||^2 / (2 sigma^2) + const``.""" + residual = torch.as_tensor(g, dtype=torch.float64) - torch.as_tensor(mean, dtype=torch.float64) + return (residual**2).sum() / (2.0 * self.variance()) + + def sample(self, mean: Tensor, generator: torch.Generator | None = None) -> Tensor: + """``mean + sigma * N(0, 1)``.""" + mean = torch.as_tensor(mean, dtype=torch.float64) + noise = torch.randn(mean.shape, generator=generator, dtype=torch.float64) + return mean + self.sigma * noise + + +class Poisson(Noise): + """Poisson (shot) noise with a linear gain. + + TODO (next pass): the full Poisson evidence / Richardson-Lucy path. For + now ``variance`` returns the gain so that a Gaussian-approximate evidence + can still be scored, and ``sample`` draws true Poisson counts for the + simulator. + + Parameters + ---------- + gain : float + Electrons (or ADU) per photon. + offset : float + Constant dark offset in data units. + """ + + def __init__(self, gain: float = 1.0, offset: float = 0.0) -> None: + super().__init__() + self.register_buffer("gain", torch.as_tensor(float(gain), dtype=torch.float64)) + self.register_buffer("offset", torch.as_tensor(float(offset), dtype=torch.float64)) + + def variance(self) -> Tensor: + """Gaussian-approximate variance (the gain); a true shot-noise map is per-voxel.""" + return self.gain + + def nll(self, g: Tensor, mean: Tensor) -> Tensor: + """Poisson negative log likelihood ``sum(mean - g log mean)`` (up to const).""" + g = torch.as_tensor(g, dtype=torch.float64) + rate = (torch.as_tensor(mean, dtype=torch.float64) - self.offset) / self.gain + rate = rate.clamp(min=1e-12) + return (rate - (g / self.gain) * torch.log(rate)).sum() + + def sample(self, mean: Tensor, generator: torch.Generator | None = None) -> Tensor: + """Draw Poisson counts ``offset + gain * Poisson((mean - offset) / gain)``.""" + rate = (torch.as_tensor(mean, dtype=torch.float64) - self.offset) / self.gain + counts = torch.poisson(rate.clamp(min=0.0), generator=generator) + return self.offset + self.gain * counts + + +# --------------------------------------------------------------------------- # +# evidence / calibrate / reconstruct # +# --------------------------------------------------------------------------- # + + +def evidence( + g: Tensor, + system: ImagingSystem, + obj: ObjectPrior, + noise: Noise, + s: Sampling, + *, + exclude_dc: bool = False, + data_fft: Tensor | None = None, +) -> Tensor: + """Negative log evidence ``-log p(g | H)`` of the imaging system. + + Exactly Sri's snippet: per frequency ``v_k = |H_k|^2 S_k + sigma^2`` and + the Whittle NLL. It calls :func:`whittle_evidence` with the module-derived + ``H_k``, ``S_k`` and ``sigma`` so the value matches the function API. + + Parameters + ---------- + g : Tensor + Observed stack ``(Z, Y, X)``. + system : ImagingSystem + ``p(H)``; supplies ``H_k`` via :meth:`ImagingSystem.transfer`. + obj : ObjectPrior + ``p(f)``; supplies ``S_k`` via :meth:`ObjectPrior.spectrum`. + noise : Noise + ``p(n)``; supplies ``sigma``. + s : Sampling + Grid geometry. + exclude_dc : bool + Drop the ``k = 0`` term (for real data with a large background). + data_fft : Tensor, optional + Precomputed ``fftn(g, norm="ortho")``. + + Returns + ------- + Tensor + Scalar negative log evidence. + """ + otf = system.transfer(s) + s_k = obj.spectrum(s, noise) + sigma = torch.sqrt(noise.variance()) + return whittle_evidence( + data=None if data_fft is not None else g, + data_fft=data_fft, + otf=otf, + noise_sigma=sigma, + object_spectrum=s_k, + exclude_dc=exclude_dc, + ) + + +def _laplace_covariance(objective, params: list[nn.Parameter]) -> Tensor: + """Inverse Hessian of ``objective`` w.r.t. ``params`` by double backprop. + + Parameters + ---------- + objective : callable + Scalar objective evaluated at the current parameter values. + params : list of nn.Parameter + The free parameters; the Hessian is taken over their flattened stack. + + Returns + ------- + Tensor + ``[P, P]`` Laplace covariance (inverse Hessian at the optimum). + """ + value = objective() + grads = torch.autograd.grad(value, params, create_graph=True) + flat_grad = torch.cat([g.reshape(-1) for g in grads]) + rows = [] + for i in range(flat_grad.numel()): + second = torch.autograd.grad(flat_grad[i], params, retain_graph=True) + rows.append(torch.cat([g.reshape(-1) for g in second])) + hessian = torch.stack(rows).to(torch.float64) + return torch.linalg.inv(hessian).detach() + + +@dataclass +class CalibrationResult: + """Output of :func:`calibrate`.""" + + parameters: dict[str, float] + parameter_names: list[str] + covariance: np.ndarray + score_trace: list[float] + best_score: float + + +def calibrate( + g: Tensor, + system: ImagingSystem, + prior: ParameterPrior, + obj: ObjectPrior, + noise: Noise, + s: Sampling, + *, + score=evidence, + max_iter: int = 200, + lr: float = 0.01, + exclude_dc: bool = False, +) -> CalibrationResult: + """Estimate ``H``'s free parameters by minimizing ``score + prior.nll``. + + The outer argmin over ``system``'s :meth:`torch.nn.Module.named_parameters` + (the design's free params ``P``), by Adam, with the Laplace covariance + (inverse Hessian of the total objective at the optimum). + + Parameters + ---------- + g : Tensor + Observed stack. + system : ImagingSystem + The imaging system; its free parameters are optimized in place. + prior : ParameterPrior + ``C(H)``, added to the data score. + obj : ObjectPrior + ``p(f)``. + noise : Noise + ``p(n)``. + s : Sampling + Grid geometry. + score : callable + Data score, ``score(g, system, obj, noise, s, exclude_dc=...)``; + defaults to :func:`evidence`. + max_iter : int + Adam iterations. + lr : float + Adam step size. + exclude_dc : bool + Passed to the score. + + Returns + ------- + CalibrationResult + """ + names = system.free_parameter_names() + params = [param for _, param in system.named_parameters()] + if not params: + raise ValueError("system has no free parameters to calibrate; mark values with free(...)") + + g = torch.as_tensor(g, dtype=torch.float64) + data_fft = torch.fft.fftn(g, norm="ortho") + + def objective() -> Tensor: + data_term = score(g, system, obj, noise, s, exclude_dc=exclude_dc, data_fft=data_fft) + return data_term + prior.nll(system.parameter_values()) + + optimizer = torch.optim.Adam(params, lr=lr) + score_trace: list[float] = [] + for _ in range(max_iter): + optimizer.zero_grad() + value = objective() + value.backward() + optimizer.step() + score_trace.append(float(value.detach())) + + with torch.no_grad(): + best_score = float(objective()) + estimated = {name: float(param.detach()) for name, param in zip(names, params)} + + # Laplace covariance: inverse Hessian of the total objective at the + # optimum. Build it by double backprop through the live parameters so the + # graph stays intact (a copy_/no_grad probe would detach it). + try: + covariance = _laplace_covariance(objective, params).numpy() + except RuntimeError: + covariance = np.full((len(params), len(params)), np.nan) + + return CalibrationResult( + parameters=estimated, + parameter_names=names, + covariance=covariance, + score_trace=score_trace, + best_score=best_score, + ) + + +def reconstruct( + g: Tensor, + system: ImagingSystem, + obj: ObjectPrior, + noise: Noise, + s: Sampling, +) -> Tensor: + """Inner argmin over ``f``: the closed-form Wiener estimate. + + For Gaussian noise and a Tikhonov object prior this is exactly the + Wiener filter ``f_hat = H* / (|H|^2 + lambda) . g``, identical to + :func:`waveorder.reconstruct.tikhonov_regularized_inverse_filter` applied + to ``g`` with ``lambda = sigma^2 / S_k``. + + Parameters + ---------- + g : Tensor + Observed stack ``(Z, Y, X)``. + system : ImagingSystem + Supplies ``H_k``. + obj : ObjectPrior + Supplies ``S_k`` (hence ``lambda = sigma^2 / S_k``). + noise : Noise + Supplies ``sigma^2``. + s : Sampling + Grid geometry. + + Returns + ------- + Tensor + Real-space object estimate ``f_hat``, same shape as ``g``. + + Raises + ------ + NotImplementedError + For a non-Gaussian noise or an improper (``None``) object spectrum; + the iterative / Richardson-Lucy path lands in a later pass. + """ + if not isinstance(noise, Gaussian): + raise NotImplementedError("reconstruct's closed form is Gaussian-only for now; RL path follows") + s_k = obj.spectrum(s, noise) + if s_k is None: + raise NotImplementedError("improper object prior (S_k=None) has no closed-form Wiener estimate") + + otf = system.transfer(s) + g = torch.as_tensor(g) + if g.shape[-3:] != otf.shape[-3:]: + raise ValueError( + f"g shape {tuple(g.shape)} does not match the transfer function {tuple(otf.shape)}; " + f"z-padded reconstruction is not handled by the closed form yet " + f"(use s.z_padding == 0, or pad g to the padded_shape)" + ) + regularization = noise.variance() / s_k + inverse_filter = tikhonov_regularized_inverse_filter(otf, regularization) + g = g.to(torch.complex64) + g_fft = torch.fft.fftn(g) + return torch.real(torch.fft.ifftn(g_fft * inverse_filter)) + + +# --------------------------------------------------------------------------- # +# Config adapter # +# --------------------------------------------------------------------------- # + + +def _value_or_free(value): + """Map a config float or ``FreeParameter`` to a plain number or :func:`free`.""" + from waveorder.config.schema import FreeParameter + + if isinstance(value, FreeParameter): + return free(value.nominal, value.tolerance) + return float(value) + + +def build_system_from_config(cfg): + """Build ``(system, obj, noise, prior, sampling)`` from a 4.0 config. + + A thin adapter from the ``imaging_system`` / ``object`` / ``noise`` / + ``solver`` blocks of a :class:`waveorder.config.schema.V4Config` to the + distribution modules of this file. A config value ``{nominal, tolerance}`` + (a :class:`waveorder.config.schema.FreeParameter`) becomes an + :class:`torch.nn.Parameter`; a plain number becomes a buffer. + + Parameters + ---------- + cfg : V4Config + The configuration. Shift-invariant fluorescence or phase contrast. + + Returns + ------- + system : ImagingSystem + obj : ObjectPrior + noise : Noise + prior : ParameterPrior + sampling : Sampling + + Raises + ------ + NotImplementedError + For tiled field dependence, non-gaussian noise with a spectral + reconstruction, or object priors without a closed-form spectrum. These + follow in later passes; the evidence path already restricts to the + supported subset. + """ + from waveorder.config.schema import ( + PriorTikhonovSpectral, + TikhonovSpectralPrior, + ) + + imaging = cfg.imaging_system + if imaging.field_dependence != "shift_invariant": + raise NotImplementedError("build_system_from_config is shift-invariant for now; tiled follows") + + yx = cfg.data.pixel_size.yx + yx_pixel_size = float(yx) if isinstance(yx, (int, float)) else float(yx.y) + zyx_shape_source = cfg.simulate.shape if cfg.simulate is not None else None + sampling_geometry = Sampling( + zyx_shape=tuple(zyx_shape_source) if zyx_shape_source is not None else (0, 0, 0), + yx_pixel_size=yx_pixel_size, + z_pixel_size=float(cfg.data.pixel_size.z), + z_padding=int(cfg.solver.padding.z), + ) + + detection = imaging.detection + if imaging.contrast == "fluorescence": + if detection.wavelength is None: + raise ValueError("fluorescence imaging_system.detection.wavelength is required") + zernike_coefficients: dict[str | int, object] = {} + if detection.pupil_aberrations is not None: + if detection.pupil_aberrations.zernike_up_to_order is not None: + from waveorder.calibrate.zernike_names import noll_indices_up_to_order + + shorthand = detection.pupil_aberrations.zernike_up_to_order + for noll in noll_indices_up_to_order(shorthand.order): + zernike_coefficients[noll] = free(0.0, shorthand.tolerance) + for name, value in detection.pupil_aberrations.terms.items(): + zernike_coefficients[name] = _value_or_free(value) + system: ImagingSystem = Fluorescence3D( + wavelength_emission=detection.wavelength, + index_of_refraction_media=imaging.index_of_refraction_media, + numerical_aperture_detection=_value_or_free(detection.numerical_aperture), + zernike_coefficients=zernike_coefficients, + z_focus_offset=_value_or_free(detection.z_focus_offset), + ) + elif imaging.contrast == "phase": + illumination = imaging.illumination + if illumination.wavelength is None: + raise ValueError("phase imaging_system.illumination.wavelength is required") + na_ill = illumination.numerical_aperture + system = Phase3D( + wavelength_illumination=illumination.wavelength, + index_of_refraction_media=imaging.index_of_refraction_media, + numerical_aperture_illumination=_value_or_free(na_ill) if na_ill is not None else 0.9, + numerical_aperture_detection=_value_or_free(detection.numerical_aperture), + tilt_angle_zenith=_value_or_free(illumination.tilt.zenith), + tilt_angle_azimuth=_value_or_free(illumination.tilt.azimuth), + invert_phase_contrast=illumination.invert_phase_contrast, + ) + else: # pragma: no cover - schema restricts contrast + raise NotImplementedError(f"contrast {imaging.contrast!r} is not supported") + + # Noise. + if cfg.noise.kind == "gaussian": + gaussian = cfg.noise.gaussian + sigma = gaussian.sigma if isinstance(gaussian.sigma, (int, float)) else 1.0 + noise: Noise = Gaussian(float(sigma)) + elif cfg.noise.kind == "poisson": + poisson = cfg.noise.poisson + noise = Poisson(gain=poisson.gain, offset=poisson.offset) + else: + raise NotImplementedError(f"noise kind {cfg.noise.kind!r} is not supported by the builder yet") + + # Object prior. + prior_spec = cfg.object.prior + weight = cfg.object.weight if isinstance(cfg.object.weight, (int, float)) else 1e-3 + if isinstance(prior_spec, PriorTikhonovSpectral) or isinstance(prior_spec, TikhonovSpectralPrior): + # A fitted / explicit spectrum: deferred to FittedSpectrum (TODO). + obj: ObjectPrior = FittedSpectrum() + elif prior_spec in ("tikhonov", "tv", "none") or prior_spec is None: + obj = Tikhonov(float(weight)) + else: + obj = Tikhonov(float(weight)) + + prior = ParameterPrior.from_system(system) + return system, obj, noise, prior, sampling_geometry From 90f04377232caf6c4f3999553bde4f3dfc884cfe Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Sun, 8 Nov 2026 14:52:49 -0800 Subject: [PATCH 17/24] feat(4.0): FittedSpectrum iso-Whittle object prior (the winning blind strategy) Data-profiled empirical-Bayes object spectrum S(|k|): an |H|^2-weighted shell fit of the data periodogram with a noise floor, broadcast onto the H_k grid. This is the object prior that beat flat Tikhonov on real blind data. evidence() auto-profiles data_fft/otf for data-profiled priors; fixed priors ignore the kwargs. Sampling.shell_index() supplies the isotropic |k| shells. 4 new tests incl. FittedSpectrum < flat Tikhonov on a structured object spectrum. --- tests/calibrate/test_distributions.py | 66 +++++++++- waveorder/distributions.py | 167 +++++++++++++++++++++----- 2 files changed, 203 insertions(+), 30 deletions(-) diff --git a/tests/calibrate/test_distributions.py b/tests/calibrate/test_distributions.py index 949c9523..9fd7d4b8 100644 --- a/tests/calibrate/test_distributions.py +++ b/tests/calibrate/test_distributions.py @@ -337,14 +337,76 @@ def test_build_system_from_config_phase_free_na(): # --- stubs / Poisson / FittedSpectrum -------------------------------------- # -def test_fitted_spectrum_is_a_todo_stub(): +def test_fitted_spectrum_requires_data_and_otf(): import pytest s = _sampling() - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): D.FittedSpectrum().spectrum(s, D.Gaussian(0.1)) +def test_fitted_spectrum_matches_reference_shell_fit(): + """The |H|^2-weighted shell fit matches a direct scatter-add reference.""" + s = _sampling() + system = _fluorescence_system({"astigmatism_0": 0.1}) + otf = system.transfer(s) + g = torch.randn(ZYX, dtype=torch.float64) + noise = D.Gaussian(0.1) + prior = D.FittedSpectrum(n_shells=24) + + data_fft = torch.fft.fftn(g, norm="ortho") + s_k = prior.spectrum(s, noise, data_fft=data_fft, otf=otf) + + # direct reference: S_shell = sum h2 (P - sigma^2) / sum h2^2, clamped + shell, _ = s.shell_index(24) + shell = shell.reshape(-1) + power = (data_fft.abs() ** 2).reshape(-1).to(torch.float64) + h2 = (otf.abs() ** 2).reshape(-1).to(torch.float64) + sigma2 = noise.variance().to(torch.float64) + num = torch.zeros(24, dtype=torch.float64) + den = torch.zeros(24, dtype=torch.float64) + num.scatter_add_(0, shell, h2 * torch.clamp(power - sigma2, min=0.0)) + den.scatter_add_(0, shell, h2 * h2) + ref_shell = (num / den.clamp_min(1e-30)).clamp_min(1e-10) + assert torch.allclose(s_k.reshape(-1), ref_shell[shell], atol=0, rtol=0) + # S_k is positive everywhere (keeps the Gaussian evidence / Wiener finite) + assert (s_k > 0).all() + + +def test_evidence_runs_with_fitted_spectrum(): + """evidence() auto-profiles the data for a data-profiled prior.""" + s = _sampling() + system = _fluorescence_system({"astigmatism_0": 0.1}) + g = torch.randn(ZYX, dtype=torch.float64) + value = D.evidence(g, system, D.FittedSpectrum(n_shells=24), D.Gaussian(0.1), s) + assert torch.isfinite(value) + + +def test_fitted_spectrum_beats_flat_tikhonov_on_structured_object(): + """On an object with a non-flat (power-law) spectrum, the data-profiled + FittedSpectrum gives a higher evidence (lower NLL) at the true system than + a mis-specified flat Tikhonov prior, i.e. it fits the object spectrum the + flat prior gets wrong. This is the winning-strategy property.""" + s = _sampling() + system = _fluorescence_system({"astigmatism_0": 0.15, "coma_x": 0.1}) + otf = system.transfer(s) + + # simulate g = H f + n with a STRUCTURED (power-law) object spectrum + gen = torch.Generator().manual_seed(3) + fz = np.fft.fftfreq(ZYX[0]) + fy = np.fft.fftfreq(ZYX[1]) + fx = np.fft.fftfreq(ZYX[2]) + kr = np.sqrt(fz[:, None, None] ** 2 + fy[None, :, None] ** 2 + fx[None, None, :] ** 2) + s_true = torch.as_tensor(1.0 / (kr**2 + 1e-2), dtype=torch.float64) # red spectrum + f = torch.randn(ZYX, dtype=torch.float64, generator=gen) * s_true.sqrt() + g = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf)) + g = g + torch.randn(ZYX, dtype=torch.float64, generator=gen) * 0.05 + + nll_flat = D.evidence(g, system, D.Tikhonov(1e-2), D.Gaussian(0.05), s) + nll_fit = D.evidence(g, system, D.FittedSpectrum(n_shells=32), D.Gaussian(0.05), s) + assert nll_fit < nll_flat # data-profiled spectrum explains the data better + + def test_gaussian_noise_sample_and_variance(): noise = D.Gaussian(0.5) assert abs(float(noise.variance()) - 0.25) < 1e-12 diff --git a/waveorder/distributions.py b/waveorder/distributions.py index 4879adad..300c84bf 100644 --- a/waveorder/distributions.py +++ b/waveorder/distributions.py @@ -23,7 +23,7 @@ class where there is nothing to carry), mirroring Sri's design. It is a thin, ``p(H)`` parameter prior :class:`ParameterPrior` ``imaging_system`` tolerances ``p(f)`` object prior :class:`Tikhonov` ``object.prior: tikhonov`` ``p(f)`` object prior :class:`TikhonovSpectral` ``object.prior.tikhonov_spectral`` -``p(f)`` object prior :class:`FittedSpectrum` (TODO) ``object.prior`` fitted +``p(f)`` object prior :class:`FittedSpectrum` ``object.prior`` fitted (iso-Whittle) ``p(n)`` noise :class:`Gaussian` ``noise.gaussian`` ``p(n)`` noise :class:`Poisson` ``noise.poisson`` ====================== ========================= ================================ @@ -126,6 +126,43 @@ def axial_nyquist(self, na_detection: float, wavelength: float, index_of_refract """Axial Nyquist spacing (see :func:`waveorder.sampling.axial_nyquist`).""" return sampling.axial_nyquist(wavelength, na_detection, index_of_refraction_media) + def shell_index(self, n_shells: int, device=None) -> tuple[Tensor, Tensor]: + """Isotropic ``|k|`` shell index per voxel, on the transfer-function grid. + + Builds the physical radial frequency ``|k|`` (cycles/um) over the + padded ``(Z, Y, X)`` grid in FFT order and bins it into ``n_shells`` + equal-width shells. Used by :class:`FittedSpectrum` to profile the + object spectrum per shell. + + Parameters + ---------- + n_shells : int + Number of equal-width ``|k|`` shells. + device : torch.device, optional + Device for the returned tensors. + + Returns + ------- + idx : Tensor + Integer shell index per voxel, shape ``padded_shape``, in ``[0, + n_shells)``. + centers : Tensor + Shell center frequencies (cycles/um), shape ``(n_shells,)``. + """ + zz, yy, xx = self.padded_shape + fz = np.fft.fftfreq(zz, d=self.z_pixel_size) + fy = np.fft.fftfreq(yy, d=self.yx_pixel_size) + fx = np.fft.fftfreq(xx, d=self.yx_pixel_size) + fr = np.sqrt(fz[:, None, None] ** 2 + fy[None, :, None] ** 2 + fx[None, None, :] ** 2) + kmax = float(fr.max()) * 1.0001 + 1e-12 + edges = np.linspace(0.0, kmax, n_shells + 1) + idx = np.clip(np.digitize(fr.ravel(), edges) - 1, 0, n_shells - 1).reshape(fr.shape) + centers = 0.5 * (edges[:-1] + edges[1:]) + return ( + torch.as_tensor(idx, dtype=torch.long, device=device), + torch.as_tensor(centers, dtype=torch.float64, device=device), + ) + # --------------------------------------------------------------------------- # # Parameter bookkeeping: {nominal, tolerance} -> Parameter, number -> buffer # @@ -438,9 +475,28 @@ def sample(self, generator: torch.Generator | None = None) -> dict[str, float]: class ObjectPrior(nn.Module): - """Base class for ``p(f)``: returns the object power spectrum ``S_k``.""" + """Base class for ``p(f)``: returns the object power spectrum ``S_k``. + + Priors come in two flavours. A *fixed* prior (``Tikhonov``, + ``TikhonovSpectral``) ignores the data and returns ``S_k`` from stored + parameters. A *data-profiled* (empirical-Bayes) prior (``FittedSpectrum``, + the winning iso-Whittle method) estimates ``S_k`` from the data + periodogram and the current transfer function; :func:`evidence` passes + ``data_fft`` and ``otf`` through for those. The optional keywords keep the + fixed priors' call sites unchanged. + """ + + #: True for priors whose ``spectrum`` consumes ``data_fft``/``otf``. + data_profiled: bool = False - def spectrum(self, s: Sampling, noise: "Noise") -> Tensor | None: # pragma: no cover - abstract + def spectrum( + self, + s: Sampling, + noise: "Noise", + *, + data_fft: Tensor | None = None, + otf: Tensor | None = None, + ) -> Tensor | None: # pragma: no cover - abstract """Object power spectrum ``S_k`` (same shape as ``H_k``), or ``None``.""" raise NotImplementedError @@ -458,7 +514,7 @@ def __init__(self, weight: float) -> None: super().__init__() self.register_buffer("weight", torch.as_tensor(float(weight), dtype=torch.float64)) - def spectrum(self, s: Sampling, noise: "Noise") -> Tensor: + def spectrum(self, s: Sampling, noise: "Noise", **_: Tensor) -> Tensor: """``S_k = sigma^2 / lambda`` (a scalar; broadcast against ``H_k``).""" return noise.variance() / self.weight @@ -476,7 +532,7 @@ def __init__(self, spectrum: Tensor) -> None: super().__init__() self.register_buffer("spectrum_buffer", torch.as_tensor(spectrum, dtype=torch.float64)) - def spectrum(self, s: Sampling, noise: "Noise") -> Tensor: + def spectrum(self, s: Sampling, noise: "Noise", **_: Tensor) -> Tensor: """Return the stored ``S_k``.""" return self.spectrum_buffer @@ -484,32 +540,83 @@ def spectrum(self, s: Sampling, noise: "Noise") -> Tensor: class FittedSpectrum(ObjectPrior): """Isotropic shell-fit object spectrum ``S(|k|)`` estimated from the data. - TODO (next pass): the winning iso-Whittle strategy of the 2026-10-30 - sprint. Fit an isotropic object power spectrum ``S(|k|)`` by a - weighted-least-squares shell fit of the data periodogram with ridge - shrinkage, then broadcast it onto the full ``H_k`` grid. See - ``winning_method_spec.md``. The interface is fixed: - - def spectrum(self, s: Sampling, noise: Noise) -> Tensor: - # 1. periodogram P_k = |FFT(g)|^2 (ortho) of the prepared data, - # 2. bin by |k| shells on the Sampling grid, - # 3. WLS fit S(|k|) = (P_shell - sigma^2) / |H_shell|^2 with - # ridge shrinkage toward a smooth power law, - # 4. scatter S(|k|) back to every frequency -> S_k. - - This flat-prior-beats-nothing placeholder is intentionally inert so that - the interface exists and the config builder can reference it before the - numeric fit lands. + This is the winning "iso-Whittle" object prior of the 2026-10-30 blind + sprint, ported onto the module interface. It is a data-profiled + (empirical-Bayes) prior: instead of a fixed ``S_k``, it estimates the + object power spectrum per isotropic frequency shell directly from the + data periodogram and the current transfer function, which is what beat a + flat Tikhonov prior on real blind data (see the theory page and + ``winning_method_spec.md``). + + Per shell of ``|k|``, with periodogram ``P_k = |FFT(g)|^2`` (ortho), + transfer magnitude ``|H_k|^2`` and noise floor ``sigma^2``, the fit is the + threshold-free, ``|H|^2``-weighted least squares estimate + + S_shell = sum_k |H_k|^2 (P_k - sigma^2) / sum_k |H_k|^4 (clamped >= 0) + + (``whittle3d._wls_shell_S``). The weighting by ``|H|^2`` is what makes it + threshold-free: a fixed ``|H|^2`` cutoff would turn the evidence into an + OTF-support indicator (the round-1 failure). The shell profile is then + broadcast back to every frequency. Being a positive, radially-binned + spectrum, ``S_k`` keeps the Gaussian evidence exact. + + Parameters + ---------- + n_shells : int + Number of isotropic ``|k|`` shells. + min_power : float + Floor on the per-shell spectrum (keeps ``S_k > 0`` so ``log v_k`` and + the Wiener inverse stay finite). """ - def __init__(self) -> None: + data_profiled = True + + def __init__(self, n_shells: int = 64, min_power: float = 1e-10) -> None: super().__init__() + self.n_shells = int(n_shells) + self.min_power = float(min_power) - def spectrum(self, s: Sampling, noise: "Noise") -> Tensor | None: - """Not implemented yet; see the class docstring for the planned fit.""" - raise NotImplementedError( - "FittedSpectrum (iso-Whittle shell fit) lands in the next pass; use Tikhonov or TikhonovSpectral for now." - ) + def spectrum( + self, + s: Sampling, + noise: "Noise", + *, + data_fft: Tensor | None = None, + otf: Tensor | None = None, + ) -> Tensor: + """Data-profiled ``S_k`` via an ``|H|^2``-weighted shell fit. + + Parameters + ---------- + s : Sampling + Grid geometry (supplies the physical ``|k|`` grid). + noise : Noise + Supplies ``sigma^2``. + data_fft : Tensor + ``fftn(g, norm="ortho")`` of the prepared data (``evidence`` passes + this). Required. + otf : Tensor + Current transfer function ``H_k`` (``evidence`` passes this). + Required. + """ + if data_fft is None or otf is None: + raise ValueError( + "FittedSpectrum needs data_fft and otf; call it through distributions.evidence (which passes both)." + ) + sigma2 = noise.variance().to(torch.float64) + power = (data_fft.abs() ** 2).to(torch.float64) + h2 = (otf.abs() ** 2).to(torch.float64) + + shell, _centers = s.shell_index(self.n_shells, device=power.device) + shell = shell.reshape(-1) + p = power.reshape(-1) + h = h2.reshape(-1) + num = torch.zeros(self.n_shells, dtype=torch.float64, device=p.device) + den = torch.zeros(self.n_shells, dtype=torch.float64, device=p.device) + num.scatter_add_(0, shell, h * torch.clamp(p - sigma2, min=0.0)) + den.scatter_add_(0, shell, h * h) + s_shell = (num / den.clamp_min(1e-30)).clamp_min(self.min_power) + return s_shell[shell].reshape(power.shape) # --------------------------------------------------------------------------- # @@ -645,7 +752,11 @@ def evidence( Scalar negative log evidence. """ otf = system.transfer(s) - s_k = obj.spectrum(s, noise) + # data-profiled priors (FittedSpectrum) need the periodogram and H_k; fixed + # priors (Tikhonov, TikhonovSpectral) ignore these kwargs. + if getattr(obj, "data_profiled", False) and data_fft is None: + data_fft = torch.fft.fftn(g, norm="ortho") + s_k = obj.spectrum(s, noise, data_fft=data_fft, otf=otf) sigma = torch.sqrt(noise.variance()) return whittle_evidence( data=None if data_fft is not None else g, From 6bf19879e41e99e683e579249446e6cba7d2c0e7 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 10:50:32 -0800 Subject: [PATCH 18/24] test(config): classify tile-stitch configs by structure not filename; fixes neuromast_tile migration after #591 merge --- tests/config/repo_configs.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/config/repo_configs.py b/tests/config/repo_configs.py index fef8e577..ffaf79aa 100644 --- a/tests/config/repo_configs.py +++ b/tests/config/repo_configs.py @@ -9,6 +9,8 @@ from pathlib import Path +import yaml + REPO_ROOT = Path(__file__).parent.parent.parent #: Directories holding reconstruction configuration files. @@ -23,9 +25,35 @@ #: 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 "tile-stitch" in p.name] +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 = [ From 33fa744e44af479723a2bb7b68e8b05cd2c0bf55 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 11:21:21 -0800 Subject: [PATCH 19/24] fix(bench): drop redundant cli kwarg splat duplicated in the #591 merge (run_hpc_case received cli twice; broke all hpc cases) --- waveorder/cli/bench.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/waveorder/cli/bench.py b/waveorder/cli/bench.py index ad986a27..9a082fdf 100644 --- a/waveorder/cli/bench.py +++ b/waveorder/cli/bench.py @@ -173,11 +173,6 @@ def run( save_all=save_all, reference=case.reference, crop=case.crop, - # Forward-compatible with the expanded suite - # (bench/expand-hpc-cases), whose cases may select the - # tile-stitch CLI; absent there, run_hpc_case defaults - # to "rec". - **({"cli": case.cli} if getattr(case, "cli", None) is not None else {}), # Forward-compatible with runners that predate the flag. **({"migrate_config": True} if migrate_configs else {}), ) From 010d5548ef679b5051deb3ac1ef80cce33cfe737 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 11:21:32 -0800 Subject: [PATCH 20/24] feat(calibrate): wire FittedSpectrum iso-Whittle prior into the evidence path; object.prior=tikhonov_spectral now profiles S(|k|) per shell (n_shells config), flat Tikhonov default untouched --- tests/calibrate/test_evidence_search.py | 25 +++++++++++- waveorder/calibrate/evidence_search.py | 51 +++++++++++++++++++++---- waveorder/config/schema.py | 8 +++- waveorder/distributions.py | 8 ++-- 4 files changed, 77 insertions(+), 15 deletions(-) diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py index 3b599614..f3b64f9c 100644 --- a/tests/calibrate/test_evidence_search.py +++ b/tests/calibrate/test_evidence_search.py @@ -25,7 +25,7 @@ SIGMA, LAM = 0.05, 1e-2 -def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA) -> V4Config: +def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA, prior="tikhonov") -> V4Config: return V4Config( **{ "version": 4, @@ -42,7 +42,7 @@ def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA) -> V4 }, }, }, - "object": {"prior": "tikhonov", "weight": weight}, + "object": {"prior": prior, "weight": weight}, "noise": {"gaussian": {"sigma": sigma}}, "solver": { "calibration": { @@ -88,6 +88,27 @@ def test_evidence_gradient_recovers_truth(synthetic_data): assert np.sqrt(variances).max() < 0.1 # tight posterior at this SNR +def test_fitted_spectrum_prior_is_wired(synthetic_data): + # object.prior: tikhonov_spectral routes the evidence through the + # FittedSpectrum iso-Whittle prior (S(|k|) profiled per isotropic |k| + # shell) instead of the flat Tikhonov S_k = sigma^2 / lambda. Assert the + # prior is selectable from a config, takes the fitted path end to end, and + # is actually in the loop (it moves the estimate relative to the flat run). + # Fitted-vs-flat accuracy is data-dependent -- the fitted prior wins on + # structured real data, not on this white-spectrum phantom -- so recovery + # accuracy is not asserted here. + fitted = run_evidence_calibration( + _config(prior={"tikhonov_spectral": {"n_shells": 8}}), synthetic_data + ) + flat = run_evidence_calibration(_config(), synthetic_data) + assert fitted.extras["object_prior"] == "tikhonov_spectral" + assert flat.extras["object_prior"] == "tikhonov" + assert np.isfinite(fitted.best_score) + assert not np.allclose(fitted.parameters, flat.parameters) # the fitted spectrum is active + assert np.abs(fitted.parameters).max() < 0.5 # a sane calibration, not blown up + assert np.all(np.isfinite(fitted.parameters_covariance)) + + 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) diff --git a/waveorder/calibrate/evidence_search.py b/waveorder/calibrate/evidence_search.py index 6ea2187a..67f7eda7 100644 --- a/waveorder/calibrate/evidence_search.py +++ b/waveorder/calibrate/evidence_search.py @@ -28,9 +28,11 @@ 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.distributions import FittedSpectrum, Gaussian, Sampling from waveorder.config.schema import ( FreeParameter, GaussianNoise, + PriorTikhonovSpectral, SearchGradient, SearchGrid, SearchRandom, @@ -311,19 +313,49 @@ def run_evidence_calibration( else: raise ValueError(f"object.weight: {cfg.object.weight} is not usable on the evidence path") + # Object prior spectrum S_k. The default flat Tikhonov prior (S_k = + # sigma^2 / lambda) is formed inside whittle_evidence. object.prior = + # tikhonov_spectral instead profiles S(|k|) from the data per isotropic + # |k| shell with an |H|^2-weighted fit (the FittedSpectrum iso-Whittle + # prior), re-fit at every evaluation so it co-optimizes with the aberrations. + use_fitted_spectrum = isinstance(cfg.object.prior, PriorTikhonovSpectral) + if use_fitted_spectrum: + fitted_spectrum = FittedSpectrum(n_shells=cfg.object.prior.tikhonov_spectral.n_shells) + spectrum_sampling = Sampling( + zyx_shape=tuple(zyx_data.shape), + yx_pixel_size=optics_kwargs["yx_pixel_size"], + z_pixel_size=float(cfg.data.pixel_size.z), + z_padding=0, + ) + 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, - ) + if use_fitted_spectrum: + object_spectrum = fitted_spectrum.spectrum( + spectrum_sampling, + Gaussian(float(torch.exp(log_sigma))), + data_fft=data_fft, + otf=otf, + ) + data_term = whittle_evidence( + data_fft=data_fft, + otf=otf, + noise_sigma=torch.exp(log_sigma), + object_spectrum=object_spectrum, + exclude_dc=True, + ) + else: + 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) @@ -476,5 +508,8 @@ def score_of(values: torch.Tensor) -> torch.Tensor: noise_sigma=float(torch.exp(log_sigma)), weight=float(torch.exp(log_weight)), best_score=best_score, - extras={"score": score_name}, + extras={ + "score": score_name, + "object_prior": "tikhonov_spectral" if use_fitted_spectrum else "tikhonov", + }, ) diff --git a/waveorder/config/schema.py b/waveorder/config/schema.py index 9e988454..7f25bc0b 100644 --- a/waveorder/config/schema.py +++ b/waveorder/config/schema.py @@ -274,8 +274,12 @@ def _na_below_index(self): class TikhonovSpectralPrior(V4BaseModel): - exponent: Optional[float] = Field(default=None, description="power-law exponent of the object spectrum") - fit_from: Optional[str] = Field(default=None, description="path to fit the spectrum from") + n_shells: PositiveInt = Field( + default=64, + description="number of isotropic |k| shells for the data-profiled iso-Whittle object " + "spectrum; the per-shell |H|^2-weighted fit (FittedSpectrum) replaces the flat Tikhonov " + "prior S_k = sigma^2 / lambda and is re-fit as the imaging system is calibrated", + ) class BeadsPrior(V4BaseModel): diff --git a/waveorder/distributions.py b/waveorder/distributions.py index 300c84bf..360773f1 100644 --- a/waveorder/distributions.py +++ b/waveorder/distributions.py @@ -1068,9 +1068,11 @@ def build_system_from_config(cfg): # Object prior. prior_spec = cfg.object.prior weight = cfg.object.weight if isinstance(cfg.object.weight, (int, float)) else 1e-3 - if isinstance(prior_spec, PriorTikhonovSpectral) or isinstance(prior_spec, TikhonovSpectralPrior): - # A fitted / explicit spectrum: deferred to FittedSpectrum (TODO). - obj: ObjectPrior = FittedSpectrum() + if isinstance(prior_spec, PriorTikhonovSpectral): + # Data-profiled iso-Whittle spectrum (same prior the evidence path uses). + obj: ObjectPrior = FittedSpectrum(n_shells=prior_spec.tikhonov_spectral.n_shells) + elif isinstance(prior_spec, TikhonovSpectralPrior): + obj = FittedSpectrum(n_shells=prior_spec.n_shells) elif prior_spec in ("tikhonov", "tv", "none") or prior_spec is None: obj = Tikhonov(float(weight)) else: From 0b2e4006fbc139953edad58762da9816843a583c Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 11:36:37 -0800 Subject: [PATCH 21/24] refactor(config): drop unimplemented 'sure' score from the schema Literal and weight:auto guard; the schema now promises only implemented scores (evidence, gcv, sharpness) --- waveorder/calibrate/run.py | 2 +- waveorder/config/schema.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/waveorder/calibrate/run.py b/waveorder/calibrate/run.py index ce528239..b3c3f019 100644 --- a/waveorder/calibrate/run.py +++ b/waveorder/calibrate/run.py @@ -133,7 +133,7 @@ def build_recovery_settings(cfg: V4Config, yx_shape: tuple[int, int]) -> tuple[R calibration = cfg.solver.calibration score = calibration.score if calibration is not None else "sharpness" if score != "sharpness": - raise _not_yet(f"score: {score} (evidence/gcv/sure land with design step 4; use 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: diff --git a/waveorder/config/schema.py b/waveorder/config/schema.py index 7f25bc0b..b6529ac0 100644 --- a/waveorder/config/schema.py +++ b/waveorder/config/schema.py @@ -551,7 +551,7 @@ class CalibrationAdvanced(V4BaseModel): class Calibration(V4BaseModel): """How H is estimated; read only when something has a tolerance.""" - score: Literal["evidence", "gcv", "sure", "sharpness"] = "evidence" + 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 @@ -686,9 +686,9 @@ def _validation_rules(self): if self.object.weight == "auto": score = calibration.score if calibration is not None else None - if anything_estimated and score not in ("gcv", "evidence", "sure"): + if anything_estimated and score not in ("gcv", "evidence"): raise ValueError( - "weight: auto requires solver.calibration.score in {gcv, evidence, sure} " + "weight: auto requires solver.calibration.score in {gcv, evidence} " "when imaging-system parameters are estimated" ) From 24eb91a981d92218c5f4a800eba70b9c8d34dee8 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 11:55:17 -0800 Subject: [PATCH 22/24] fix(migrate): raise a clear MigrationError for tile-stitch configs (tile/blend/recon) instead of a raw pydantic error; they stay 3.x like birefringence, so --migrate-configs skips them and wo tile-stitch keeps running. Add a guard test over TILE_STITCH_CONFIG_PATHS --- tests/config/test_migrate.py | 13 +++++++++++++ waveorder/config/migrate.py | 10 +++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/config/test_migrate.py b/tests/config/test_migrate.py index 450dbb53..55a56a8c 100644 --- a/tests/config/test_migrate.py +++ b/tests/config/test_migrate.py @@ -87,6 +87,19 @@ def test_birefringence_not_migratable(config_path): 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 = { diff --git a/waveorder/config/migrate.py b/waveorder/config/migrate.py index 401e66f8..e3622c04 100644 --- a/waveorder/config/migrate.py +++ b/waveorder/config/migrate.py @@ -82,8 +82,16 @@ def migrate_3x_dict(raw: dict) -> schema.V4Config: Raises ------ MigrationError - For birefringence configurations, which stay 3.x in this release. + 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) From 780a13398d1f26fc7e252ed68950410373caae80 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 13:14:25 -0800 Subject: [PATCH 23/24] style: ruff check + format (import order, line length, blank line) to pass the Lint CI --- tests/calibrate/test_evidence_search.py | 4 +--- tests/config/repo_configs.py | 1 + waveorder/calibrate/evidence_search.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py index f3b64f9c..f81c7a97 100644 --- a/tests/calibrate/test_evidence_search.py +++ b/tests/calibrate/test_evidence_search.py @@ -97,9 +97,7 @@ def test_fitted_spectrum_prior_is_wired(synthetic_data): # Fitted-vs-flat accuracy is data-dependent -- the fitted prior wins on # structured real data, not on this white-spectrum phantom -- so recovery # accuracy is not asserted here. - fitted = run_evidence_calibration( - _config(prior={"tikhonov_spectral": {"n_shells": 8}}), synthetic_data - ) + fitted = run_evidence_calibration(_config(prior={"tikhonov_spectral": {"n_shells": 8}}), synthetic_data) flat = run_evidence_calibration(_config(), synthetic_data) assert fitted.extras["object_prior"] == "tikhonov_spectral" assert flat.extras["object_prior"] == "tikhonov" diff --git a/tests/config/repo_configs.py b/tests/config/repo_configs.py index ffaf79aa..1ed8920c 100644 --- a/tests/config/repo_configs.py +++ b/tests/config/repo_configs.py @@ -25,6 +25,7 @@ #: 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. diff --git a/waveorder/calibrate/evidence_search.py b/waveorder/calibrate/evidence_search.py index 67f7eda7..9525abcb 100644 --- a/waveorder/calibrate/evidence_search.py +++ b/waveorder/calibrate/evidence_search.py @@ -28,7 +28,6 @@ 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.distributions import FittedSpectrum, Gaussian, Sampling from waveorder.config.schema import ( FreeParameter, GaussianNoise, @@ -39,6 +38,7 @@ V4Config, free_parameters, ) +from waveorder.distributions import FittedSpectrum, Gaussian, Sampling @dataclass From f31d6c688d00b9ea6b4f95b7c1bc153493fd1730 Mon Sep 17 00:00:00 2001 From: Talon Chandler Date: Tue, 10 Nov 2026 13:39:50 -0800 Subject: [PATCH 24/24] revert(4.0): drop the unintegrated OO distribution layer and its tikhonov_spectral prior Removes waveorder/distributions.py (nn.Module-per-distribution layer: ImagingSystem/ObjectPrior/Noise/calibrate/reconstruct/build_system_from_config) and its tests. Only FittedSpectrum was wired anywhere (the tikhonov_spectral object prior on the evidence path); that hook is reverted too, so the schema, evidence_search, and tests no longer reference it. The flat-Tikhonov evidence/gcv/sharpness/beads calibration paths are unchanged. This scopes #592 to the config/migrate/calibrate work; the OO layer stays in git history for a future dedicated PR once it drives a CLI path. --- tests/calibrate/test_distributions.py | 422 --------- tests/calibrate/test_evidence_search.py | 23 +- waveorder/calibrate/evidence_search.py | 51 +- waveorder/calibrate/scores.py | 4 +- waveorder/config/schema.py | 20 +- waveorder/distributions.py | 1082 ----------------------- 6 files changed, 14 insertions(+), 1588 deletions(-) delete mode 100644 tests/calibrate/test_distributions.py delete mode 100644 waveorder/distributions.py diff --git a/tests/calibrate/test_distributions.py b/tests/calibrate/test_distributions.py deleted file mode 100644 index 9fd7d4b8..00000000 --- a/tests/calibrate/test_distributions.py +++ /dev/null @@ -1,422 +0,0 @@ -"""The nn.Module-per-distribution layer wraps the existing numerics exactly. - -These are the equivalence/behavior tests of the WaveOrder 4.0 distribution -layer (``waveorder.distributions``). Each one pins the OO layer to the tested -function API it wraps, so the regression suite and the new layer cannot drift. -""" - -import numpy as np -import torch -import yaml - -from waveorder import distributions as D -from waveorder.calibrate.scores import calibration_penalty, whittle_evidence -from waveorder.calibrate.synthesis import fluorescence_otf -from waveorder.config.schema import V4Config -from waveorder.models import phase_thick_3d -from waveorder.reconstruct import tikhonov_regularized_inverse_filter - -ZYX = (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, -) - - -def _fluorescence_system(coefficients=None, normalize="peak"): - return D.Fluorescence3D( - wavelength_emission=OPTICS["wavelength_emission"], - index_of_refraction_media=OPTICS["index_of_refraction_media"], - numerical_aperture_detection=OPTICS["numerical_aperture_detection"], - zernike_coefficients=coefficients or {}, - normalize=normalize, - ) - - -def _sampling(shape=ZYX, z_padding=0): - return D.Sampling( - zyx_shape=shape, - yx_pixel_size=OPTICS["yx_pixel_size"], - z_pixel_size=OPTICS["z_pixel_size"], - z_padding=z_padding, - ) - - -def _simulate(otf_truth, noise_sigma, regularization_strength, seed=0): - """Draw g = H f + n with f from the Tikhonov prior S_k = sigma^2 / lambda.""" - generator = torch.Generator().manual_seed(seed) - prior_std = noise_sigma / regularization_strength**0.5 - f = torch.randn(otf_truth.shape, generator=generator) * prior_std - g = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf_truth)) - g = g + torch.randn(otf_truth.shape, generator=generator) * noise_sigma - return g.double() - - -# --- 1. evidence == whittle_evidence ---------------------------------------- # - - -def test_evidence_matches_whittle_evidence(): - s = _sampling() - coefficients = {"astigmatism_0": 0.12, "spherical": -0.05} - system = _fluorescence_system(coefficients) - sigma, weight = 0.1, 1e-2 - g = torch.randn(ZYX, dtype=torch.float64) - - value = D.evidence(g, system, D.Tikhonov(weight), D.Gaussian(sigma), s) - reference = whittle_evidence( - data=g, - otf=fluorescence_otf(zyx_shape=ZYX, zernike_coefficients=coefficients, **OPTICS), - noise_sigma=sigma, - regularization_strength=weight, - ) - assert abs(float(value) - float(reference)) < 1e-10 - - -def test_evidence_matches_whittle_evidence_exclude_dc(): - s = _sampling() - system = _fluorescence_system({"coma_x": 0.07}) - sigma, weight = 0.2, 5e-3 - g = torch.randn(ZYX, dtype=torch.float64) - - value = D.evidence(g, system, D.Tikhonov(weight), D.Gaussian(sigma), s, exclude_dc=True) - reference = whittle_evidence( - data=g, - otf=fluorescence_otf(zyx_shape=ZYX, zernike_coefficients={"coma_x": 0.07}, **OPTICS), - noise_sigma=sigma, - regularization_strength=weight, - exclude_dc=True, - ) - assert abs(float(value) - float(reference)) < 1e-10 - - -# --- 2. Fluorescence3D.transfer == fluorescence_otf ------------------------- # - - -def test_fluorescence_transfer_matches_otf(): - s = _sampling() - coefficients = {"astigmatism_0": 0.12, "spherical": -0.05, "defocus": 0.03} - system = _fluorescence_system(coefficients) - transfer = system.transfer(s) - reference = fluorescence_otf(zyx_shape=ZYX, zernike_coefficients=coefficients, **OPTICS) - assert torch.equal(transfer, reference) - - -def test_fluorescence_transfer_with_z_padding(): - s = _sampling(z_padding=2) - system = _fluorescence_system({"spherical": 0.1}) - transfer = system.transfer(s) - reference = fluorescence_otf( - zyx_shape=(ZYX[0] + 4, ZYX[1], ZYX[2]), - zernike_coefficients={"spherical": 0.1}, - **OPTICS, - ) - assert torch.equal(transfer, reference) - - -# --- 3. ParameterPrior.nll == calibration_penalty -------------------------- # - - -def test_parameter_prior_matches_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"} - prior = D.ParameterPrior(nominals, tolerances) - assert torch.equal(prior.nll(values), calibration_penalty(values, nominals, tolerances)) - - -def test_parameter_prior_from_system(): - system = _fluorescence_system({"spherical": D.free(0.0, 0.1), "astigmatism_0": 0.2}) - prior = D.ParameterPrior.from_system(system) - assert prior.nominals == {"pupil_aberrations.spherical": 0.0} - assert prior.tolerances == {"pupil_aberrations.spherical": 0.1} - values = {"pupil_aberrations.spherical": 0.05} - assert torch.equal(prior.nll(values), calibration_penalty(values, prior.nominals, prior.tolerances)) - - -# --- 4. reconstruct (Tikhonov) == closed-form Wiener ----------------------- # - - -def test_reconstruct_matches_tikhonov_closed_form_phase(): - s = _sampling(shape=(12, 24, 24), z_padding=0) - system = D.Phase3D( - wavelength_illumination=0.532, - index_of_refraction_media=1.3, - numerical_aperture_illumination=0.9, - numerical_aperture_detection=1.2, - ) - sigma, weight = 0.1, 1e-2 - noise = D.Gaussian(sigma) - obj = D.Tikhonov(weight) - g = torch.randn(12, 24, 24, dtype=torch.float64) - - estimate = D.reconstruct(g, system, obj, noise, s) - otf = system.transfer(s) - regularization = float(noise.variance() / obj.spectrum(s, noise)) - inverse_filter = tikhonov_regularized_inverse_filter(otf, regularization) - reference = torch.real(torch.fft.ifftn(torch.fft.fftn(g.to(torch.complex64)) * inverse_filter)) - assert torch.allclose(estimate, reference, atol=1e-6) - - -def test_reconstruct_tikhonov_regularization_is_sigma2_over_s_k(): - # S_k = sigma^2 / lambda, so sigma^2 / S_k = lambda exactly. - s = _sampling() - system = _fluorescence_system({"spherical": 0.1}) - sigma, weight = 0.3, 7e-3 - noise = D.Gaussian(sigma) - obj = D.Tikhonov(weight) - g = torch.randn(ZYX, dtype=torch.float64) - - estimate = D.reconstruct(g, system, obj, noise, s) - inverse_filter = tikhonov_regularized_inverse_filter(system.transfer(s), weight) - reference = torch.real(torch.fft.ifftn(torch.fft.fftn(g.to(torch.complex64)) * inverse_filter)) - assert torch.allclose(estimate, reference, atol=1e-6) - - -# --- 5. Phase3D.transfer == phase_thick_3d.calculate_transfer_function ----- # - - -def test_phase_transfer_matches_model(): - s = _sampling(shape=(8, 16, 16), z_padding=1) - system = D.Phase3D( - wavelength_illumination=0.532, - index_of_refraction_media=1.3, - numerical_aperture_illumination=0.9, - numerical_aperture_detection=1.2, - tilt_angle_zenith=0.05, - ) - real_tf, imag_tf = system.transfer_real_imag(s) - reference_real, reference_imag = phase_thick_3d.calculate_transfer_function( - (8, 16, 16), - 0.2, - 0.4, - 0.532, - 1, - 1.3, - 0.9, - 1.2, - tilt_angle_zenith=0.05, - ) - assert torch.equal(real_tf, reference_real) - assert torch.equal(imag_tf, reference_imag) - # transfer() returns the phase (real-potential) WOTF as the single H_k. - assert torch.equal(system.transfer(s), reference_real) - - -# --- 6. going blind flips an argument and calibrate recovers --------------- # - - -def test_going_blind_flips_one_argument_to_parameter(): - fixed = _fluorescence_system({"spherical": 0.0}) - assert fixed.free_parameter_names() == [] - blind = _fluorescence_system({"spherical": D.free(0.0, 0.5)}) - assert blind.free_parameter_names() == ["pupil_aberrations.spherical"] - # The free value is an nn.Parameter; a plain number is a buffer. - names = dict(blind.named_parameters()) - assert isinstance(names["pupil_aberrations.spherical"], torch.nn.Parameter) - assert "pupil_aberrations.spherical" not in dict(fixed.named_parameters()) - - -def test_calibrate_recovers_injected_aberration(): - # Only the magnitude of a lone astigmatism term is identifiable (the - # PSF-flip sign ambiguity), matching test_synthesis_scores.py. - s = _sampling() - sigma, weight = 0.05, 1e-2 - truth = 0.2 - otf_truth = fluorescence_otf(zyx_shape=ZYX, zernike_coefficients={"astigmatism_0": truth}, **OPTICS) - g = _simulate(otf_truth, sigma, weight, seed=3) - - system = _fluorescence_system({"astigmatism_0": D.free(0.0, 0.5)}) - prior = D.ParameterPrior.from_system(system) - result = D.calibrate( - g, system, prior, D.Tikhonov(weight), D.Gaussian(sigma), s, max_iter=200, lr=0.02, exclude_dc=True - ) - recovered = result.parameters["pupil_aberrations.astigmatism_0"] - assert abs(abs(recovered) - truth) < 0.11, result.parameters - assert result.covariance.shape == (1, 1) - assert np.isfinite(result.covariance).all() - - -# --- 7. state_dict round-trip, .to, dotted tf.zarr names ------------------- # - - -def test_state_dict_round_trip(): - s = _sampling() - system = _fluorescence_system({"spherical": D.free(0.17, 0.1), "astigmatism_0": 0.05}) - restored = _fluorescence_system({"spherical": D.free(0.0, 0.1), "astigmatism_0": 0.05}) - restored.load_state_dict(system.state_dict()) - assert torch.equal(system.transfer(s), restored.transfer(s)) - - -def test_to_device_cpu(): - system = _fluorescence_system({"spherical": D.free(0.1, 0.1)}) - moved = system.to("cpu") - assert moved is system - s = _sampling() - assert moved.transfer(s).shape == ZYX - - -def test_dotted_names_usable_as_tf_zarr_parameter_names(): - import tempfile - - from waveorder.calibrate.tfzarr import read_tf_zarr, write_tf_zarr - - system = _fluorescence_system({"astigmatism_0": D.free(0.1, 0.05), "spherical": D.free(-0.02, 0.05)}) - names = system.free_parameter_names() - assert names == ["pupil_aberrations.astigmatism_0", "pupil_aberrations.spherical"] - values = np.array([[float(p.detach()) for _, p in system.named_parameters()]]) - with tempfile.TemporaryDirectory() as tmp: - path = f"{tmp}/tf.zarr" - write_tf_zarr(path, values, names, tile_grid={"size_yx": [32, 32]}) - store = read_tf_zarr(path) - assert store.parameter_names == names - assert np.allclose(store.parameters, values) - - -# --- config adapter -------------------------------------------------------- # - - -def test_build_system_from_config_fluorescence(): - cfg = V4Config( - **yaml.safe_load( - """ -version: 4 -data: - 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: {nominal: 0.0, tolerance: 0.1} - astigmatism_0: 0.05 -object: {prior: tikhonov, weight: 0.01} -noise: {gaussian: {sigma: 0.1}} -""" - ) - ) - system, obj, noise, prior, sampling = D.build_system_from_config(cfg) - assert isinstance(system, D.Fluorescence3D) - assert system.free_parameter_names() == ["pupil_aberrations.spherical"] - assert isinstance(obj, D.Tikhonov) - assert isinstance(noise, D.Gaussian) - assert prior.tolerances == {"pupil_aberrations.spherical": 0.1} - assert float(obj.weight) == 0.01 - assert float(noise.sigma) == 0.1 - - -def test_build_system_from_config_phase_free_na(): - cfg = V4Config( - **yaml.safe_load( - """ -version: 4 -data: - pixel_size: {yx: 0.2, z: 0.4} -imaging_system: - contrast: phase - index_of_refraction_media: 1.3 - detection: {numerical_aperture: 1.2} - illumination: - wavelength: 0.532 - numerical_aperture: {nominal: 0.9, tolerance: 0.05} -object: {prior: tikhonov, weight: 0.001} -noise: {gaussian: {sigma: 1.0}} -""" - ) - ) - system, _, _, prior, _ = D.build_system_from_config(cfg) - assert isinstance(system, D.Phase3D) - assert system.free_parameter_names() == ["numerical_aperture_illumination"] - assert prior.tolerances == {"numerical_aperture_illumination": 0.05} - - -# --- stubs / Poisson / FittedSpectrum -------------------------------------- # - - -def test_fitted_spectrum_requires_data_and_otf(): - import pytest - - s = _sampling() - with pytest.raises(ValueError): - D.FittedSpectrum().spectrum(s, D.Gaussian(0.1)) - - -def test_fitted_spectrum_matches_reference_shell_fit(): - """The |H|^2-weighted shell fit matches a direct scatter-add reference.""" - s = _sampling() - system = _fluorescence_system({"astigmatism_0": 0.1}) - otf = system.transfer(s) - g = torch.randn(ZYX, dtype=torch.float64) - noise = D.Gaussian(0.1) - prior = D.FittedSpectrum(n_shells=24) - - data_fft = torch.fft.fftn(g, norm="ortho") - s_k = prior.spectrum(s, noise, data_fft=data_fft, otf=otf) - - # direct reference: S_shell = sum h2 (P - sigma^2) / sum h2^2, clamped - shell, _ = s.shell_index(24) - shell = shell.reshape(-1) - power = (data_fft.abs() ** 2).reshape(-1).to(torch.float64) - h2 = (otf.abs() ** 2).reshape(-1).to(torch.float64) - sigma2 = noise.variance().to(torch.float64) - num = torch.zeros(24, dtype=torch.float64) - den = torch.zeros(24, dtype=torch.float64) - num.scatter_add_(0, shell, h2 * torch.clamp(power - sigma2, min=0.0)) - den.scatter_add_(0, shell, h2 * h2) - ref_shell = (num / den.clamp_min(1e-30)).clamp_min(1e-10) - assert torch.allclose(s_k.reshape(-1), ref_shell[shell], atol=0, rtol=0) - # S_k is positive everywhere (keeps the Gaussian evidence / Wiener finite) - assert (s_k > 0).all() - - -def test_evidence_runs_with_fitted_spectrum(): - """evidence() auto-profiles the data for a data-profiled prior.""" - s = _sampling() - system = _fluorescence_system({"astigmatism_0": 0.1}) - g = torch.randn(ZYX, dtype=torch.float64) - value = D.evidence(g, system, D.FittedSpectrum(n_shells=24), D.Gaussian(0.1), s) - assert torch.isfinite(value) - - -def test_fitted_spectrum_beats_flat_tikhonov_on_structured_object(): - """On an object with a non-flat (power-law) spectrum, the data-profiled - FittedSpectrum gives a higher evidence (lower NLL) at the true system than - a mis-specified flat Tikhonov prior, i.e. it fits the object spectrum the - flat prior gets wrong. This is the winning-strategy property.""" - s = _sampling() - system = _fluorescence_system({"astigmatism_0": 0.15, "coma_x": 0.1}) - otf = system.transfer(s) - - # simulate g = H f + n with a STRUCTURED (power-law) object spectrum - gen = torch.Generator().manual_seed(3) - fz = np.fft.fftfreq(ZYX[0]) - fy = np.fft.fftfreq(ZYX[1]) - fx = np.fft.fftfreq(ZYX[2]) - kr = np.sqrt(fz[:, None, None] ** 2 + fy[None, :, None] ** 2 + fx[None, None, :] ** 2) - s_true = torch.as_tensor(1.0 / (kr**2 + 1e-2), dtype=torch.float64) # red spectrum - f = torch.randn(ZYX, dtype=torch.float64, generator=gen) * s_true.sqrt() - g = torch.real(torch.fft.ifftn(torch.fft.fftn(f) * otf)) - g = g + torch.randn(ZYX, dtype=torch.float64, generator=gen) * 0.05 - - nll_flat = D.evidence(g, system, D.Tikhonov(1e-2), D.Gaussian(0.05), s) - nll_fit = D.evidence(g, system, D.FittedSpectrum(n_shells=32), D.Gaussian(0.05), s) - assert nll_fit < nll_flat # data-profiled spectrum explains the data better - - -def test_gaussian_noise_sample_and_variance(): - noise = D.Gaussian(0.5) - assert abs(float(noise.variance()) - 0.25) < 1e-12 - generator = torch.Generator().manual_seed(0) - draw = noise.sample(torch.zeros(5000, dtype=torch.float64), generator=generator) - assert abs(float(draw.std()) - 0.5) < 0.05 - - -def test_poisson_sample_mean(): - noise = D.Poisson(gain=1.0) - generator = torch.Generator().manual_seed(0) - draw = noise.sample(torch.full((5000,), 10.0, dtype=torch.float64), generator=generator) - assert abs(float(draw.mean()) - 10.0) < 0.3 diff --git a/tests/calibrate/test_evidence_search.py b/tests/calibrate/test_evidence_search.py index f81c7a97..3b599614 100644 --- a/tests/calibrate/test_evidence_search.py +++ b/tests/calibrate/test_evidence_search.py @@ -25,7 +25,7 @@ SIGMA, LAM = 0.05, 1e-2 -def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA, prior="tikhonov") -> V4Config: +def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA) -> V4Config: return V4Config( **{ "version": 4, @@ -42,7 +42,7 @@ def _config(score: str = "evidence", search=None, weight=LAM, sigma=SIGMA, prior }, }, }, - "object": {"prior": prior, "weight": weight}, + "object": {"prior": "tikhonov", "weight": weight}, "noise": {"gaussian": {"sigma": sigma}}, "solver": { "calibration": { @@ -88,25 +88,6 @@ def test_evidence_gradient_recovers_truth(synthetic_data): assert np.sqrt(variances).max() < 0.1 # tight posterior at this SNR -def test_fitted_spectrum_prior_is_wired(synthetic_data): - # object.prior: tikhonov_spectral routes the evidence through the - # FittedSpectrum iso-Whittle prior (S(|k|) profiled per isotropic |k| - # shell) instead of the flat Tikhonov S_k = sigma^2 / lambda. Assert the - # prior is selectable from a config, takes the fitted path end to end, and - # is actually in the loop (it moves the estimate relative to the flat run). - # Fitted-vs-flat accuracy is data-dependent -- the fitted prior wins on - # structured real data, not on this white-spectrum phantom -- so recovery - # accuracy is not asserted here. - fitted = run_evidence_calibration(_config(prior={"tikhonov_spectral": {"n_shells": 8}}), synthetic_data) - flat = run_evidence_calibration(_config(), synthetic_data) - assert fitted.extras["object_prior"] == "tikhonov_spectral" - assert flat.extras["object_prior"] == "tikhonov" - assert np.isfinite(fitted.best_score) - assert not np.allclose(fitted.parameters, flat.parameters) # the fitted spectrum is active - assert np.abs(fitted.parameters).max() < 0.5 # a sane calibration, not blown up - assert np.all(np.isfinite(fitted.parameters_covariance)) - - 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) diff --git a/waveorder/calibrate/evidence_search.py b/waveorder/calibrate/evidence_search.py index 9525abcb..6ea2187a 100644 --- a/waveorder/calibrate/evidence_search.py +++ b/waveorder/calibrate/evidence_search.py @@ -31,14 +31,12 @@ from waveorder.config.schema import ( FreeParameter, GaussianNoise, - PriorTikhonovSpectral, SearchGradient, SearchGrid, SearchRandom, V4Config, free_parameters, ) -from waveorder.distributions import FittedSpectrum, Gaussian, Sampling @dataclass @@ -313,49 +311,19 @@ def run_evidence_calibration( else: raise ValueError(f"object.weight: {cfg.object.weight} is not usable on the evidence path") - # Object prior spectrum S_k. The default flat Tikhonov prior (S_k = - # sigma^2 / lambda) is formed inside whittle_evidence. object.prior = - # tikhonov_spectral instead profiles S(|k|) from the data per isotropic - # |k| shell with an |H|^2-weighted fit (the FittedSpectrum iso-Whittle - # prior), re-fit at every evaluation so it co-optimizes with the aberrations. - use_fitted_spectrum = isinstance(cfg.object.prior, PriorTikhonovSpectral) - if use_fitted_spectrum: - fitted_spectrum = FittedSpectrum(n_shells=cfg.object.prior.tikhonov_spectral.n_shells) - spectrum_sampling = Sampling( - zyx_shape=tuple(zyx_data.shape), - yx_pixel_size=optics_kwargs["yx_pixel_size"], - z_pixel_size=float(cfg.data.pixel_size.z), - z_padding=0, - ) - 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": - if use_fitted_spectrum: - object_spectrum = fitted_spectrum.spectrum( - spectrum_sampling, - Gaussian(float(torch.exp(log_sigma))), - data_fft=data_fft, - otf=otf, - ) - data_term = whittle_evidence( - data_fft=data_fft, - otf=otf, - noise_sigma=torch.exp(log_sigma), - object_spectrum=object_spectrum, - exclude_dc=True, - ) - else: - 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, - ) + 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) @@ -508,8 +476,5 @@ def score_of(values: torch.Tensor) -> torch.Tensor: noise_sigma=float(torch.exp(log_sigma)), weight=float(torch.exp(log_weight)), best_score=best_score, - extras={ - "score": score_name, - "object_prior": "tikhonov_spectral" if use_fitted_spectrum else "tikhonov", - }, + extras={"score": score_name}, ) diff --git a/waveorder/calibrate/scores.py b/waveorder/calibrate/scores.py index 6c924c1b..825730d6 100644 --- a/waveorder/calibrate/scores.py +++ b/waveorder/calibrate/scores.py @@ -59,8 +59,8 @@ def whittle_evidence( 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``), - for ``tikhonov_spectral`` priors. + 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. diff --git a/waveorder/config/schema.py b/waveorder/config/schema.py index b6529ac0..db23b58b 100644 --- a/waveorder/config/schema.py +++ b/waveorder/config/schema.py @@ -273,15 +273,6 @@ def _na_below_index(self): # --- object ---------------------------------------------------------------- -class TikhonovSpectralPrior(V4BaseModel): - n_shells: PositiveInt = Field( - default=64, - description="number of isotropic |k| shells for the data-profiled iso-Whittle object " - "spectrum; the per-shell |H|^2-weighted fit (FittedSpectrum) replaces the flat Tikhonov " - "prior S_k = sigma^2 / lambda and is re-fit as the imaging system is calibrated", - ) - - class BeadsPrior(V4BaseModel): diameter: PositiveFloat = Field(description="bead diameter in um; sparse spheres at unknown positions") snr_threshold: PositiveFloat = Field( @@ -297,10 +288,6 @@ class DenoiserPrior(V4BaseModel): checkpoint: str -class PriorTikhonovSpectral(V4BaseModel): - tikhonov_spectral: TikhonovSpectralPrior - - class PriorBeads(V4BaseModel): beads: BeadsPrior @@ -311,7 +298,6 @@ class PriorDenoiser(V4BaseModel): Prior = Union[ Literal["tikhonov", "tv", "none"], - PriorTikhonovSpectral, PriorBeads, PriorDenoiser, ] @@ -641,11 +627,9 @@ def _validation_rules(self): UserWarning, ) if calibration.score == "evidence": - prior_ok = self.object.prior in ("tikhonov",) or isinstance( - self.object.prior, (PriorTikhonovSpectral, PriorBeads) - ) + 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, tikhonov_spectral, beads}") + 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: diff --git a/waveorder/distributions.py b/waveorder/distributions.py deleted file mode 100644 index 360773f1..00000000 --- a/waveorder/distributions.py +++ /dev/null @@ -1,1082 +0,0 @@ -"""Object-oriented, one-``nn.Module``-per-distribution layer over the calibrate numerics. - -WaveOrder 4.0 writes the forward model as ``g = H f + n`` with three -distributions: - -* ``p(H)`` -- the imaging system and its prior, -* ``p(f)`` -- the object prior, -* ``p(n)`` -- the noise. - -This module gives each distribution a :class:`torch.nn.Module` (or a plain -class where there is nothing to carry), mirroring Sri's design. It is a thin, -*additive* layer: every numeric here calls the existing, tested functions in -:mod:`waveorder.calibrate` and :mod:`waveorder.models.phase_thick_3d`. Nothing -reimplements the math, so the regression suite stays byte-for-byte valid. - -Distribution -> module -> config map ------------------------------------- -====================== ========================= ================================ -distribution module config block -====================== ========================= ================================ -``p(H)`` imaging system :class:`Fluorescence3D` ``imaging_system`` (fluorescence) -``p(H)`` imaging system :class:`Phase3D` ``imaging_system`` (phase) -``p(H)`` parameter prior :class:`ParameterPrior` ``imaging_system`` tolerances -``p(f)`` object prior :class:`Tikhonov` ``object.prior: tikhonov`` -``p(f)`` object prior :class:`TikhonovSpectral` ``object.prior.tikhonov_spectral`` -``p(f)`` object prior :class:`FittedSpectrum` ``object.prior`` fitted (iso-Whittle) -``p(n)`` noise :class:`Gaussian` ``noise.gaussian`` -``p(n)`` noise :class:`Poisson` ``noise.poisson`` -====================== ========================= ================================ - -A configuration value written ``{nominal, tolerance}`` becomes an -:class:`torch.nn.Parameter` (a free parameter ``P`` of the design, estimated -from the data); a plain number becomes a registered buffer. The dotted names -from :meth:`torch.nn.Module.named_parameters` are the ``parameter_names`` of a -``tf.zarr`` store (:mod:`waveorder.calibrate.tfzarr`). -""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass - -import numpy as np -import torch -from torch import Tensor, nn - -from waveorder import sampling -from waveorder.calibrate.scores import calibration_penalty, whittle_evidence -from waveorder.calibrate.synthesis import fluorescence_otf -from waveorder.calibrate.zernike_names import name_to_noll, noll_to_name -from waveorder.models import phase_thick_3d -from waveorder.reconstruct import tikhonov_regularized_inverse_filter - -__all__ = [ - "Sampling", - "ImagingSystem", - "Fluorescence3D", - "Phase3D", - "ParameterPrior", - "ObjectPrior", - "Tikhonov", - "TikhonovSpectral", - "FittedSpectrum", - "Noise", - "Gaussian", - "Poisson", - "evidence", - "calibrate", - "reconstruct", - "CalibrationResult", -] - - -# --------------------------------------------------------------------------- # -# Sampling: grid / geometry # -# --------------------------------------------------------------------------- # - - -@dataclass -class Sampling: - """Grid geometry shared by an imaging system and an object prior. - - A thin container around the pixel sizes and shape; the Nyquist helpers - delegate to :mod:`waveorder.sampling` so there is a single source of truth - for the sample-spacing formulas. - - Parameters - ---------- - zyx_shape : tuple[int, int, int] - ``(Z, Y, X)`` data shape (before axial padding). - yx_pixel_size : float - Transverse pixel size, um. - z_pixel_size : float - Axial pixel size, um. - z_padding : int - Axial padding added on each side before building the transfer - function. - """ - - zyx_shape: tuple[int, int, int] - yx_pixel_size: float - z_pixel_size: float - z_padding: int = 0 - - @property - def shape(self) -> tuple[int, int, int]: - """The unpadded ``(Z, Y, X)`` shape.""" - return tuple(self.zyx_shape) - - @property - def padded_shape(self) -> tuple[int, int, int]: - """The ``(Z + 2 z_padding, Y, X)`` shape used for transfer functions.""" - z, y, x = self.zyx_shape - return (z + 2 * self.z_padding, y, x) - - @property - def yx(self) -> tuple[int, int]: - """The transverse ``(Y, X)`` shape.""" - return tuple(self.zyx_shape[1:]) - - def transverse_nyquist(self, na_illumination: float, na_detection: float, wavelength: float) -> float: - """Transverse Nyquist spacing (see :func:`waveorder.sampling.transverse_nyquist`).""" - return sampling.transverse_nyquist(wavelength, na_illumination, na_detection) - - def axial_nyquist(self, na_detection: float, wavelength: float, index_of_refraction_media: float) -> float: - """Axial Nyquist spacing (see :func:`waveorder.sampling.axial_nyquist`).""" - return sampling.axial_nyquist(wavelength, na_detection, index_of_refraction_media) - - def shell_index(self, n_shells: int, device=None) -> tuple[Tensor, Tensor]: - """Isotropic ``|k|`` shell index per voxel, on the transfer-function grid. - - Builds the physical radial frequency ``|k|`` (cycles/um) over the - padded ``(Z, Y, X)`` grid in FFT order and bins it into ``n_shells`` - equal-width shells. Used by :class:`FittedSpectrum` to profile the - object spectrum per shell. - - Parameters - ---------- - n_shells : int - Number of equal-width ``|k|`` shells. - device : torch.device, optional - Device for the returned tensors. - - Returns - ------- - idx : Tensor - Integer shell index per voxel, shape ``padded_shape``, in ``[0, - n_shells)``. - centers : Tensor - Shell center frequencies (cycles/um), shape ``(n_shells,)``. - """ - zz, yy, xx = self.padded_shape - fz = np.fft.fftfreq(zz, d=self.z_pixel_size) - fy = np.fft.fftfreq(yy, d=self.yx_pixel_size) - fx = np.fft.fftfreq(xx, d=self.yx_pixel_size) - fr = np.sqrt(fz[:, None, None] ** 2 + fy[None, :, None] ** 2 + fx[None, None, :] ** 2) - kmax = float(fr.max()) * 1.0001 + 1e-12 - edges = np.linspace(0.0, kmax, n_shells + 1) - idx = np.clip(np.digitize(fr.ravel(), edges) - 1, 0, n_shells - 1).reshape(fr.shape) - centers = 0.5 * (edges[:-1] + edges[1:]) - return ( - torch.as_tensor(idx, dtype=torch.long, device=device), - torch.as_tensor(centers, dtype=torch.float64, device=device), - ) - - -# --------------------------------------------------------------------------- # -# Parameter bookkeeping: {nominal, tolerance} -> Parameter, number -> buffer # -# --------------------------------------------------------------------------- # - - -@dataclass -class _Free: - """A ``{nominal, tolerance}`` configuration value: becomes an nn.Parameter.""" - - nominal: float - tolerance: float | str = "free" - - -def free(nominal: float, tolerance: float | str = "free") -> _Free: - """Mark a value free for calibration: ``free(nominal, tolerance)``. - - A :class:`Fluorescence3D` or :class:`Phase3D` built with a :class:`_Free` - value registers it as an :class:`torch.nn.Parameter` (appears in - ``named_parameters``); a plain number becomes a buffer. - - Parameters - ---------- - nominal : float - Prior center. - tolerance : float or "free" - Gaussian prior standard deviation, or ``"free"`` for an improper - (flat) prior. - - Returns - ------- - _Free - """ - return _Free(float(nominal), tolerance) - - -class ImagingSystem(nn.Module): - """Base class for ``p(H)``: an imaging system with free parameters. - - Subclasses own named optical parameters as buffers (fixed numbers) or - :class:`torch.nn.Parameter` (free ``{nominal, tolerance}`` values), build - the complex transfer function ``H_k`` with :meth:`transfer`, and apply it - with :meth:`forward`. The tolerances/nominals of the free parameters are - carried alongside so a :class:`ParameterPrior` can score them. - """ - - #: Optical-parameter dtype. ``float32`` matches the function API - #: (:func:`fluorescence_otf`, :func:`phase_thick_3d.calculate_transfer_function`) - #: exactly, so :meth:`transfer` reproduces it to the bit; the scores still - #: accumulate in ``float64`` internally. - param_dtype = torch.float32 - - def __init__(self) -> None: - super().__init__() - # name -> (nominal, tolerance) for the free parameters only. - self._prior: dict[str, tuple[float, float | str]] = {} - - def _register_value(self, name: str, value: float | _Free | Tensor) -> None: - """Register ``value`` as a Parameter (if free) or a buffer (if a number).""" - if isinstance(value, _Free): - tensor = torch.tensor(float(value.nominal), dtype=self.param_dtype) - self.register_parameter(name, nn.Parameter(tensor)) - self._prior[name] = (float(value.nominal), value.tolerance) - else: - tensor = torch.as_tensor(float(value), dtype=self.param_dtype) - self.register_buffer(name, tensor) - - def free_parameter_names(self) -> list[str]: - """Dotted names of the free parameters (``tf.zarr`` ``parameter_names``).""" - return [name for name, _ in self.named_parameters()] - - def parameter_values(self) -> dict[str, Tensor]: - """Current values of the free parameters by dotted name.""" - return {name: param for name, param in self.named_parameters()} - - def nominals(self) -> dict[str, float]: - """Prior centers of the free parameters by dotted name.""" - return {name: nominal for name, (nominal, _) in self._prior.items()} - - def tolerances(self) -> dict[str, float | str]: - """Prior tolerances of the free parameters by dotted name.""" - return {name: tolerance for name, (_, tolerance) in self._prior.items()} - - def transfer(self, s: Sampling) -> Tensor: # pragma: no cover - abstract - """Complex transfer function ``H_k`` in FFT order.""" - raise NotImplementedError - - def prepare(self, raw: Tensor, s: Sampling) -> Tensor: - """Turn a raw stack into the modeled data ``g`` (identity by default).""" - return torch.as_tensor(raw) - - def forward(self, f: Tensor, s: Sampling) -> Tensor: - """Apply the transfer function: ``g = H f`` (real part of the IFFT).""" - otf = self.transfer(s) - f_fft = torch.fft.fftn(torch.as_tensor(f, dtype=otf.dtype if otf.is_complex() else torch.complex64)) - return torch.real(torch.fft.ifftn(f_fft * otf)) - - -class Fluorescence3D(ImagingSystem): - """Widefield fluorescence ``p(H)``; wraps :func:`fluorescence_otf`. - - Owns the detection NA, named Zernike coefficients, focus offset, emission - wavelength, medium index and pixel sizes. Any value passed as - :func:`free` becomes a calibratable :class:`torch.nn.Parameter`; a plain - number becomes a buffer. :meth:`transfer` calls :func:`fluorescence_otf` - with the current values, so the OTF is identical to the function API. - - Parameters - ---------- - wavelength_emission : float - Emission wavelength, um. - index_of_refraction_media : float - Medium refractive index. - numerical_aperture_detection : float or _Free - Detection NA. - zernike_coefficients : mapping, optional - ``{name_or_noll: value}`` in waves; each value a number (buffer) or a - :func:`free` value (Parameter). The dotted parameter name is - ``pupil_aberrations.`` (e.g. - ``pupil_aberrations.astigmatism_0``). - z_focus_offset : float or _Free - Focus offset in slices. - normalize : {"peak", "none"} - OTF normalization, passed to :func:`fluorescence_otf`. - """ - - def __init__( - self, - wavelength_emission: float, - index_of_refraction_media: float, - numerical_aperture_detection: float | _Free, - zernike_coefficients: Mapping[str | int, float | _Free] | None = None, - z_focus_offset: float | _Free = 0.0, - normalize: str = "peak", - ) -> None: - super().__init__() - self.wavelength_emission = float(wavelength_emission) - self.index_of_refraction_media = float(index_of_refraction_media) - self.normalize = normalize - - self._register_value("numerical_aperture_detection", numerical_aperture_detection) - self._register_value("z_focus_offset", z_focus_offset) - - # Zernike coefficients live under a dotted "pupil_aberrations." prefix - # so their free names round-trip through tf.zarr parameter_names. - self._zernike_nolls: list[int] = [] - self.pupil_aberrations = nn.Module() - aberrations_prior: dict[str, tuple[float, float | str]] = {} - for key, value in (zernike_coefficients or {}).items(): - noll = key if isinstance(key, int) else name_to_noll(key) - canonical = noll_to_name(noll) - self._zernike_nolls.append(noll) - if isinstance(value, _Free): - tensor = torch.tensor(float(value.nominal), dtype=self.param_dtype) - self.pupil_aberrations.register_parameter(canonical, nn.Parameter(tensor)) - aberrations_prior[f"pupil_aberrations.{canonical}"] = (float(value.nominal), value.tolerance) - else: - tensor = torch.as_tensor(float(value), dtype=self.param_dtype) - self.pupil_aberrations.register_buffer(canonical, tensor) - self._prior.update(aberrations_prior) - - def _coefficients(self) -> dict[int, float | Tensor]: - """Current Zernike coefficients keyed by Noll index.""" - out: dict[int, float | Tensor] = {} - for noll in self._zernike_nolls: - canonical = noll_to_name(noll) - value = getattr(self.pupil_aberrations, canonical) - out[noll] = value - return out - - def transfer(self, s: Sampling) -> Tensor: - """Complex OTF from :func:`fluorescence_otf` at the current parameters.""" - return fluorescence_otf( - zyx_shape=s.padded_shape, - yx_pixel_size=s.yx_pixel_size, - z_pixel_size=s.z_pixel_size, - wavelength_emission=self.wavelength_emission, - index_of_refraction_media=self.index_of_refraction_media, - numerical_aperture_detection=self.numerical_aperture_detection, - zernike_coefficients=self._coefficients(), - z_focus_offset=self.z_focus_offset, - normalize=self.normalize, - ) - - -class Phase3D(ImagingSystem): - """Thick-phase ``p(H)``; wraps :func:`phase_thick_3d.calculate_transfer_function`. - - The weak-object model has two transfer functions: the real (phase) WOTF - and the imaginary (absorption) WOTF, both returned by WaveOrder as complex - tensors. :meth:`transfer_real_imag` exposes the pair directly for the - label-free reconstruction path. :meth:`transfer` returns the phase WOTF as - the single ``H_k`` (the usual pure-phase case; combine with the absorption - WOTF via an ``absorption_ratio`` at reconstruction time, as - :func:`phase_thick_3d.apply_inverse_transfer_function` does). - - Parameters - ---------- - wavelength_illumination : float - Illumination wavelength, um. - index_of_refraction_media : float - Medium refractive index. - numerical_aperture_illumination : float or _Free - Condenser NA. - numerical_aperture_detection : float or _Free - Detection NA. - tilt_angle_zenith, tilt_angle_azimuth : float or _Free - Illumination tilt angles, radians. - invert_phase_contrast : bool - Flip the defocus axis (matches 3.x). - pupil_steepness : float - Sigmoid pupil-edge steepness. - """ - - def __init__( - self, - wavelength_illumination: float, - index_of_refraction_media: float, - numerical_aperture_illumination: float | _Free = 0.9, - numerical_aperture_detection: float | _Free = 1.2, - tilt_angle_zenith: float | _Free = 0.0, - tilt_angle_azimuth: float | _Free = 0.0, - invert_phase_contrast: bool = False, - pupil_steepness: float = 1e4, - ) -> None: - super().__init__() - self.wavelength_illumination = float(wavelength_illumination) - self.index_of_refraction_media = float(index_of_refraction_media) - self.invert_phase_contrast = bool(invert_phase_contrast) - self.pupil_steepness = float(pupil_steepness) - - self._register_value("numerical_aperture_illumination", numerical_aperture_illumination) - self._register_value("numerical_aperture_detection", numerical_aperture_detection) - self._register_value("tilt_angle_zenith", tilt_angle_zenith) - self._register_value("tilt_angle_azimuth", tilt_angle_azimuth) - - def transfer_real_imag(self, s: Sampling) -> tuple[Tensor, Tensor]: - """Real and imaginary potential WOTFs from the phase model. - - The optical parameters are ``float32`` (``param_dtype``), matching the - model's internal cast, so the returned WOTFs reproduce the function - API exactly while preserving autograd through the free parameters. - """ - return phase_thick_3d.calculate_transfer_function( - zyx_shape=s.shape, - yx_pixel_size=s.yx_pixel_size, - z_pixel_size=s.z_pixel_size, - wavelength_illumination=self.wavelength_illumination, - z_padding=s.z_padding, - index_of_refraction_media=self.index_of_refraction_media, - numerical_aperture_illumination=self.numerical_aperture_illumination, - numerical_aperture_detection=self.numerical_aperture_detection, - invert_phase_contrast=self.invert_phase_contrast, - tilt_angle_zenith=self.tilt_angle_zenith, - tilt_angle_azimuth=self.tilt_angle_azimuth, - pupil_steepness=self.pupil_steepness, - ) - - def transfer(self, s: Sampling) -> Tensor: - """The phase (real-potential) WOTF ``H_k`` for the pure-phase case.""" - real_tf, _ = self.transfer_real_imag(s) - return real_tf - - -# --------------------------------------------------------------------------- # -# ParameterPrior: C(H) = -log p(H) # -# --------------------------------------------------------------------------- # - - -class ParameterPrior: - """The calibration penalty ``C(H) = -log p(H)``; wraps :func:`calibration_penalty`. - - Parameters - ---------- - nominals : dict[str, float] - Prior centers by parameter name. - tolerances : dict[str, float or "free"] - Gaussian prior standard deviations by parameter name (or ``"free"``). - """ - - def __init__(self, nominals: Mapping[str, float], tolerances: Mapping[str, float | str]) -> None: - self.nominals = dict(nominals) - self.tolerances = dict(tolerances) - - @classmethod - def from_system(cls, system: ImagingSystem) -> "ParameterPrior": - """Build the prior from an imaging system's free-parameter nominals/tolerances.""" - return cls(system.nominals(), system.tolerances()) - - def nll(self, params: Mapping[str, float | Tensor]) -> Tensor: - """Negative log prior ``C = sum_i (h_i - h_i0)^2 / (2 tau_i^2)``.""" - return calibration_penalty(dict(params), self.nominals, self.tolerances) - - def sample(self, generator: torch.Generator | None = None) -> dict[str, float]: - """Draw a parameter dict from the Gaussian prior (``free`` -> nominal).""" - out: dict[str, float] = {} - for name, nominal in self.nominals.items(): - tolerance = self.tolerances.get(name, "free") - if tolerance == "free": - out[name] = float(nominal) - else: - noise = torch.randn((), generator=generator, dtype=torch.float64) - out[name] = float(nominal) + float(tolerance) * float(noise) - return out - - -# --------------------------------------------------------------------------- # -# ObjectPrior: p(f) -> S_k # -# --------------------------------------------------------------------------- # - - -class ObjectPrior(nn.Module): - """Base class for ``p(f)``: returns the object power spectrum ``S_k``. - - Priors come in two flavours. A *fixed* prior (``Tikhonov``, - ``TikhonovSpectral``) ignores the data and returns ``S_k`` from stored - parameters. A *data-profiled* (empirical-Bayes) prior (``FittedSpectrum``, - the winning iso-Whittle method) estimates ``S_k`` from the data - periodogram and the current transfer function; :func:`evidence` passes - ``data_fft`` and ``otf`` through for those. The optional keywords keep the - fixed priors' call sites unchanged. - """ - - #: True for priors whose ``spectrum`` consumes ``data_fft``/``otf``. - data_profiled: bool = False - - def spectrum( - self, - s: Sampling, - noise: "Noise", - *, - data_fft: Tensor | None = None, - otf: Tensor | None = None, - ) -> Tensor | None: # pragma: no cover - abstract - """Object power spectrum ``S_k`` (same shape as ``H_k``), or ``None``.""" - raise NotImplementedError - - -class Tikhonov(ObjectPrior): - """Flat Tikhonov prior ``S_k = sigma^2 / lambda`` (the current default). - - Parameters - ---------- - weight : float - Regularization weight ``lambda``. - """ - - def __init__(self, weight: float) -> None: - super().__init__() - self.register_buffer("weight", torch.as_tensor(float(weight), dtype=torch.float64)) - - def spectrum(self, s: Sampling, noise: "Noise", **_: Tensor) -> Tensor: - """``S_k = sigma^2 / lambda`` (a scalar; broadcast against ``H_k``).""" - return noise.variance() / self.weight - - -class TikhonovSpectral(ObjectPrior): - """Explicit per-frequency object spectrum ``S_k``. - - Parameters - ---------- - spectrum : Tensor - The object power spectrum ``S_k`` in FFT order, same shape as ``H_k``. - """ - - def __init__(self, spectrum: Tensor) -> None: - super().__init__() - self.register_buffer("spectrum_buffer", torch.as_tensor(spectrum, dtype=torch.float64)) - - def spectrum(self, s: Sampling, noise: "Noise", **_: Tensor) -> Tensor: - """Return the stored ``S_k``.""" - return self.spectrum_buffer - - -class FittedSpectrum(ObjectPrior): - """Isotropic shell-fit object spectrum ``S(|k|)`` estimated from the data. - - This is the winning "iso-Whittle" object prior of the 2026-10-30 blind - sprint, ported onto the module interface. It is a data-profiled - (empirical-Bayes) prior: instead of a fixed ``S_k``, it estimates the - object power spectrum per isotropic frequency shell directly from the - data periodogram and the current transfer function, which is what beat a - flat Tikhonov prior on real blind data (see the theory page and - ``winning_method_spec.md``). - - Per shell of ``|k|``, with periodogram ``P_k = |FFT(g)|^2`` (ortho), - transfer magnitude ``|H_k|^2`` and noise floor ``sigma^2``, the fit is the - threshold-free, ``|H|^2``-weighted least squares estimate - - S_shell = sum_k |H_k|^2 (P_k - sigma^2) / sum_k |H_k|^4 (clamped >= 0) - - (``whittle3d._wls_shell_S``). The weighting by ``|H|^2`` is what makes it - threshold-free: a fixed ``|H|^2`` cutoff would turn the evidence into an - OTF-support indicator (the round-1 failure). The shell profile is then - broadcast back to every frequency. Being a positive, radially-binned - spectrum, ``S_k`` keeps the Gaussian evidence exact. - - Parameters - ---------- - n_shells : int - Number of isotropic ``|k|`` shells. - min_power : float - Floor on the per-shell spectrum (keeps ``S_k > 0`` so ``log v_k`` and - the Wiener inverse stay finite). - """ - - data_profiled = True - - def __init__(self, n_shells: int = 64, min_power: float = 1e-10) -> None: - super().__init__() - self.n_shells = int(n_shells) - self.min_power = float(min_power) - - def spectrum( - self, - s: Sampling, - noise: "Noise", - *, - data_fft: Tensor | None = None, - otf: Tensor | None = None, - ) -> Tensor: - """Data-profiled ``S_k`` via an ``|H|^2``-weighted shell fit. - - Parameters - ---------- - s : Sampling - Grid geometry (supplies the physical ``|k|`` grid). - noise : Noise - Supplies ``sigma^2``. - data_fft : Tensor - ``fftn(g, norm="ortho")`` of the prepared data (``evidence`` passes - this). Required. - otf : Tensor - Current transfer function ``H_k`` (``evidence`` passes this). - Required. - """ - if data_fft is None or otf is None: - raise ValueError( - "FittedSpectrum needs data_fft and otf; call it through distributions.evidence (which passes both)." - ) - sigma2 = noise.variance().to(torch.float64) - power = (data_fft.abs() ** 2).to(torch.float64) - h2 = (otf.abs() ** 2).to(torch.float64) - - shell, _centers = s.shell_index(self.n_shells, device=power.device) - shell = shell.reshape(-1) - p = power.reshape(-1) - h = h2.reshape(-1) - num = torch.zeros(self.n_shells, dtype=torch.float64, device=p.device) - den = torch.zeros(self.n_shells, dtype=torch.float64, device=p.device) - num.scatter_add_(0, shell, h * torch.clamp(p - sigma2, min=0.0)) - den.scatter_add_(0, shell, h * h) - s_shell = (num / den.clamp_min(1e-30)).clamp_min(self.min_power) - return s_shell[shell].reshape(power.shape) - - -# --------------------------------------------------------------------------- # -# Noise: p(n) # -# --------------------------------------------------------------------------- # - - -class Noise(nn.Module): - """Base class for ``p(n)``.""" - - def variance(self) -> Tensor: # pragma: no cover - abstract - """Noise variance ``sigma^2``.""" - raise NotImplementedError - - def nll(self, g: Tensor, mean: Tensor) -> Tensor: # pragma: no cover - abstract - """Negative log likelihood of ``g`` given its mean (the data misfit ``D``).""" - raise NotImplementedError - - def sample(self, mean: Tensor, generator: torch.Generator | None = None) -> Tensor: # pragma: no cover - """Draw a noisy observation around ``mean``.""" - raise NotImplementedError - - -class Gaussian(Noise): - """Additive white Gaussian noise ``n ~ N(0, sigma^2 I)``. - - Parameters - ---------- - sigma : float - Noise standard deviation in data units. - """ - - def __init__(self, sigma: float) -> None: - super().__init__() - self.register_buffer("sigma", torch.as_tensor(float(sigma), dtype=torch.float64)) - - def variance(self) -> Tensor: - """``sigma^2``.""" - return self.sigma**2 - - def nll(self, g: Tensor, mean: Tensor) -> Tensor: - """Gaussian data misfit ``||g - mean||^2 / (2 sigma^2) + const``.""" - residual = torch.as_tensor(g, dtype=torch.float64) - torch.as_tensor(mean, dtype=torch.float64) - return (residual**2).sum() / (2.0 * self.variance()) - - def sample(self, mean: Tensor, generator: torch.Generator | None = None) -> Tensor: - """``mean + sigma * N(0, 1)``.""" - mean = torch.as_tensor(mean, dtype=torch.float64) - noise = torch.randn(mean.shape, generator=generator, dtype=torch.float64) - return mean + self.sigma * noise - - -class Poisson(Noise): - """Poisson (shot) noise with a linear gain. - - TODO (next pass): the full Poisson evidence / Richardson-Lucy path. For - now ``variance`` returns the gain so that a Gaussian-approximate evidence - can still be scored, and ``sample`` draws true Poisson counts for the - simulator. - - Parameters - ---------- - gain : float - Electrons (or ADU) per photon. - offset : float - Constant dark offset in data units. - """ - - def __init__(self, gain: float = 1.0, offset: float = 0.0) -> None: - super().__init__() - self.register_buffer("gain", torch.as_tensor(float(gain), dtype=torch.float64)) - self.register_buffer("offset", torch.as_tensor(float(offset), dtype=torch.float64)) - - def variance(self) -> Tensor: - """Gaussian-approximate variance (the gain); a true shot-noise map is per-voxel.""" - return self.gain - - def nll(self, g: Tensor, mean: Tensor) -> Tensor: - """Poisson negative log likelihood ``sum(mean - g log mean)`` (up to const).""" - g = torch.as_tensor(g, dtype=torch.float64) - rate = (torch.as_tensor(mean, dtype=torch.float64) - self.offset) / self.gain - rate = rate.clamp(min=1e-12) - return (rate - (g / self.gain) * torch.log(rate)).sum() - - def sample(self, mean: Tensor, generator: torch.Generator | None = None) -> Tensor: - """Draw Poisson counts ``offset + gain * Poisson((mean - offset) / gain)``.""" - rate = (torch.as_tensor(mean, dtype=torch.float64) - self.offset) / self.gain - counts = torch.poisson(rate.clamp(min=0.0), generator=generator) - return self.offset + self.gain * counts - - -# --------------------------------------------------------------------------- # -# evidence / calibrate / reconstruct # -# --------------------------------------------------------------------------- # - - -def evidence( - g: Tensor, - system: ImagingSystem, - obj: ObjectPrior, - noise: Noise, - s: Sampling, - *, - exclude_dc: bool = False, - data_fft: Tensor | None = None, -) -> Tensor: - """Negative log evidence ``-log p(g | H)`` of the imaging system. - - Exactly Sri's snippet: per frequency ``v_k = |H_k|^2 S_k + sigma^2`` and - the Whittle NLL. It calls :func:`whittle_evidence` with the module-derived - ``H_k``, ``S_k`` and ``sigma`` so the value matches the function API. - - Parameters - ---------- - g : Tensor - Observed stack ``(Z, Y, X)``. - system : ImagingSystem - ``p(H)``; supplies ``H_k`` via :meth:`ImagingSystem.transfer`. - obj : ObjectPrior - ``p(f)``; supplies ``S_k`` via :meth:`ObjectPrior.spectrum`. - noise : Noise - ``p(n)``; supplies ``sigma``. - s : Sampling - Grid geometry. - exclude_dc : bool - Drop the ``k = 0`` term (for real data with a large background). - data_fft : Tensor, optional - Precomputed ``fftn(g, norm="ortho")``. - - Returns - ------- - Tensor - Scalar negative log evidence. - """ - otf = system.transfer(s) - # data-profiled priors (FittedSpectrum) need the periodogram and H_k; fixed - # priors (Tikhonov, TikhonovSpectral) ignore these kwargs. - if getattr(obj, "data_profiled", False) and data_fft is None: - data_fft = torch.fft.fftn(g, norm="ortho") - s_k = obj.spectrum(s, noise, data_fft=data_fft, otf=otf) - sigma = torch.sqrt(noise.variance()) - return whittle_evidence( - data=None if data_fft is not None else g, - data_fft=data_fft, - otf=otf, - noise_sigma=sigma, - object_spectrum=s_k, - exclude_dc=exclude_dc, - ) - - -def _laplace_covariance(objective, params: list[nn.Parameter]) -> Tensor: - """Inverse Hessian of ``objective`` w.r.t. ``params`` by double backprop. - - Parameters - ---------- - objective : callable - Scalar objective evaluated at the current parameter values. - params : list of nn.Parameter - The free parameters; the Hessian is taken over their flattened stack. - - Returns - ------- - Tensor - ``[P, P]`` Laplace covariance (inverse Hessian at the optimum). - """ - value = objective() - grads = torch.autograd.grad(value, params, create_graph=True) - flat_grad = torch.cat([g.reshape(-1) for g in grads]) - rows = [] - for i in range(flat_grad.numel()): - second = torch.autograd.grad(flat_grad[i], params, retain_graph=True) - rows.append(torch.cat([g.reshape(-1) for g in second])) - hessian = torch.stack(rows).to(torch.float64) - return torch.linalg.inv(hessian).detach() - - -@dataclass -class CalibrationResult: - """Output of :func:`calibrate`.""" - - parameters: dict[str, float] - parameter_names: list[str] - covariance: np.ndarray - score_trace: list[float] - best_score: float - - -def calibrate( - g: Tensor, - system: ImagingSystem, - prior: ParameterPrior, - obj: ObjectPrior, - noise: Noise, - s: Sampling, - *, - score=evidence, - max_iter: int = 200, - lr: float = 0.01, - exclude_dc: bool = False, -) -> CalibrationResult: - """Estimate ``H``'s free parameters by minimizing ``score + prior.nll``. - - The outer argmin over ``system``'s :meth:`torch.nn.Module.named_parameters` - (the design's free params ``P``), by Adam, with the Laplace covariance - (inverse Hessian of the total objective at the optimum). - - Parameters - ---------- - g : Tensor - Observed stack. - system : ImagingSystem - The imaging system; its free parameters are optimized in place. - prior : ParameterPrior - ``C(H)``, added to the data score. - obj : ObjectPrior - ``p(f)``. - noise : Noise - ``p(n)``. - s : Sampling - Grid geometry. - score : callable - Data score, ``score(g, system, obj, noise, s, exclude_dc=...)``; - defaults to :func:`evidence`. - max_iter : int - Adam iterations. - lr : float - Adam step size. - exclude_dc : bool - Passed to the score. - - Returns - ------- - CalibrationResult - """ - names = system.free_parameter_names() - params = [param for _, param in system.named_parameters()] - if not params: - raise ValueError("system has no free parameters to calibrate; mark values with free(...)") - - g = torch.as_tensor(g, dtype=torch.float64) - data_fft = torch.fft.fftn(g, norm="ortho") - - def objective() -> Tensor: - data_term = score(g, system, obj, noise, s, exclude_dc=exclude_dc, data_fft=data_fft) - return data_term + prior.nll(system.parameter_values()) - - optimizer = torch.optim.Adam(params, lr=lr) - score_trace: list[float] = [] - for _ in range(max_iter): - optimizer.zero_grad() - value = objective() - value.backward() - optimizer.step() - score_trace.append(float(value.detach())) - - with torch.no_grad(): - best_score = float(objective()) - estimated = {name: float(param.detach()) for name, param in zip(names, params)} - - # Laplace covariance: inverse Hessian of the total objective at the - # optimum. Build it by double backprop through the live parameters so the - # graph stays intact (a copy_/no_grad probe would detach it). - try: - covariance = _laplace_covariance(objective, params).numpy() - except RuntimeError: - covariance = np.full((len(params), len(params)), np.nan) - - return CalibrationResult( - parameters=estimated, - parameter_names=names, - covariance=covariance, - score_trace=score_trace, - best_score=best_score, - ) - - -def reconstruct( - g: Tensor, - system: ImagingSystem, - obj: ObjectPrior, - noise: Noise, - s: Sampling, -) -> Tensor: - """Inner argmin over ``f``: the closed-form Wiener estimate. - - For Gaussian noise and a Tikhonov object prior this is exactly the - Wiener filter ``f_hat = H* / (|H|^2 + lambda) . g``, identical to - :func:`waveorder.reconstruct.tikhonov_regularized_inverse_filter` applied - to ``g`` with ``lambda = sigma^2 / S_k``. - - Parameters - ---------- - g : Tensor - Observed stack ``(Z, Y, X)``. - system : ImagingSystem - Supplies ``H_k``. - obj : ObjectPrior - Supplies ``S_k`` (hence ``lambda = sigma^2 / S_k``). - noise : Noise - Supplies ``sigma^2``. - s : Sampling - Grid geometry. - - Returns - ------- - Tensor - Real-space object estimate ``f_hat``, same shape as ``g``. - - Raises - ------ - NotImplementedError - For a non-Gaussian noise or an improper (``None``) object spectrum; - the iterative / Richardson-Lucy path lands in a later pass. - """ - if not isinstance(noise, Gaussian): - raise NotImplementedError("reconstruct's closed form is Gaussian-only for now; RL path follows") - s_k = obj.spectrum(s, noise) - if s_k is None: - raise NotImplementedError("improper object prior (S_k=None) has no closed-form Wiener estimate") - - otf = system.transfer(s) - g = torch.as_tensor(g) - if g.shape[-3:] != otf.shape[-3:]: - raise ValueError( - f"g shape {tuple(g.shape)} does not match the transfer function {tuple(otf.shape)}; " - f"z-padded reconstruction is not handled by the closed form yet " - f"(use s.z_padding == 0, or pad g to the padded_shape)" - ) - regularization = noise.variance() / s_k - inverse_filter = tikhonov_regularized_inverse_filter(otf, regularization) - g = g.to(torch.complex64) - g_fft = torch.fft.fftn(g) - return torch.real(torch.fft.ifftn(g_fft * inverse_filter)) - - -# --------------------------------------------------------------------------- # -# Config adapter # -# --------------------------------------------------------------------------- # - - -def _value_or_free(value): - """Map a config float or ``FreeParameter`` to a plain number or :func:`free`.""" - from waveorder.config.schema import FreeParameter - - if isinstance(value, FreeParameter): - return free(value.nominal, value.tolerance) - return float(value) - - -def build_system_from_config(cfg): - """Build ``(system, obj, noise, prior, sampling)`` from a 4.0 config. - - A thin adapter from the ``imaging_system`` / ``object`` / ``noise`` / - ``solver`` blocks of a :class:`waveorder.config.schema.V4Config` to the - distribution modules of this file. A config value ``{nominal, tolerance}`` - (a :class:`waveorder.config.schema.FreeParameter`) becomes an - :class:`torch.nn.Parameter`; a plain number becomes a buffer. - - Parameters - ---------- - cfg : V4Config - The configuration. Shift-invariant fluorescence or phase contrast. - - Returns - ------- - system : ImagingSystem - obj : ObjectPrior - noise : Noise - prior : ParameterPrior - sampling : Sampling - - Raises - ------ - NotImplementedError - For tiled field dependence, non-gaussian noise with a spectral - reconstruction, or object priors without a closed-form spectrum. These - follow in later passes; the evidence path already restricts to the - supported subset. - """ - from waveorder.config.schema import ( - PriorTikhonovSpectral, - TikhonovSpectralPrior, - ) - - imaging = cfg.imaging_system - if imaging.field_dependence != "shift_invariant": - raise NotImplementedError("build_system_from_config is shift-invariant for now; tiled follows") - - yx = cfg.data.pixel_size.yx - yx_pixel_size = float(yx) if isinstance(yx, (int, float)) else float(yx.y) - zyx_shape_source = cfg.simulate.shape if cfg.simulate is not None else None - sampling_geometry = Sampling( - zyx_shape=tuple(zyx_shape_source) if zyx_shape_source is not None else (0, 0, 0), - yx_pixel_size=yx_pixel_size, - z_pixel_size=float(cfg.data.pixel_size.z), - z_padding=int(cfg.solver.padding.z), - ) - - detection = imaging.detection - if imaging.contrast == "fluorescence": - if detection.wavelength is None: - raise ValueError("fluorescence imaging_system.detection.wavelength is required") - zernike_coefficients: dict[str | int, object] = {} - if detection.pupil_aberrations is not None: - if detection.pupil_aberrations.zernike_up_to_order is not None: - from waveorder.calibrate.zernike_names import noll_indices_up_to_order - - shorthand = detection.pupil_aberrations.zernike_up_to_order - for noll in noll_indices_up_to_order(shorthand.order): - zernike_coefficients[noll] = free(0.0, shorthand.tolerance) - for name, value in detection.pupil_aberrations.terms.items(): - zernike_coefficients[name] = _value_or_free(value) - system: ImagingSystem = Fluorescence3D( - wavelength_emission=detection.wavelength, - index_of_refraction_media=imaging.index_of_refraction_media, - numerical_aperture_detection=_value_or_free(detection.numerical_aperture), - zernike_coefficients=zernike_coefficients, - z_focus_offset=_value_or_free(detection.z_focus_offset), - ) - elif imaging.contrast == "phase": - illumination = imaging.illumination - if illumination.wavelength is None: - raise ValueError("phase imaging_system.illumination.wavelength is required") - na_ill = illumination.numerical_aperture - system = Phase3D( - wavelength_illumination=illumination.wavelength, - index_of_refraction_media=imaging.index_of_refraction_media, - numerical_aperture_illumination=_value_or_free(na_ill) if na_ill is not None else 0.9, - numerical_aperture_detection=_value_or_free(detection.numerical_aperture), - tilt_angle_zenith=_value_or_free(illumination.tilt.zenith), - tilt_angle_azimuth=_value_or_free(illumination.tilt.azimuth), - invert_phase_contrast=illumination.invert_phase_contrast, - ) - else: # pragma: no cover - schema restricts contrast - raise NotImplementedError(f"contrast {imaging.contrast!r} is not supported") - - # Noise. - if cfg.noise.kind == "gaussian": - gaussian = cfg.noise.gaussian - sigma = gaussian.sigma if isinstance(gaussian.sigma, (int, float)) else 1.0 - noise: Noise = Gaussian(float(sigma)) - elif cfg.noise.kind == "poisson": - poisson = cfg.noise.poisson - noise = Poisson(gain=poisson.gain, offset=poisson.offset) - else: - raise NotImplementedError(f"noise kind {cfg.noise.kind!r} is not supported by the builder yet") - - # Object prior. - prior_spec = cfg.object.prior - weight = cfg.object.weight if isinstance(cfg.object.weight, (int, float)) else 1e-3 - if isinstance(prior_spec, PriorTikhonovSpectral): - # Data-profiled iso-Whittle spectrum (same prior the evidence path uses). - obj: ObjectPrior = FittedSpectrum(n_shells=prior_spec.tikhonov_spectral.n_shells) - elif isinstance(prior_spec, TikhonovSpectralPrior): - obj = FittedSpectrum(n_shells=prior_spec.n_shells) - elif prior_spec in ("tikhonov", "tv", "none") or prior_spec is None: - obj = Tikhonov(float(weight)) - else: - obj = Tikhonov(float(weight)) - - prior = ParameterPrior.from_system(system) - return system, obj, noise, prior, sampling_geometry