feat(phase_thick_3d_tilt): Re-factor of ops_process PR#99 into OPS agnostic waveorder optimizations.

[mark-a-potts]

Headline

This is a refactor of the optimizations from https://github.com/royerlab/ops_process/pull/99. Optimizations have been pushed down to the waveorder level in order to make them more generically applicable to other projects. The idea is to offer this as a starting point for Shalin's 3-month "Scaling WaveOrder 3.0 to Biohub data" project (proposal dated 2026-05-29). That project's 20× target needs a bit more work, but this achieves roughly 4-16× per-GPU and 6.6-6.8× wallclock on ops0154 — at PROD cell-count parity across all wells (track ±0.18%, pheno ±0.05%) once the OPS adapter applies the meniscus override (PR3).

WaveOrder-on-OPS tilt-recon is 4-16× more GPU-efficient per cell than the vanilla per-position SLURM-array path (16× on pheno from skip_opt + warmstart map, 4× on track since ~40% of positions run the full 25-iter PROD recipe via the meniscus override to hit cell-count parity). Pipeline wallclock speedups are larger because we additionally amortize SLURM-array dispatch via INPROC and run more GPUs in parallel.

Performance — vs vanilla per-position SLURM-array on ops0154

Test data: 7,035 pheno positions × 1 timepoint × 3 wells; 296 track positions × 2 timepoints × 2 wells.

The Vanilla GPU-h numbers come from a direct per-position SLURM-array benchmark (100 pheno tasks, 50 track tasks, 1 GPU per task on H100/H200/Blackwell). Per-position median wall × position count gives the theoretical compute floor of the vanilla path.

Workload	Vanilla per-pos wall	Vanilla GPU-h	Ours GPU-h	Per-GPU speedup	Our wallclock	Pipeline wallclock
Pheno (7,035 pos × 1 T)	74 s	~145	9	~16×	45 min (12 GPU)	6.8× vs PROD 5 h 5 min
Track (296 pos × 2 T, meniscus@0.70)	79 s	~6.4	~1.6	~4×	~50 min (2 GPU)	6.6× vs PROD 5 h 30 min

Per-GPU speedup is the algorithmic compute-per-cell improvement, independent of parallelism. Pheno gains more (16×) because warmstart_map + skip_opt skips NAdam entirely on most positions; track gains less (4×) because the meniscus override at r/r_max > 0.70 puts ~40% of positions on the full 25-iter PROD recipe to hit cell-count parity.

Pipeline wallclock includes additional benefits of INPROC dispatch (eliminates SLURM-array per-task startup tax — PROD's 5h30m track wall was dominated by scheduling overhead, not compute) and running multiple GPUs in parallel.

Multi-timepoint amortization. The numbers above are for ops0154, where pheno has 1 timepoint per position and track has 2. The T-cache (memoize optimized tilt params per position — first call runs full optim, subsequent T's at the same position skip the 2D Adam phase) means per-position cost drops as more timepoints are added. Track's 4× at N=2 already includes ~2× from T-cache; many-T workloads (e.g. a LiveScreen-style experiment) will see substantially larger savings.

Cell-count parity vs PROD

Pheno: PASS, all 3 wells within ±0.05% of PROD on ops0154 cell-count verify (-0.018% / -0.026% / -0.022%). Acceptance ±0.5%.

Track: PASS, all 4 well-timepoints within ±0.18% of PROD with the OPS-side meniscus override at OPS_TILT_MENISCUS_R_THRESHOLD=0.70:

Well/T	PROD	ours	% diff
A/1 T0	510,607	511,431	+0.16%
A/1 T1	530,736	531,690	+0.18%
A/2 T0	236,989	236,972	-0.01%
A/2 T1	541,318	541,445	+0.02%

The track cell-count residual at edge positions was traced to the speed recipe (zen_blend + frozen angles + radius-scaled 3-25 iters + T-cache) diverging from PROD's full 10-25-iter unfrozen NAdam. The OPS adapter PR (royerlab/ops_process#104) detects meniscus positions (r/r_max > T) and routes them to the PROD recipe (full 25-iter NAdam, no zen_blend, no T-cache, no warmstart skip). Threshold sweep:

Threshold	Meniscus pos/well	max |Δ vs PROD|	±0.5% accept
0.85	11 (7%)	2.78%	FAIL
0.75	39 (26%)	0.66%	FAIL
0.70	59 (40%)	0.18%	PASS
0.60	87 (59%)	0.17%	PASS
0.50	111 (75%)	0.12%	PASS

0.70 is the highest threshold that passes ±0.5% — the optimum (lower thresholds add compute without improving accuracy).

A five-way env-flag bisect on meniscus positions (toggling WAVEORDER_FAST_2D_TIKHONOV, GPU-shared-optics, batched-vs-per-subtile NAdam — and including a "vanilla waveorder 3.0.2 + ops_process main" reference run) showed all variants within ±0.02% of each other in meniscus-band cell counts. So none of this PR's three numerical optimizations (closed-form Tikhonov, GPU-resident optics, batched NAdam) is the cause of the residual — the residual was purely the recipe choice at edge positions, which the OPS-side override fixes.

Tilt-recon quality examples — best / median / worst per branch

A/1 well, T=0. Pearson(ours vs PROD) computed per FOV on the 2D phase output. Each panel: PROD (top) / ours (middle) / abs-diff (bottom). Images live in the OPS adapter PR's pr_artifacts/.

Track (148 positions, threshold=0.70, Pearson: min 0.941, median 0.972, max 1.000):

Pheno (2,345 positions in A/1, Pearson: min 0.937, median 0.979, max 0.994):

The "worst" panels show disagreement concentrated at high-contrast cell-edge pixels — biological structure is preserved across all panels.

What's in this PR

Six pieces, transfer cleanly to the upstream project:

1. waveorder.models.phase_thick_3d_tilt.optimize_subtile_tilt_params — batched per-subtile NAdam over N subtiles × 3 params (z_offset, zenith, azimuth) each. Existing optimize_reconstruction is structurally single-parameter-set; the batched-subtile shape is genuinely new API.

2. waveorder.models.phase_thick_3d_tilt.radial_blend_zenith_init — radial blend of per-subtile zen init from a baseline formula. Used for low-NA track tilt to prevent zen=0 init from drifting to the wrong focal plane on well-edge positions.

3. GPU-resident shared optics — _compute_shared_optics(device=...) builds propagation kernel + Green's function on-device instead of CPU torch.exp + .to(device). Verified 1.79× on ops0154 pheno full-scale (89 min → 49.7 min on Blackwell 8-GPU). GPU util 13% → 78%.

4. Closed-form 2×2 Tikhonov inverse in isotropic_thin_3d.reconstruct — replaces SVD path with a closed-form Hermitian inverse. Mathematically equivalent, ~18× faster per call. Bit-identical Pearson on ops0154.

5. CachedTiltOptics — angle-fixed optics computed once per well and reused when freeze_axes=("zenith","azimuth"). ~2× speedup on track-tilt where angles are frozen post-calibration.

6. LM-damped diagonal Newton added as method='newton' option to optimize_reconstruction.

Tests

tests/models/test_phase_thick_3d_tilt.py — 8 tests covering shape invariants, synthetic recovery, frozen-axis behavior, warmstart-skip roundtrip, and the radial-blend zero-r edge case. All pass on CPU and CUDA.

Anti-patterns deliberately avoided (per OPS memory)

• Convergence stays on absolute-0.5 threshold; no rtol added.
• No best-params restoration on plateau.
• No CUDA streams within a worker (PyTorch allocator serializes).
• No geometric z correction in defaults.

Open caveats

• Some pieces in the sibling ops_process PR (per-well SLURM fan-out, universal warmstart map, meniscus override OPS_TILT_MENISCUS_R_THRESHOLD) bake in OPS-specific assumptions and are likely candidates for rework when the WaveOrder 3.0 engineer designs the generalizable layer. Those are intentionally kept out of this PR.

Draft

This is a draft while the OPS-side ops_process / ops_monorepo PRs that pin to this branch land. Numbers above are reproducible from mark/ops0154_cell_count_verify (this-stack run at threshold 0.85), mark/ops0154_bisect_men_0p{50,60,70,75} (threshold sweep), mark/ops0154_bisect_vanilla (PROD-stack reference), mark/ops0154_vanilla_pheno_benchmark / mark/ops0154_vanilla_track_benchmark (vanilla extrapolations) on the Biohub HPC.

🤖 Generated with Claude Code