Optimization infrastructure: pluggable backends, losses, batched GPU processing (#531)

* add pluggable optimizer backends (lbfgs, nelder_mead, grid_search), convergence criteria, and `use_gradients` flag

* rename num_iterations to max_iterations

* add CacheSpec and disk-backed TransferFunctionCache for multi-tile TF reuse

* add benchmark_optimizers for comparing optimization strategies"

* add optimizer and TF cache benchmark examples"

* add pluggable, normalized loss functions with 3D support

* fix widget tests: use model_validator for loss config instead of discriminated union

* Fix: import PrintLogger in optimize.py for nelder_mead/grid_search defaults

* add batched (B,Z,Y,X) reconstruction support to all 4 models

Accept both (Z,Y,X) and (B,Z,Y,X) inputs following PyTorch convention.
Internal code unsqueezes/squeezes at boundaries so all operations use
(B,Z,Y,X). Unbatched path is bit-exact with prior behavior.

- util: pad_zyx_along_z, inten_normalization, inten_normalization_3D
- isotropic_thin_3d: calculate_transfer_function (batched tilt angles),
  calculate_singular_system, apply_inverse_transfer_function, reconstruct
- isotropic_fluorescent_thin_3d: apply_inverse_transfer_function, reconstruct
- phase_thick_3d: apply_inverse_transfer_function, reconstruct
- isotropic_fluorescent_thick_3d: apply_inverse_transfer_function, reconstruct
- optics: compute_weak_object_transfer_function_2d now supports broadcasting
- optics: generate_tilted_pupil documents batched output shape

* add batched API input and per-tile optimization support

API layer: phase and fluorescence apply_inverse_transfer_function accept
list[xr.DataArray] for batched processing. Adds _wrap_output_tensor
helper to reduce wrapping duplication.

Optimizer: when data.ndim == 4, parameters become (B,) tensors for
independent per-tile optimization. Standard Adam with (B,) tensors
gives B independent optimizers (verified bit-exact vs sequential).
Loss is summed per-tile; cross-tile gradient terms are zero since
tiles are decoupled in the forward pass.

* add batched reconstruction tests for all 4 models

Tests B=4 batched vs sequential bit-exactness, B=1 matches unbatched,
and per-tile tilt angles for isotropic_thin_3d.

* restore torch.linalg.svd in calculate_singular_system

The norm-based decomposition (U=identity) introduced a regression
vs waveorder 3.0.0 by ignoring cross-channel coupling between
absorption and phase transfer functions. Restoring the real SVD
reduces the mean reconstruction error from 0.077 to 0.0005 on
OPS data (verified against production phenotyping_phase_2d.zarr).

torch.linalg.svd supports backpropagation through complex tensors,
so gradient-based optimization still works.

* add use_svd flag to calculate_singular_system for gradient compatibility

torch.linalg.svd backward fails with complex tensors on some GPU types
due to singular vector phase ambiguity. When gradients are needed
(optimization), reconstruct() auto-selects the norm-based decomposition
(use_svd=False) which supports backpropagation on all devices. When no
gradients are needed (final reconstruction), full SVD is used for best
accuracy.

The no-grad path remains bit-exact with the previous SVD-only behavior.

* catch SVD/NaN errors in optimization loop and revert to last good params

* support batched parameters in calculate_transfer_function and apply_inverse_transfer_function

* Remove TF caching and benchmarking harness

Removes cache.py, benchmark.py, CacheSpec, and associated tests/examples.
TF caching showed marginal speedup with ~24% accuracy loss (nearest-neighbor)
and the benchmarking harness is better suited as standalone scripts.

* Remove CacheSpec from OptimizableFloat

* Add NAdam optimizer backend

* Validate optimization method name

* Use default lr=1.0 for L-BFGS line search step size

* Eliminate L-BFGS double forward pass

* Use np.linspace for grid search to avoid float-step accumulation

Author: talonchandler

.gitignore

@@ -155,3 +155,4 @@ waveorder/_version.py
 # example data
 /examples/data_temp/*
 /logs/*
+scripts/

docs/examples/optimization/optimize_fluorescence_2d.py

@@ -8,6 +8,7 @@
 
 from waveorder.api import fluorescence
 from waveorder.optim import OptimizableFloat
+from waveorder.optim.losses import MidbandPowerLossSettings
 
 # Ground truth parameters
 gt_z_offset = 0.6
@@ -32,8 +33,8 @@
 optimized_settings, recon = fluorescence.optimize(
     data,
     settings=opt_settings,
-    num_iterations=50,
-    midband_fractions=(0.01, 0.5),
+    max_iterations=50,
+    loss_settings=MidbandPowerLossSettings(midband_fractions=[0.01, 0.5]),
     log_dir=log_dir,
     log_images=True,
 )

docs/examples/optimization/optimize_phase_2d.py

@@ -11,6 +11,7 @@
 
 from waveorder.api import phase
 from waveorder.optim import OptimizableFloat
+from waveorder.optim.losses import MidbandPowerLossSettings
 
 # To use your own data instead of simulated data, create a CZYX xr.DataArray:
 #
@@ -53,8 +54,8 @@
 optimized_settings, recon = phase.optimize(
     data,
     settings=opt_settings,
-    num_iterations=50,
-    midband_fractions=(0.1, 0.5),
+    max_iterations=50,
+    loss_settings=MidbandPowerLossSettings(midband_fractions=[0.1, 0.5]),
     log_dir=log_dir,
     log_images=True,
 )

docs/examples/optimization/phase_2d_optimized.yml

@@ -24,8 +24,9 @@ phase:
     reconstruction_algorithm: Tikhonov
     regularization_strength: 0.01
 optimization:
-  num_iterations: 10
+  max_iterations: 10
+  method: adam                          # adam, lbfgs, nelder_mead, grid_search
   loss:
-    type: midband_power
-    midband_fractions: [0.125, 0.25]
+    type: midband_power                 # midband_power, total_variation, laplacian_variance,
+    midband_fractions: [0.125, 0.25]    # normalized_variance, spectral_flatness
   log_dir: ./runs/phase_optim

tests/cli_tests/test_reconstruct.py

@@ -270,7 +270,7 @@ def test_optimization_cli(tmp_path):
         phase=settings.PhaseSettings(
             transfer_function=PhaseTFSettings(z_focus_offset={"init": 0, "lr": 0.1}),
         ),
-        optimization=OptimizationSettings(num_iterations=2),
+        optimization=OptimizationSettings(max_iterations=2),
     )
     config_path = tmp_path / "optim.yml"
     utils.model_to_yaml(recon_settings, config_path)

tests/models/test_batched_reconstruction.py

@@ -0,0 +1,238 @@
+"""Test batched reconstruction is bit-exact against sequential single-tile calls."""
+
+import pytest
+import torch
+
+from waveorder.models import (
+    isotropic_fluorescent_thick_3d,
+    isotropic_fluorescent_thin_3d,
+    isotropic_thin_3d,
+    phase_thick_3d,
+)
+
+B = 4
+Z, Y, X = 5, 32, 32
+
+
+# --- isotropic_thin_3d ---
+
+
+class TestIsotropicThin3DBatched:
+    """Batched vs sequential tests for isotropic_thin_3d."""
+
+    @pytest.fixture
+    def optical_params(self):
+        return dict(
+            yx_pixel_size=6.5 / 40,
+            wavelength_illumination=0.532,
+            index_of_refraction_media=1.33,
+            numerical_aperture_illumination=0.9,
+            numerical_aperture_detection=1.2,
+        )
+
+    @pytest.fixture
+    def z_position_list(self):
+        return (-torch.arange(Z) + Z // 2).float() * 0.1
+
+    def test_batched_reconstruct_matches_sequential(self, optical_params, z_position_list):
+        torch.manual_seed(42)
+        bzyx = torch.randn(B, Z, Y, X)
+
+        # Sequential
+        seq_abs, seq_phase = [], []
+        for b in range(B):
+            a, p = isotropic_thin_3d.reconstruct(bzyx[b], z_position_list=z_position_list, **optical_params)
+            seq_abs.append(a)
+            seq_phase.append(p)
+        seq_abs = torch.stack(seq_abs)
+        seq_phase = torch.stack(seq_phase)
+
+        # Batched
+        bat_abs, bat_phase = isotropic_thin_3d.reconstruct(bzyx, z_position_list=z_position_list, **optical_params)
+
+        torch.testing.assert_close(bat_abs, seq_abs, atol=0, rtol=0)
+        torch.testing.assert_close(bat_phase, seq_phase, atol=0, rtol=0)
+
+    def test_b1_matches_unbatched(self, optical_params, z_position_list):
+        torch.manual_seed(42)
+        zyx = torch.randn(Z, Y, X)
+
+        ref_abs, ref_phase = isotropic_thin_3d.reconstruct(zyx, z_position_list=z_position_list, **optical_params)
+
+        bat_abs, bat_phase = isotropic_thin_3d.reconstruct(
+            zyx.unsqueeze(0), z_position_list=z_position_list, **optical_params
+        )
+
+        torch.testing.assert_close(bat_abs.squeeze(0), ref_abs, atol=0, rtol=0)
+        torch.testing.assert_close(bat_phase.squeeze(0), ref_phase, atol=0, rtol=0)
+
+    def test_batched_tf_with_per_tile_tilt(self, optical_params, z_position_list):
+        torch.manual_seed(42)
+        bzyx = torch.randn(B, Z, Y, X)
+
+        zeniths = torch.tensor([0.0, 0.05, 0.1, 0.15])
+        azimuths = torch.tensor([0.0, 0.5, 1.0, 1.5])
+
+        # Sequential with per-tile tilt
+        seq_abs, seq_phase = [], []
+        for b in range(B):
+            a, p = isotropic_thin_3d.reconstruct(
+                bzyx[b],
+                z_position_list=z_position_list,
+                tilt_angle_zenith=zeniths[b].item(),
+                tilt_angle_azimuth=azimuths[b].item(),
+                **optical_params,
+            )
+            seq_abs.append(a)
+            seq_phase.append(p)
+        seq_abs = torch.stack(seq_abs)
+        seq_phase = torch.stack(seq_phase)
+
+        # Batched TF with per-tile tilt
+        abs_tf, phase_tf = isotropic_thin_3d.calculate_transfer_function(
+            (Y, X),
+            z_position_list=z_position_list,
+            tilt_angle_zenith=zeniths,
+            tilt_angle_azimuth=azimuths,
+            **optical_params,
+        )
+        singular_system = isotropic_thin_3d.calculate_singular_system(abs_tf, phase_tf)
+        bat_abs, bat_phase = isotropic_thin_3d.apply_inverse_transfer_function(bzyx, singular_system)
+
+        # Batched SVD produces slightly different singular values at
+        # near-zero frequencies due to floating-point ordering, which gets
+        # amplified by the Tikhonov regularized inverse. Match to within
+        # 1% of the signal standard deviation.
+        for b in range(B):
+            atol_p = 0.01 * seq_phase[b].std()
+            torch.testing.assert_close(bat_phase[b], seq_phase[b], rtol=1e-3, atol=atol_p)
+            atol_a = 0.01 * seq_abs[b].std()
+            torch.testing.assert_close(bat_abs[b], seq_abs[b], rtol=1e-3, atol=atol_a)
+
+
+# --- isotropic_fluorescent_thin_3d ---
+
+
+class TestFluorescentThin3DBatched:
+    """Batched vs sequential tests for isotropic_fluorescent_thin_3d."""
+
+    @pytest.fixture
+    def optical_params(self):
+        return dict(
+            yx_pixel_size=6.5 / 40,
+            wavelength_emission=0.532,
+            index_of_refraction_media=1.33,
+            numerical_aperture_detection=1.2,
+        )
+
+    @pytest.fixture
+    def z_position_list(self):
+        return (-torch.arange(Z) + Z // 2).float() * 0.1
+
+    def test_batched_reconstruct_matches_sequential(self, optical_params, z_position_list):
+        torch.manual_seed(42)
+        bzyx = torch.randn(B, Z, Y, X).abs() + 1  # fluorescence is positive
+
+        seq = []
+        for b in range(B):
+            r = isotropic_fluorescent_thin_3d.reconstruct(bzyx[b], z_position_list=z_position_list, **optical_params)
+            seq.append(r)
+        seq = torch.stack(seq)
+
+        bat = isotropic_fluorescent_thin_3d.reconstruct(bzyx, z_position_list=z_position_list, **optical_params)
+
+        torch.testing.assert_close(bat, seq, atol=0, rtol=0)
+
+    def test_b1_matches_unbatched(self, optical_params, z_position_list):
+        torch.manual_seed(42)
+        zyx = torch.randn(Z, Y, X).abs() + 1
+
+        ref = isotropic_fluorescent_thin_3d.reconstruct(zyx, z_position_list=z_position_list, **optical_params)
+        bat = isotropic_fluorescent_thin_3d.reconstruct(
+            zyx.unsqueeze(0), z_position_list=z_position_list, **optical_params
+        )
+
+        torch.testing.assert_close(bat.squeeze(0), ref, atol=0, rtol=0)
+
+
+# --- phase_thick_3d ---
+
+
+class TestPhaseThick3DBatched:
+    """Batched vs sequential tests for phase_thick_3d."""
+
+    @pytest.fixture
+    def optical_params(self):
+        return dict(
+            yx_pixel_size=6.5 / 40,
+            z_pixel_size=0.25,
+            wavelength_illumination=0.532,
+            z_padding=5,
+            index_of_refraction_media=1.33,
+            numerical_aperture_illumination=0.9,
+            numerical_aperture_detection=1.2,
+        )
+
+    def test_batched_reconstruct_matches_sequential(self, optical_params):
+        torch.manual_seed(42)
+        bzyx = torch.randn(B, Z, Y, X)
+
+        seq = []
+        for b in range(B):
+            r = phase_thick_3d.reconstruct(bzyx[b], **optical_params)
+            seq.append(r)
+        seq = torch.stack(seq)
+
+        bat = phase_thick_3d.reconstruct(bzyx, **optical_params)
+
+        torch.testing.assert_close(bat, seq, atol=1e-6, rtol=1e-6)
+
+    def test_b1_matches_unbatched(self, optical_params):
+        torch.manual_seed(42)
+        zyx = torch.randn(Z, Y, X)
+
+        ref = phase_thick_3d.reconstruct(zyx, **optical_params)
+        bat = phase_thick_3d.reconstruct(zyx.unsqueeze(0), **optical_params)
+
+        torch.testing.assert_close(bat.squeeze(0), ref, atol=1e-6, rtol=1e-6)
+
+
+# --- isotropic_fluorescent_thick_3d ---
+
+
+class TestFluorescentThick3DBatched:
+    """Batched vs sequential tests for isotropic_fluorescent_thick_3d."""
+
+    @pytest.fixture
+    def optical_params(self):
+        return dict(
+            yx_pixel_size=6.5 / 40,
+            z_pixel_size=0.25,
+            wavelength_emission=0.532,
+            z_padding=5,
+            index_of_refraction_media=1.33,
+            numerical_aperture_detection=1.2,
+        )
+
+    def test_batched_reconstruct_matches_sequential(self, optical_params):
+        torch.manual_seed(42)
+        bzyx = torch.randn(B, Z, Y, X).abs() + 1
+
+        seq = []
+        for b in range(B):
+            r = isotropic_fluorescent_thick_3d.reconstruct(bzyx[b], **optical_params)
+            seq.append(r)
+        seq = torch.stack(seq)
+
+        bat = isotropic_fluorescent_thick_3d.reconstruct(bzyx, **optical_params)
+
+        torch.testing.assert_close(bat, seq, atol=1e-6, rtol=1e-6)
+
+    def test_b1_matches_unbatched(self, optical_params):
+        torch.manual_seed(42)
+        zyx = torch.randn(Z, Y, X).abs() + 1
+
+        ref = isotropic_fluorescent_thick_3d.reconstruct(zyx, **optical_params)
+        bat = isotropic_fluorescent_thick_3d.reconstruct(zyx.unsqueeze(0), **optical_params)
+
+        torch.testing.assert_close(bat.squeeze(0), ref, atol=1e-6, rtol=1e-6)