diff --git a/.gitignore b/.gitignore index cc908aeb..5ef7a182 100644 --- a/.gitignore +++ b/.gitignore @@ -151,6 +151,8 @@ waveorder/_version.py *.npz *.tif[f] *.pdf +# the screening tutorial ships its vector figures and whitepaper PDF +!docs/guide/labelfree-screening/**/*.pdf # example data /examples/data_temp/* diff --git a/docs/guide.rst b/docs/guide.rst index 2954c6cc..ddf596d2 100644 --- a/docs/guide.rst +++ b/docs/guide.rst @@ -10,6 +10,7 @@ User Guides guide/QLIPP guide/data-schema + guide/labelfree-screening/labelfree-screening guide/buyers-guide guide/microscope-installation-guide guide/software-installation-guide diff --git a/docs/guide/labelfree-screening/README.md b/docs/guide/labelfree-screening/README.md new file mode 100644 index 00000000..ba5792ef --- /dev/null +++ b/docs/guide/labelfree-screening/README.md @@ -0,0 +1,27 @@ +--- +orphan: true +--- + +# Label-free screening tutorial: source files + +This folder is self-contained: the tutorial page, the 2-page whitepaper PDF +source, the figures, the configs the tutorial teaches, and the scripts that +regenerate everything. + +| Path | Contents | +| --- | --- | +| `labelfree-screening.md` | The tutorial page (rendered on the docs site). | +| `whitepaper.tex` | 2-page print layout of the same content (boxed failure modes, Biohub footer). | +| `make_pdf.sh` | Builds `whitepaper.pdf` with pdflatex (base texlive packages only). | +| `figures/` | Tutorial figures: vector PDF (whitepaper) and PNG (docs site). | +| `figures/make_figures.py` | Regenerates the figures. Requires Biohub HPC access to the OPS atlas data and reconstruction outputs; see the module docstring for the analysis directory. | +| `configs/phase.yml` | Baseline 2D phase reconstruction config. | +| `configs/phase_autofocus.yml` | Baseline plus `z_focus_offset: auto`. | +| `configs/phase_autofocus_tilt.yml` | Autofocus plus warm-started illumination tilt fitting. | +| `brand/` | CZ Biohub logo for the whitepaper footer. | + +To update a figure: edit `figures/make_figures.py`, run it on the Biohub HPC +(it writes into the analysis directory), and copy the outputs here. All +reconstructions behind the figures are plain `waveorder reconstruct` calls +with the configs in `configs/` on positions `A/1/029029` (well center) and +`A/1/002029` (north edge) of the OPS atlas store. diff --git a/docs/guide/labelfree-screening/brand/biohub-logo-black.png b/docs/guide/labelfree-screening/brand/biohub-logo-black.png new file mode 100644 index 00000000..84232814 Binary files /dev/null and b/docs/guide/labelfree-screening/brand/biohub-logo-black.png differ diff --git a/docs/guide/labelfree-screening/configs/phase.yml b/docs/guide/labelfree-screening/configs/phase.yml new file mode 100644 index 00000000..99199e0a --- /dev/null +++ b/docs/guide/labelfree-screening/configs/phase.yml @@ -0,0 +1,13 @@ +input_channel_names: [BF] +reconstruction_dimension: 2 +phase: + transfer_function: + yx_pixel_size: 0.325 + z_pixel_size: 2.0 + wavelength_illumination: 0.450 + index_of_refraction_media: 1.0 + numerical_aperture_detection: 0.55 + numerical_aperture_illumination: 0.4 + invert_phase_contrast: false + apply_inverse: + regularization_strength: 0.001 diff --git a/docs/guide/labelfree-screening/configs/phase_autofocus.yml b/docs/guide/labelfree-screening/configs/phase_autofocus.yml new file mode 100644 index 00000000..553ad88b --- /dev/null +++ b/docs/guide/labelfree-screening/configs/phase_autofocus.yml @@ -0,0 +1,14 @@ +input_channel_names: [BF] +reconstruction_dimension: 2 +phase: + transfer_function: + yx_pixel_size: 0.325 + z_pixel_size: 2.0 + wavelength_illumination: 0.450 + index_of_refraction_media: 1.0 + numerical_aperture_detection: 0.55 + numerical_aperture_illumination: 0.4 + invert_phase_contrast: false + z_focus_offset: auto + apply_inverse: + regularization_strength: 0.001 diff --git a/docs/guide/labelfree-screening/configs/phase_autofocus_tilt.yml b/docs/guide/labelfree-screening/configs/phase_autofocus_tilt.yml new file mode 100644 index 00000000..db8480ee --- /dev/null +++ b/docs/guide/labelfree-screening/configs/phase_autofocus_tilt.yml @@ -0,0 +1,19 @@ +input_channel_names: [BF] +reconstruction_dimension: 2 +phase: + transfer_function: + yx_pixel_size: 0.325 + z_pixel_size: 2.0 + wavelength_illumination: 0.450 + index_of_refraction_media: 1.0 + numerical_aperture_detection: 0.55 + numerical_aperture_illumination: 0.4 + invert_phase_contrast: false + z_focus_offset: auto + tilt_angle_zenith: {init: 0.13, lr: 0.05} + tilt_angle_azimuth: {init: 0.0, lr: 0.05} + apply_inverse: + regularization_strength: 0.001 +optimization: + max_iterations: 50 + method: adam diff --git a/docs/guide/labelfree-screening/figures/fig_flip.pdf b/docs/guide/labelfree-screening/figures/fig_flip.pdf new file mode 100644 index 00000000..c38aec4c Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_flip.pdf differ diff --git a/docs/guide/labelfree-screening/figures/fig_flip.png b/docs/guide/labelfree-screening/figures/fig_flip.png new file mode 100644 index 00000000..fddfe7ad Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_flip.png differ diff --git a/docs/guide/labelfree-screening/figures/fig_focus.pdf b/docs/guide/labelfree-screening/figures/fig_focus.pdf new file mode 100644 index 00000000..e5868e35 Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_focus.pdf differ diff --git a/docs/guide/labelfree-screening/figures/fig_focus.png b/docs/guide/labelfree-screening/figures/fig_focus.png new file mode 100644 index 00000000..bbe556b0 Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_focus.png differ diff --git a/docs/guide/labelfree-screening/figures/fig_overview.pdf b/docs/guide/labelfree-screening/figures/fig_overview.pdf new file mode 100644 index 00000000..f7d0992b Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_overview.pdf differ diff --git a/docs/guide/labelfree-screening/figures/fig_overview.png b/docs/guide/labelfree-screening/figures/fig_overview.png new file mode 100644 index 00000000..23f3a54f Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_overview.png differ diff --git a/docs/guide/labelfree-screening/figures/fig_reg.pdf b/docs/guide/labelfree-screening/figures/fig_reg.pdf new file mode 100644 index 00000000..63074e25 Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_reg.pdf differ diff --git a/docs/guide/labelfree-screening/figures/fig_reg.png b/docs/guide/labelfree-screening/figures/fig_reg.png new file mode 100644 index 00000000..5846fc62 Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_reg.png differ diff --git a/docs/guide/labelfree-screening/figures/fig_tilt.pdf b/docs/guide/labelfree-screening/figures/fig_tilt.pdf new file mode 100644 index 00000000..76a7c56b Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_tilt.pdf differ diff --git a/docs/guide/labelfree-screening/figures/fig_tilt.png b/docs/guide/labelfree-screening/figures/fig_tilt.png new file mode 100644 index 00000000..a4b9670b Binary files /dev/null and b/docs/guide/labelfree-screening/figures/fig_tilt.png differ diff --git a/docs/guide/labelfree-screening/figures/make_figures.py b/docs/guide/labelfree-screening/figures/make_figures.py new file mode 100644 index 00000000..16b83716 --- /dev/null +++ b/docs/guide/labelfree-screening/figures/make_figures.py @@ -0,0 +1,276 @@ +"""Generate figures for the label-free screening tutorial. + +Reads the raw OPS position and the reconstructions in ./recons/, writes +vector PDF (for the whitepaper) and PNG (for the docs site) figure pairs +into ./figures/. Figures are drawn at final print width so fonts match +the 9 pt document when included at 100% scale. +""" + +from pathlib import Path + +import matplotlib +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from iohub import open_ome_zarr + +matplotlib.use("Agg") +plt.rcParams["font.family"] = "monospace" +plt.rcParams["font.size"] = 8 + +RAW_PATH = "/hpc/projects/icd.fast.ops/ops0105_20260106/0-convert/live_imaging/phenotyping_transform.zarr/A/1/029029" +# Biohub HPC analysis directory holding the raw crops and reconstructions. +WORK = Path("/hpc/mydata/talon.chandler/2026-11-04-wo-hcs-whitepaper") +RECONS = WORK / "recons" +FIGURES = WORK / "figures" + +FULL_W = 7.3 # print width (in) of a full-width figure at 1.3 cm margins +HALF_W = 3.52 # print width of a boxed half-width figure + +# Zoom regions chosen from full-FOV overviews: clean monolayer, no debris. +YS = slice(384, 896) +XS = slice(1152, 1664) +T_YS, T_XS = slice(256, 768), slice(1280, 1792) # edge-tile crop +PX_UM = 0.325 # yx pixel size in micrometers + + +def load_phase(name, position="029029"): + """Load the single 2D phase image from a reconstruction zarr. + + Parameters + ---------- + name : str + Basename of the zarr store in ./recons/ (without extension). + position : str, optional + Position name inside well A/1 of the store. + + Returns + ------- + numpy.ndarray + 2D phase image, float32. + """ + with open_ome_zarr(RECONS / f"{name}.zarr" / f"A/1/{position}") as ds: + return np.asarray(ds.data[0, 0, 0]).astype(np.float32) + + +def midband_power(img): + """Mean squared FFT magnitude in the 0.25-0.5 Nyquist annulus.""" + f = np.abs(np.fft.fftshift(np.fft.fft2(img))) + n = img.shape[0] + yy, xx = np.mgrid[-n // 2 : n // 2, -n // 2 : n // 2] + r = np.hypot(yy, xx) / (n / 2) + return float((f[(r > 0.25) & (r < 0.5)] ** 2).mean()) + + +def add_scalebar(ax, n_px, length_um=25): + """Draw a scale bar in the lower-right corner of an image axis. + + The bar and label carry a thin black outline so they stay legible + over bright cells, and the label right-aligns with the bar end so + a label wider than the bar overhangs inward, never off the panel. + """ + bar_px = length_um / PX_UM + x0, y0 = n_px - bar_px - 45, n_px - 55 + ax.plot( + [x0, x0 + bar_px], + [y0, y0], + color="white", + lw=1.5, + path_effects=[pe.withStroke(linewidth=3, foreground="black")], + ) + ax.text( + x0 + bar_px, + y0 - 10, + f"{length_um} µm", + color="white", + ha="right", + va="bottom", + fontsize=7, + path_effects=[pe.withStroke(linewidth=1.5, foreground="black")], + ) + + +def panel(ax, img, title, vmin, vmax, scalebar=False): + """Render one grayscale image panel without ticks.""" + ax.imshow(img, cmap="gray", vmin=vmin, vmax=vmax, interpolation="nearest") + ax.set_title(title, fontsize=8) + ax.set_xticks([]) + ax.set_yticks([]) + if scalebar: + add_scalebar(ax, img.shape[-1]) + + +def save(fig, stem): + """Save a figure as vector PDF (whitepaper) and PNG (docs site).""" + for ext, dpi in (("pdf", 300), ("png", 220)): + fig.savefig(FIGURES / f"{stem}.{ext}", dpi=dpi, bbox_inches="tight", transparent=True) + plt.close(fig) + print(f"wrote {stem}.pdf/.png") + + +with open_ome_zarr(RAW_PATH) as raw_ds: + bf_idx = raw_ds.channel_names.index("BF") + bf = np.asarray(raw_ds.data[0, bf_idx]).astype(np.float32) # (7, 2048, 2048) + +nominal = load_phase("nominal") +lo, hi = np.percentile(nominal[YS, XS], [1, 99.5]) + + +# --- Figure 1: overview flowchart -------------------------------------------- +fig = plt.figure(figsize=(FULL_W, 2.1)) + + +def icon_axes(rect, aspect_equal=False): + ax = fig.add_axes(rect) + ax.set_xlim(0, 1) + ax.set_ylim(0, 1) + if aspect_equal: + ax.set_aspect("equal", anchor="W") + ax.axis("off") + return ax + + +ICON_TEXT_X = 0.085 # left-aligned icon labels, all in one column + +# Koehler alignment on top: centered concentric circles, text to the right. +ax = icon_axes([0.015, 0.68, 0.055, 0.27], aspect_equal=True) +for r in (0.15, 0.30, 0.46): + ax.add_patch(plt.Circle((0.5, 0.5), r, fill=False, color="0.25", lw=1.0)) +ax.plot(0.5, 0.5, marker="+", color="0.25", markersize=5) +fig.text(ICON_TEXT_X, 0.815, "Koehler\nalignment", fontsize=7, va="center", ha="left") + +# Transmission optics: cone angles drawn to scale for each NA. +# half-angle = asin(NA): 23.6 deg for 0.4 (illumination, top), +# 33.4 deg for 0.55 (detection, bottom). +ax = icon_axes([0.015, 0.36, 0.055, 0.27], aspect_equal=True) +tan_ill = np.tan(np.arcsin(0.4)) +tan_det = np.tan(np.arcsin(0.55)) +# trapezoids: cone sides truncated near the sample plane (y = 0.5) +ill_near, ill_far = 0.70, 0.95 # condenser above the sample +det_near, det_far = 0.30, 0.05 # objective below the sample +ax.add_patch( + plt.Polygon( + [ + [0.5 - (ill_near - 0.5) * tan_ill, ill_near], + [0.5 + (ill_near - 0.5) * tan_ill, ill_near], + [0.5 + (ill_far - 0.5) * tan_ill, ill_far], + [0.5 - (ill_far - 0.5) * tan_ill, ill_far], + ], + color="0.45", + ) +) +ax.add_patch( + plt.Polygon( + [ + [0.5 - (0.5 - det_near) * tan_det, det_near], + [0.5 + (0.5 - det_near) * tan_det, det_near], + [0.5 + (0.5 - det_far) * tan_det, det_far], + [0.5 - (0.5 - det_far) * tan_det, det_far], + ], + color="0.25", + ) +) +ax.plot([0.08, 0.92], [0.5, 0.5], color="black", lw=1.0) +fig.text(ICON_TEXT_X, 0.555, "0.4 NA condenser", fontsize=7, va="center", ha="left") +fig.text(ICON_TEXT_X, 0.425, "0.55 NA, 20$\\times$\ndetection", fontsize=7, va="center", ha="left") + +# z stack: 7 slices. +ax = icon_axes([0.015, 0.05, 0.055, 0.26], aspect_equal=True) +for i in range(7): + ax.plot([0.15, 0.85], [0.08 + i * 0.14] * 2, color="0.35", lw=1.0) +fig.text(ICON_TEXT_X, 0.18, "7-slice z stack\n2 µm spacing", fontsize=7, va="center", ha="left") + +# Defocus stack: all 7 planes fanned bottom-to-top in acquisition order, +# shared grayscale so the background stays consistent across slices. +stack_lo, stack_hi = np.percentile(bf[:, 640:1152, 1280:1792], [0.5, 99.5]) +for k, z in enumerate(range(7)): + ax = fig.add_axes([0.268 + 0.006 * k, 0.085 + 0.053 * k, 0.135, 0.47]) + img = bf[z, 640:1152, 1280:1792] + ax.imshow(img, cmap="gray", vmin=stack_lo, vmax=stack_hi) + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_color("0.3") +fig.text(0.345, 0.025, "7-slice defocus stack", fontsize=8, ha="center") + +# WaveOrder box: config + CLI. +ax = icon_axes([0.525, 0.05, 0.245, 0.9]) +ax.add_patch(plt.Rectangle((0.02, 0.02), 0.96, 0.96, fill=False, color="0.3", lw=1.0)) +ax.text(0.5, 0.885, "WaveOrder", fontsize=9, ha="center", weight="bold") +ax.text( + 0.08, + 0.56, + "$ cat config.yml\n...\nyx_pixel_size: 0.325\nz_pixel_size: 2.0\nNA_detection: 0.55\nNA_illumination: 0.4\nz_focus_offset: auto\n...", + fontsize=6.5, + va="center", +) +ax.text( + 0.08, + 0.19, + "$ waveorder reconstruct \\\n -i data.zarr \\\n -c config.yml \\\n -o phase.zarr", + fontsize=6.5, + va="center", +) + +# Reconstruction. +ax = fig.add_axes([0.825, 0.13, 0.155, 0.72]) +ax.imshow(nominal[YS, XS], cmap="gray", vmin=lo, vmax=hi) +ax.set_xticks([]) +ax.set_yticks([]) +fig.text(0.9025, 0.025, "quantitative phase", fontsize=8, ha="center") + +# Arrows in the clear gaps between columns. +for x0, x1 in ((0.215, 0.255), (0.475, 0.515), (0.785, 0.818)): + fig.add_artist( + matplotlib.patches.FancyArrowPatch( + (x0, 0.5), + (x1, 0.5), + transform=fig.transFigure, + arrowstyle="-|>", + mutation_scale=10, + color="0.15", + lw=1.2, + ) + ) + +save(fig, "fig_overview") + + +# --- Figure: contrast flip (half-width box) ----------------------------------- +flipped = load_phase("flipped") +fig, axs = plt.subplots(1, 2, figsize=(HALF_W, 2.0)) +panel(axs[0], nominal[YS, XS], "correct\n(nucleoli bright)", lo, hi, scalebar=True) +panel(axs[1], flipped[YS, XS], "flipped z axis\n(nucleoli dark)", lo, hi) +save(fig, "fig_flip") + +# --- Figure: regularization sweep (half-width box) ---------------------------- +fig, axs = plt.subplots(1, 3, figsize=(HALF_W, 1.45)) +for ax, reg, note in zip(axs, ["1e-5", "1e-3", "1e-1"], ["too noisy", "chosen", "too smooth"]): + img = nominal if reg == "1e-3" else load_phase(f"reg_{reg}") + plo, phi = np.percentile(img[YS, XS], [1, 99.5]) + panel(ax, img[YS, XS], f"reg {reg}\n({note})", plo, phi, scalebar=(reg == "1e-5")) +save(fig, "fig_reg") + +# --- Figure: focus offset (half-width box) ------------------------------------- +F_YS, F_XS = slice(384, 704), slice(1344, 1664) +fig, axs = plt.subplots(1, 3, figsize=(HALF_W, 1.45)) +for ax, off, name in zip(axs, ["-2", "0", "+2"], ["offset_-2.0", None, "defocused"]): + img = nominal if name is None else load_phase(name) + sharp = midband_power(img[YS, XS]) + note = " (chosen)" if off == "0" else "" + panel(ax, img[F_YS, F_XS], f"{off} slices{note}", lo, hi, scalebar=(off == "-2")) +save(fig, "fig_focus") + +# --- Figure: illumination tilt on an edge FOV (half-width box) ----------------- +edge_base = load_phase("edge_baseline", position="002029") +edge_tilt = load_phase("edge_apply", position="002029") +fig, axs = plt.subplots(1, 2, figsize=(HALF_W, 2.0)) +for ax, img, title in zip( + axs, + [edge_base, edge_tilt], + ["tilt not modeled", "tilt fitted\n(0.49 rad zenith)"], +): + crop = img[T_YS, T_XS] + plo, phi = np.percentile(crop, [1, 99.5]) + panel(ax, crop, title, plo, phi, scalebar=(img is edge_base)) +save(fig, "fig_tilt") diff --git a/docs/guide/labelfree-screening/labelfree-screening.md b/docs/guide/labelfree-screening/labelfree-screening.md new file mode 100644 index 00000000..b71b8ea2 --- /dev/null +++ b/docs/guide/labelfree-screening/labelfree-screening.md @@ -0,0 +1,168 @@ +# WaveOrder for label-free high-throughput screening + +![From acquisition settings to quantitative phase: a brightfield defocus stack plus one WaveOrder config yields label-free contrast at screening scale](./figures/fig_overview.png) + +This guide demonstrates WaveOrder for label-free high-throughput screening +using optical pooled screening (OPS) data from the +[Biohub's multimodal perturbation atlas](https://www.biorxiv.org/content/10.64898/2026.06.01.728087v1) +[1], where phase from defocus enabled label-free tracking and phenotyping of +CRISPR-perturbed A549 cells. Brightfield defocus stacks from a standard +motorized widefield microscope, together with a +[WaveOrder](https://github.com/biohub/waveorder) config [2], yield +quantitative phase at screening scale (see the overview figure). + +After reading this guide and acquiring data, use the +[first-reconstruction guide](../first-reconstruction.md) and connect your +agent to the [Biohub MCP](https://biohub.ai/mcp) to run the +`waveorder-tutorial` skill. Methods are described in the +[WaveOrder preprint](https://arxiv.org/abs/2412.09775) [3]. + +## Acquisition + +The three settings that matter most are: + +- **Koehler illumination.** Align the condenser before every session. + Misalignment is the most common source of artifacts. Ron Vale's + five-minute + [iBiology tutorial](https://www.ibiology.org/talks/set-up-koehler-illumination/) + [4] is a good refresher. +- **Condenser NA of roughly 2/3 the detection NA.** We used a 0.4 NA + condenser and a 0.55 NA objective. A larger condenser aperture reduces + contrast, while a smaller condenser aperture reduces contrast and + resolution. +- **Z spacing and range.** At 20x, start at 2 µm spacing and cover the full + cell height with visibly blurred features on either end of the stack + (here, 7 slices = 12 µm). For different objectives, scale inversely with + magnification: e.g. 4 µm at 10x, 1 µm at 40x. + +| Parameter | Perturbation atlas value | Notes | +| --- | --- | --- | +| Objective | 20x, 0.55 NA, air | any objective works; match the config | +| Condenser | 0.4 NA | ~2/3 of detection NA, Koehler aligned | +| Z stack | 7 slices, 2 µm | confirm defocused features on top and bottom slices; use more slices for thicker samples | +| Camera | 6.5 µm pixels | 6.5 µm / 20x = 0.325 µm at the sample | +| Wavelength | 450 nm | any wavelength is fine; ideally use a 10 nm notch filter | +| Exposure | a few ms | fill 1/2-2/3 of the dynamic range | + +## Reconstruction + +Convert to OME-Zarr with [iohub](https://github.com/czbiohub-sf/iohub) +(`iohub convert -i -o data.zarr`), then record the acquisition +parameters in a file named `config.yml`: + +```yaml +input_channel_names: [BF] +reconstruction_dimension: 2 +phase: + transfer_function: + yx_pixel_size: 0.325 # 6.5 um / 20x + z_pixel_size: 2.0 + wavelength_illumination: 0.450 + index_of_refraction_media: 1.0 + numerical_aperture_detection: 0.55 + numerical_aperture_illumination: 0.4 + invert_phase_contrast: false + z_focus_offset: auto + apply_inverse: + regularization_strength: 0.001 +``` + +Reconstruct one position and compare it to the raw data in napari: + +```bash +waveorder reconstruct -i data.zarr/A/1/029029 -c config.yml -o phase.zarr +waveorder view data.zarr/A/1/029029 phase.zarr/A/1/029029 +``` + +Dense structures (nucleoli, nuclear envelopes, cell margins) should appear +brighter than background, as in the overview figure. +`reconstruction_dimension: 2` returns one phase image per defocus stack, the +best choice for samples that are thin compared to the depth of field of the +microscope. If the sample appears to blur all at once as you defocus, use 2. +If different regions of the sample blur at different defocus positions, +consider trialling 3. + +## Failure modes and their fixes + +**Failure mode 1 | Flipped z axis.** Stages disagree on the z sign +convention (is the first slice above or below focus?). If the convention is +flipped, contrast will be inverted and downstream comparison will be +challenging. Dense organelles should be bright. Fix: +`invert_phase_contrast: true`. Check once per microscope by reconstructing +one stack both ways. + +![Correct versus inverted phase contrast](./figures/fig_flip.png) + +**Failure mode 2 | Incorrect regularization.** `regularization_strength` +balances noise and resolution and varies by orders of magnitude across +samples and exposures. Sweep decades on one position and choose manually. +Too small amplifies blotchy low-frequency noise; too large leaves a +washed-out image. The optimum is usually within a decade of `1e-3`. + +```bash +for reg in 1e-5 1e-4 1e-3 1e-2 1e-1; do + sed "s/regularization_strength: .*/regularization_strength: $reg/" config.yml > config_$reg.yml + waveorder reconstruct -i data.zarr/A/1/029029 -c config_$reg.yml -o phase_$reg.zarr +done +waveorder view phase_*.zarr/A/1/029029 +``` + +![Regularization sweep](./figures/fig_reg.png) + +**Failure mode 3 | Autofocus drift.** When the stack center (e.g. the 4th +slice in a 7 slice stack) is not in focus, the reconstruction will be +blurry. Setting `z_focus_offset: auto` runs the search that can recover the +defocus offset: + +``` +$ waveorder reconstruct -i data.zarr/A/1/029029 -c config.yml -o phase.zarr +Resolving z_focus_offset with a per-slice focus search... +Resolved z_focus_offset = 0 (slice units) from the data; per-crop focus slices [3, 3, 3, 3, 3] +``` + +![z_focus_offset sweep](./figures/fig_focus.png) + +**Failure mode 4 | Tilted illumination at the well edge.** Misaligned +Koehler and imaging near a well-edge meniscus tilts the illumination, making +cells look shadow-cast with reduced contrast. Fit the tilt by gradient +descent in the same config: + +```yaml +# inside phase.transfer_function: + z_focus_offset: auto + tilt_angle_zenith: {init: 0.13, lr: 0.05} + tilt_angle_azimuth: {init: 0.0, lr: 0.05} +# and at the top level: +optimization: + max_iterations: 50 + method: adam +``` + +![Edge-of-well field with and without the fitted illumination tilt](./figures/fig_tilt.png) + +## Scaling up + +A single config file should generate good reconstructions across a screen. +We suggest parallelizing `waveorder reconstruct` over positions (one process +or cluster job per position) with `device: cuda:0` to reconstruct on a GPU. +We recommend `z_focus_offset: auto` to find the focus of each position. + +We recommend manual supervision to set up the optimized +`tilt_angle_{zenith,azimuth}` parameters across entire wells. We find that +initializing these tilt parameters at good initial tilts generates the +stablest reconstructions, so we recommend inspecting reconstructions at the +well edges, then refining the initial values until you reach stable +reconstructions that pass QC. + +## References + +1. C. Liu, A. V. Hillsley, M. Sekhar, et al., "A Multimodal Perturbation + Atlas Defines the Phenotypic Resolution of Cellular Morphology," + [bioRxiv 2026.06.01.728087](https://www.biorxiv.org/content/10.64898/2026.06.01.728087v1). +2. WaveOrder repository: + [github.com/biohub/waveorder](https://github.com/biohub/waveorder). +3. T. Chandler, I. E. Ivanov, G. Sturm, et al., "WaveOrder: A differentiable + wave-optical framework for scalable biological microscopy with diverse + modalities," [arXiv:2412.09775](https://arxiv.org/abs/2412.09775). +4. R. Vale, "Microscopy: Setting Up Koehler Illumination," + [iBiology Microscopy Series](https://www.ibiology.org/talks/set-up-koehler-illumination/). diff --git a/docs/guide/labelfree-screening/make_pdf.sh b/docs/guide/labelfree-screening/make_pdf.sh new file mode 100755 index 00000000..d54c8d8e --- /dev/null +++ b/docs/guide/labelfree-screening/make_pdf.sh @@ -0,0 +1,9 @@ +#!/bin/bash +# Render the 2-page whitepaper PDF from whitepaper.tex. +# Needs only base texlive packages (pdflatex, fancyvrb, booktabs, caption). +set -euo pipefail +cd "$(dirname "$0")" +pdflatex -interaction=nonstopmode whitepaper.tex >/dev/null +pdflatex -interaction=nonstopmode whitepaper.tex >/dev/null +rm -f whitepaper.aux whitepaper.log whitepaper.out +pdfinfo whitepaper.pdf | grep Pages diff --git a/docs/guide/labelfree-screening/whitepaper.pdf b/docs/guide/labelfree-screening/whitepaper.pdf new file mode 100644 index 00000000..e8f1cfcb Binary files /dev/null and b/docs/guide/labelfree-screening/whitepaper.pdf differ diff --git a/docs/guide/labelfree-screening/whitepaper.tex b/docs/guide/labelfree-screening/whitepaper.tex new file mode 100644 index 00000000..709e4397 --- /dev/null +++ b/docs/guide/labelfree-screening/whitepaper.tex @@ -0,0 +1,240 @@ +% WaveOrder label-free screening whitepaper: 2-page print layout. +% The docs-site version of this tutorial is labelfree-screening.md; this +% file re-lays the same content and figures for a compact PDF. +% Uses only base texlive packages (no tcolorbox/titlesec on the build host). +\documentclass[9pt]{extarticle} +\usepackage[letterpaper,margin=1.25cm,bottom=1.8cm]{geometry} +\usepackage[T1]{fontenc} +\usepackage{textcomp} +\usepackage{helvet} +\renewcommand{\familydefault}{\sfdefault} +\usepackage{graphicx} +\usepackage{xcolor} +\usepackage{fancyvrb} +\usepackage{booktabs} +\usepackage[labelfont=bf,labelsep=period,font=small,skip=3pt]{caption} +\usepackage{enumitem} +\setlist{nosep,leftmargin=1.2em} +\usepackage{fancyhdr} +\usepackage[colorlinks=true,linkcolor=black,urlcolor=blue!60!black]{hyperref} +\setlength{\parindent}{0pt} +\setlength{\parskip}{4pt} +\pagestyle{fancy} +\fancyhf{} +\renewcommand{\headrulewidth}{0pt} +\fancyfoot[L]{\raisebox{-3pt}{\includegraphics[height=13pt]{brand/biohub-logo-black.png}}} +\fancyfoot[R]{\small Computational Imaging Team \textbar{} \texttt{\{talon.chandler,shalin.mehta\}@biohub.org} \textbar{} 2026-11-13} + +\definecolor{boxgray}{gray}{0.96} +\definecolor{boxframe}{gray}{0.75} +\newcommand{\mysection}[1]{\vspace{5pt}{\Large\bfseries #1}\par\vspace{2pt}} +\setlength{\fboxsep}{6pt} +\setlength{\fboxrule}{0.4pt} + +% Verbatim blocks saved up front so they can live inside boxes. +\begin{SaveVerbatim}{cfgyml} +input_channel_names: [BF] +reconstruction_dimension: 2 +phase: + transfer_function: + yx_pixel_size: 0.325 # 6.5 um / 20x + z_pixel_size: 2.0 + wavelength_illumination: 0.450 + index_of_refraction_media: 1.0 + numerical_aperture_detection: 0.55 + numerical_aperture_illumination: 0.4 + invert_phase_contrast: false + z_focus_offset: auto + apply_inverse: + regularization_strength: 0.001 +\end{SaveVerbatim} + +\begin{SaveVerbatim}{tiltcfg} + z_focus_offset: auto + tilt_angle_zenith: {init: 0.13, lr: 0.05} + tilt_angle_azimuth: {init: 0.0, lr: 0.05} +optimization: {max_iterations: 50, method: adam} +\end{SaveVerbatim} + +\begin{document} + +{\huge\bfseries WaveOrder for label-free high-throughput screening} + +\vspace{5pt} +This guide demonstrates WaveOrder for label-free high-throughput +screening using optical pooled screening (OPS) data from the +\href{https://www.biorxiv.org/content/10.64898/2026.06.01.728087v1}{Biohub's multimodal +perturbation atlas} [1], where phase from defocus enabled label-free tracking and +phenotyping of CRISPR-perturbed A549 cells. Brightfield +defocus stacks from a standard motorized widefield microscope, together with a +\href{https://github.com/biohub/waveorder}{WaveOrder} config [2], yield +quantitative phase at screening scale (Figure~1). + +After reading this guide and acquiring data, use the +\href{https://github.com/biohub/waveorder/blob/main/docs/guide/first-reconstruction.md}{first-reconstruction +guide} and connect your agent to the +\href{https://biohub.ai/mcp}{Biohub MCP} to run the +\texttt{waveorder-tutorial} skill. Methods are described in the +\href{https://arxiv.org/abs/2412.09775}{WaveOrder preprint} [3]. + +\begin{figure}[h!] +\includegraphics[width=\linewidth]{figures/fig_overview.pdf} +\caption{Key acquisition settings (left) produce a defocus stack with +defocus-dependent contrast that WaveOrder converts into quantitative phase.} +\end{figure} + +\mysection{Acquisition} + +The three settings that matter most are: +\begin{itemize} +\item \textbf{Koehler illumination.} Align the condenser before every + session. Misalignment is the most common source of artifacts. Ron Vale's + five-minute + \href{https://www.ibiology.org/talks/set-up-koehler-illumination/}{iBiology + tutorial} [4] is a good refresher. +\item \textbf{Condenser NA $\approx$ 2/3 of detection NA.} We used a 0.4 NA condenser and + a 0.55 NA objective. A larger condenser aperture reduces contrast, while a + smaller condenser aperture reduces contrast and resolution. +\item \textbf{Z spacing and range.} At 20$\times$, start at 2 \textmu m + spacing and cover the full cell height with visibly blurred features on either end of the stack (here, 7 slices = 12 \textmu m). + For different objectives, scale inversely with magnification: e.g.\, 4 \textmu m at 10$\times$, 1 \textmu m + at 40$\times$. +\end{itemize} + +\begin{table}[h!] +\small +\centering +\begin{tabular}{llp{2.7in}} +\toprule +Parameter & Perturbation atlas value & Notes \\ +\midrule +Objective & 20$\times$, 0.55 NA, air & any objective works; match the config \\ +Condenser & 0.4 NA & $\approx$2/3 of detection NA, Koehler aligned \\ +Z stack & 7 slices, 2 \textmu m & confirm defocused features on top and bottom slices; use more slices for thicker samples \\ +Camera & 6.5 \textmu m pixels & 6.5 \textmu m / 20$\times$ = 0.325 \textmu m at the sample \\ +Wavelength & 450 nm & any wavelength is fine; ideally use a 10 nm notch filter \\ +Exposure & a few ms & fill 1/2--2/3 of the dynamic range \\ +\bottomrule +\end{tabular} +\caption{Acquisition parameters for the \href{https://www.biorxiv.org/content/10.64898/2026.06.01.728087v1}{Biohub's multimodal perturbation atlas dataset.}} +\end{table} + +\mysection{Reconstruction} + +\begin{minipage}[t]{0.55\linewidth} +Convert to OME-Zarr with +\href{https://github.com/czbiohub-sf/iohub}{iohub} +(\texttt{iohub convert -i -o data.zarr}), record the acquisition +parameters in a config, reconstruct one position, and compare against the +raw data in napari: + +\vspace{3pt} +\begin{Verbatim}[fontsize=\small] +waveorder reconstruct \ + -i data.zarr/A/1/029029 \ + -c config.yml -o phase.zarr +waveorder view data.zarr/A/1/029029 \ + phase.zarr/A/1/029029 +\end{Verbatim} +\vspace{3pt} + +Dense structures (nucleoli, nuclear envelopes, cell margins) should appear +brighter than background, as in Figure~1. \texttt{reconstruction\_dimension:~2} returns one +phase image per defocus stack, the best choice for samples that are thin compared to the depth of field of the microscope. If the sample appears to blur all at once as you defocus, use 2. If different regions of the sample blur at different defocus positions, consider trialling 3. +\end{minipage}\hfill +\begin{minipage}[t]{0.42\linewidth} +\vspace{0pt} +\fcolorbox{boxframe}{boxgray}{% +\begin{minipage}{\dimexpr\linewidth-2\fboxsep-2\fboxrule} +{\small\texttt{\$ cat config.yml}}\par\vspace{1pt} +{\small\UseVerbatim{cfgyml}} +\end{minipage}} +\end{minipage} + +\vspace{8pt} +\noindent +\begin{minipage}[t]{0.555\linewidth} +\vspace{0pt}% +\fcolorbox{boxframe}{boxgray}{% +\begin{minipage}{\dimexpr\linewidth-2\fboxsep-2\fboxrule} +{\small\textbf{Failure mode 1 \textbar{} Flipped z axis}}\par\vspace{2pt} +\centering +\includegraphics[width=0.97\linewidth]{figures/fig_flip.pdf} +\par\vspace{2pt}\raggedright\small +Stages disagree on the z sign convention---is the first slice above or below focus? If the convention is flipped, contrast +will be inverted and downstream comparison will be challenging. Dense organelles should be +bright. Fix: \texttt{invert\_phase\_contrast: true}. Check once per +microscope by reconstructing one stack both ways. +\end{minipage}} + +\vspace{6pt} +\fcolorbox{boxframe}{boxgray}{% +\begin{minipage}{\dimexpr\linewidth-2\fboxsep-2\fboxrule} +{\small\textbf{Failure mode 3 \textbar{} Autofocus drift}}\par\vspace{2pt} +\centering +\includegraphics[width=0.97\linewidth]{figures/fig_focus.pdf} +\par\vspace{2pt}\raggedright\small +When the stack center (e.g. the 4th slice in a 7 slice stack) is not in focus, the reconstruction will be blurry. +Setting \texttt{z\_focus\_offset:~auto} runs +the search that can recover the defocus offset. +\end{minipage}} +\end{minipage}\hfill +\begin{minipage}[t]{0.425\linewidth} +\vspace{0pt}% +\fcolorbox{boxframe}{boxgray}{% +\begin{minipage}{\dimexpr\linewidth-2\fboxsep-2\fboxrule} +{\small\textbf{Failure mode 2 \textbar{} Incorrect regularization}}\par\vspace{2pt} +\centering +\includegraphics[width=0.97\linewidth]{figures/fig_reg.pdf} +\par\vspace{2pt}\raggedright\small +\texttt{regularization\_strength} balances noise and resolution and varies by +orders of magnitude across samples and exposures. Sweep decades on one +position and choose manually. Too small +amplifies blotchy low-frequency noise; too large leaves a washed-out image. The optimum is usually within a decade of +\texttt{1e-3}. +\end{minipage}} + +\vspace{6pt} +\fcolorbox{boxframe}{boxgray}{% +\begin{minipage}{\dimexpr\linewidth-2\fboxsep-2\fboxrule} +{\small\textbf{Failure mode 4 \textbar{} Tilted illumination at the well edge}}\par\vspace{2pt} +\centering +\includegraphics[width=0.97\linewidth]{figures/fig_tilt.pdf} +\par\vspace{2pt}\raggedright\small +Misaligned Koehler and imaging near a well-edge meniscus tilts the illumination, making cells look shadow-cast +with reduced contrast. Fit the tilt by gradient descent in the same config: +{\footnotesize\UseVerbatim{tiltcfg}} +\end{minipage}} +\end{minipage} + +\mysection{Scaling up} + +A single config file should generate good reconstructions across a screen. We suggest parallelizing +\texttt{waveorder reconstruct} over positions (one process or cluster job +per position) with \texttt{device:~cuda:0} to reconstruct on a GPU. We recommend +\texttt{z\_focus\_offset:~auto} to find the focus of each position. + +We recommend manual supervision +to set up the optimized \texttt{tilt\_angle\_\{zenith,azimuth\}} parameters across entire wells. We find that initializing these tilt parameters at good initial tilts generates the stablest reconstructions, so we recommend inspecting +reconstructions at the well edges, then refining the initial values until you reach stable reconstructions that pass QC. + +\mysection{References} + +\sloppy +\begin{enumerate} +\item C.~Liu, A.~V.~Hillsley, M.~Sekhar, et al., ``A Multimodal Perturbation + Atlas Defines the Phenotypic Resolution of Cellular Morphology,'' + \href{https://www.biorxiv.org/content/10.64898/2026.06.01.728087v1}{bioRxiv + 2026.06.01.728087}. +\item WaveOrder repository: + \href{https://github.com/biohub/waveorder}{github.com/biohub/waveorder}. +\item T.~Chandler, I.~E.~Ivanov, G.~Sturm, et al., + ``WaveOrder: A differentiable wave-optical framework for scalable + biological microscopy with diverse modalities,'' + \href{https://arxiv.org/abs/2412.09775}{arXiv:2412.09775}. +\item R.~Vale, ``Microscopy: Setting Up Koehler Illumination,'' + \href{https://www.ibiology.org/talks/set-up-koehler-illumination/}{iBiology + Microscopy Series}. +\end{enumerate} + +\end{document} diff --git a/tests/api_tests/test_auto_focus_offset.py b/tests/api_tests/test_auto_focus_offset.py new file mode 100644 index 00000000..f92ef4c8 --- /dev/null +++ b/tests/api_tests/test_auto_focus_offset.py @@ -0,0 +1,86 @@ +"""Tests for z_focus_offset: 'auto' resolution.""" + +import pytest + +from waveorder.api import phase +from waveorder.optim._types import OptimizableFloat + +SIM_SETTINGS = phase.Settings( + transfer_function=dict( + yx_pixel_size=0.325, + z_pixel_size=2.0, + wavelength_illumination=0.45, + index_of_refraction_media=1.0, + numerical_aperture_detection=0.55, + numerical_aperture_illumination=0.4, + z_focus_offset=2.0, + ), + apply_inverse=dict(regularization_strength=1e-3), +) + + +def _settings_with_offset(offset): + dumped = SIM_SETTINGS.model_dump() + dumped["transfer_function"]["z_focus_offset"] = offset + return phase.Settings.model_validate(dumped) + + +@pytest.fixture(scope="module") +def defocused_data(): + """Simulated 7-slice stack whose focus sits 2 slices from center.""" + _, data = phase.simulate(settings=SIM_SETTINGS, recon_dim=2, zyx_shape=(7, 128, 128)) + return data + + +def test_auto_accepted_in_settings(): + settings = _settings_with_offset("auto") + assert settings.transfer_function.z_focus_offset == "auto" + + settings = _settings_with_offset({"init": "auto", "lr": 0.05}) + offset = settings.transfer_function.z_focus_offset + assert isinstance(offset, OptimizableFloat) + assert offset.initial_value == "auto" + + +def test_unresolved_auto_value_raises(): + with pytest.raises(ValueError, match="resolved from the data"): + OptimizableFloat(init="auto", lr=0.05).value + + +def test_simulate_rejects_auto(): + with pytest.raises(ValueError, match="measured data"): + phase.simulate(settings=_settings_with_offset("auto"), recon_dim=2, zyx_shape=(5, 32, 32)) + + +def test_resolve_recovers_integer_offset(defocused_data): + resolved = phase.resolve_z_focus_offset(defocused_data, _settings_with_offset("auto")) + offset = resolved.transfer_function.z_focus_offset + assert isinstance(offset, float) + assert offset == 2.0 + + +def test_resolve_keeps_learning_rate(defocused_data): + settings = _settings_with_offset(OptimizableFloat(init="auto", lr=0.05)) + resolved = phase.resolve_z_focus_offset(defocused_data, settings) + offset = resolved.transfer_function.z_focus_offset + assert isinstance(offset, OptimizableFloat) + assert offset.initial_value == 2.0 + assert offset.learning_rate == 0.05 + + +def test_resolve_passes_through_numeric(defocused_data): + settings = _settings_with_offset(1.5) + resolved = phase.resolve_z_focus_offset(defocused_data, settings) + assert resolved is settings + + +def test_compute_transfer_function_drops_auto_in_3d(defocused_data): + with pytest.warns(UserWarning, match="3D reconstructions do not use a focus offset"): + phase.compute_transfer_function(defocused_data, recon_dim=3, settings=_settings_with_offset("auto")) + + +def test_reconstruct_resolves_auto(defocused_data): + resolved = phase.resolve_z_focus_offset(defocused_data, _settings_with_offset("auto")) + fixed = phase.reconstruct(defocused_data, recon_dim=2, settings=resolved) + auto = phase.reconstruct(defocused_data, recon_dim=2, settings=_settings_with_offset("auto")) + assert abs(auto - fixed).max() == 0.0 diff --git a/tests/models/test_isotropic_thin_3d.py b/tests/models/test_isotropic_thin_3d.py index 75039001..cba8fba0 100644 --- a/tests/models/test_isotropic_thin_3d.py +++ b/tests/models/test_isotropic_thin_3d.py @@ -123,3 +123,45 @@ def test_reconstruct_apodization_rolloff_smoke(): nyquist_row_apod = np.abs(np.fft.fft2(phase_apod.numpy())[16, :]) nyquist_row_hard = np.abs(np.fft.fft2(phase_hard.numpy())[16, :]) assert nyquist_row_apod.max() < 1e-3 * nyquist_row_hard.max() + + +def test_gram_singular_system_matches_svd(): + """The gradient-mode decomposition reconstructs identically to the SVD.""" + torch.manual_seed(0) + zyx_data = torch.rand(4, 32, 32) + 10.0 + absorption_tf, phase_tf = isotropic_thin_3d.calculate_transfer_function( + yx_shape=(32, 32), + yx_pixel_size=6.5 / 20, + z_position_list=[-2, -1, 0, 1], + wavelength_illumination=0.5, + index_of_refraction_media=1.0, + numerical_aperture_illumination=0.4, + numerical_aperture_detection=0.55, + ) + + results = [] + for use_svd in (True, False): + singular_system = isotropic_thin_3d.calculate_singular_system(absorption_tf, phase_tf, use_svd=use_svd) + results.append( + isotropic_thin_3d.apply_inverse_transfer_function(zyx_data, singular_system, regularization_strength=1e-3) + ) + for svd_result, gram_result in zip(*results): + assert torch.allclose(svd_result, gram_result, atol=1e-4) + + +def test_gram_singular_system_finite_gradients(): + """Gradients stay finite where singular values are degenerate or zero.""" + z_focus_offset = torch.tensor(0.0, requires_grad=True) + z_positions = (-torch.arange(4) + 2 + z_focus_offset) * 0.25 + _, phase_yx = isotropic_thin_3d.reconstruct( + torch.rand(4, 32, 32) + 10.0, + yx_pixel_size=0.1, + z_position_list=z_positions, + wavelength_illumination=0.532, + index_of_refraction_media=1.3, + numerical_aperture_illumination=0.9, + numerical_aperture_detection=1.2, + ) + loss = (phase_yx**2).mean() + loss.backward() + assert torch.isfinite(z_focus_offset.grad) diff --git a/waveorder/api/phase.py b/waveorder/api/phase.py index 24ee9341..2eb7aee3 100644 --- a/waveorder/api/phase.py +++ b/waveorder/api/phase.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from collections.abc import Mapping from typing import Literal, Union @@ -28,6 +29,7 @@ _wrap_output_tensor, ) from waveorder.device import resolve_device +from waveorder.focus import compute_midband_power from waveorder.models import isotropic_thin_3d, phase_thick_3d from waveorder.optim import ( NullLogger, @@ -36,6 +38,7 @@ optimize_reconstruction, ) from waveorder.optim._types import OptimizableFloat +from waveorder.optim.autoreg import select_crop from waveorder.optim.losses import MidbandPowerLossSettings, build_loss_fn # --- Settings --- @@ -45,6 +48,15 @@ class TransferFunctionSettings( OptimizableFourierTransferFunctionSettings, WavelengthIllumination, ): + z_focus_offset: Union[Literal["auto"], float, OptimizableFloat] = Field( + default=0, + description=( + "(optimizable) offset from center slice in slice units; 'auto' finds the " + "in-focus slice from the data with a per-slice search over a 3D " + "reconstruction; '{init: auto, lr: 0.05}' refines the searched slice to " + "sub-slice precision with the gradient optimizer" + ), + ) numerical_aperture_illumination: Union[NonNegativeFloat, OptimizableFloat] = Field( default=0.9, description="(optimizable) condenser numerical aperture" ) @@ -74,6 +86,141 @@ class Settings(MyBaseModel): # --- Functions --- +def _z_focus_offset_is_auto(transfer_function: TransferFunctionSettings) -> bool: + """Whether ``z_focus_offset`` asks to be estimated from the data.""" + v = transfer_function.z_focus_offset + return v == "auto" or (isinstance(v, OptimizableFloat) and v.initial_value == "auto") + + +def _drop_auto_z_focus_offset(settings: Settings, reason: str) -> Settings: + """Warn and replace an unused ``z_focus_offset: "auto"`` with 0.""" + warnings.warn(f"ignoring z_focus_offset: 'auto': {reason}", UserWarning) + return Settings( + transfer_function=settings.transfer_function.model_copy(update={"z_focus_offset": 0.0}), + apply_inverse=settings.apply_inverse, + ) + + +def _select_scoring_crops(zyx: torch.Tensor, crop_size: int = 512) -> list[tuple[int, int, int]]: + """Spatially spread crops for focus scoring: one per quadrant plus center. + + A single most-structured crop is fragile: bright floating debris + dominates a variance score while sitting in a different focal plane + than the sample. Several spread-out crops, each the most structured + within its region, let the median focus index reject local outliers. + """ + _, Y, X = zyx.shape + regions = [ + (0, Y // 2, 0, X // 2), + (0, Y // 2, X // 2, X), + (Y // 2, Y, 0, X // 2), + (Y // 2, Y, X // 2, X), + (Y // 4, Y // 4 + Y // 2, X // 4, X // 4 + X // 2), + ] + crops = set() + for y_lo, y_hi, x_lo, x_hi in regions: + if y_hi - y_lo < 1 or x_hi - x_lo < 1: + continue + y0, x0, size = select_crop(zyx[:, y_lo:y_hi, x_lo:x_hi], crop_size=crop_size) + crops.add((y_lo + y0, x_lo + x0, size)) + return sorted(crops) + + +def resolve_z_focus_offset( + czyx_data: xr.DataArray, + settings: Settings = None, + device: str | torch.device | None = None, +) -> Settings: + """Resolve ``z_focus_offset: "auto"`` from the data. + + Reconstructs 3D phase volumes on a handful of structured crops of + the data, scores each slice's midband spectral power, and takes the + median sharpest slice across crops as the focus: + ``z_focus_offset = focus_idx - Z // 2``. + + This discrete integer search is robust to multi-slice focus errors, + which gradient optimization cannot fix: shifting the model focus by + a whole number of slices relabels which acquired plane is "in + focus" and still reconstructs a locally sharp image, so the + sharpness landscape has a local optimum at every integer offset. + For sub-slice precision, write + ``z_focus_offset: {init: auto, lr: 0.05}``: the integer search runs + first and seeds the gradient optimizer, which refines within the + correct basin and co-optimizes any other optimizable parameters + (e.g. tilt angles). + + Parameters + ---------- + czyx_data : xr.DataArray + CZYX brightfield data (single channel). + settings : Settings, optional + Phase settings. Returned unchanged unless ``z_focus_offset`` is + ``"auto"``, either bare or as an ``OptimizableFloat`` init. + device : str, torch.device, or None + Compute device. None = CPU, "auto" = best available. + + Returns + ------- + Settings + Copy of ``settings`` with ``z_focus_offset`` resolved to a float + (or to an ``OptimizableFloat`` with a float init). + """ + if settings is None: + settings = Settings() + if not _z_focus_offset_is_auto(settings.transfer_function): + return settings + + Z = czyx_data.shape[1] + if Z == 1: + offset = 0.0 + focus_indices = [] + else: + # The focus reconstruction needs a numeric offset (its value does not + # enter the 3D model) and enough z padding to keep wrap-around from + # biasing the per-slice score. + tf_3d = settings.transfer_function.model_copy( + update={ + "z_focus_offset": 0.0, + "z_padding": max(settings.transfer_function.z_padding, Z // 2), + } + ) + settings_3d = Settings(transfer_function=tf_3d, apply_inverse=settings.apply_inverse) + s = tf_3d.resolve_floats() + yx = YXPixelSize.from_value(s.yx_pixel_size) + + zyx = torch.as_tensor(np.asarray(czyx_data.values[0], dtype=np.float32)) + focus_indices = [] + for y0, x0, size in _select_scoring_crops(zyx): + czyx_crop = czyx_data.isel(y=slice(y0, y0 + size), x=slice(x0, x0 + size)) + phase_zyx = reconstruct(czyx_crop, recon_dim=3, settings=settings_3d, device=device).values[0] + midband_per_slice = ( + compute_midband_power( + torch.as_tensor(phase_zyx), + NA_det=s.numerical_aperture_detection, + lambda_ill=s.wavelength_illumination, + pixel_size=(yx.y + yx.x) / 2, + ) + .cpu() + .numpy() + ) + focus_indices.append(int(np.argmax(midband_per_slice))) + offset = float(int(np.median(focus_indices)) - (Z // 2)) + + current = settings.transfer_function.z_focus_offset + if isinstance(current, OptimizableFloat): + resolved = OptimizableFloat(initial_value=offset, learning_rate=current.learning_rate) + else: + resolved = offset + print( + f"Resolved z_focus_offset = {offset:g} (slice units) from the data" + + (f"; per-crop focus slices {focus_indices}" if focus_indices else "") + ) + return Settings( + transfer_function=settings.transfer_function.model_copy(update={"z_focus_offset": resolved}), + apply_inverse=settings.apply_inverse, + ) + + def simulate( settings: Settings = None, recon_dim: Literal[2, 3] = 3, @@ -108,6 +255,8 @@ def simulate( """ if settings is None: settings = Settings() + if _z_focus_offset_is_auto(settings.transfer_function): + raise ValueError("z_focus_offset: 'auto' is estimated from measured data; set a numeric value for simulate()") s = settings.transfer_function.resolve_floats() yx_pixel_size = YXPixelSize.from_value(s.yx_pixel_size) @@ -220,6 +369,12 @@ def compute_transfer_function( settings = Settings() device = resolve_device(device) + if _z_focus_offset_is_auto(settings.transfer_function): + if recon_dim == 3: + settings = _drop_auto_z_focus_offset(settings, "3D reconstructions do not use a focus offset") + else: + settings = resolve_z_focus_offset(czyx_data, settings, device=device) + zyx_shape = czyx_data.shape[1:] # CZYX -> ZYX s = settings.transfer_function.resolve_floats() @@ -386,6 +541,11 @@ def reconstruct( """ if settings is None: settings = Settings() + if _z_focus_offset_is_auto(settings.transfer_function): + if recon_dim == 3: + settings = _drop_auto_z_focus_offset(settings, "3D reconstructions do not use a focus offset") + else: + settings = resolve_z_focus_offset(czyx_data, settings, device=device) tf = compute_transfer_function(czyx_data, recon_dim, settings, device=device) return apply_inverse_transfer_function(czyx_data, tf, recon_dim, settings, device=device) @@ -446,6 +606,12 @@ def optimize( settings = Settings() device = resolve_device(device) + if _z_focus_offset_is_auto(settings.transfer_function): + if recon_dim == 3: + settings = _drop_auto_z_focus_offset(settings, "3D reconstructions do not use a focus offset") + else: + settings = resolve_z_focus_offset(czyx_data, settings, device=device) + opt_params, _ = extract_optimizable_params(settings.transfer_function) if not opt_params: diff --git a/waveorder/cli/reconstruct.py b/waveorder/cli/reconstruct.py index 3b49972e..32f05b6d 100644 --- a/waveorder/cli/reconstruct.py +++ b/waveorder/cli/reconstruct.py @@ -35,6 +35,12 @@ def _reconstruct_cli( If any parameter has an `lr` key (OptimizableFloat), an optimization loop runs before the standard reconstruction pipeline. + If `z_focus_offset` is 'auto' (bare, or as an `init`), the in-focus slice is + found with a per-slice search over a 3D reconstruction of the first + position, and the resolved settings are written alongside the config. As + `{init: auto, lr: 0.05}`, the searched slice seeds the optimization loop, + which refines it to sub-slice precision. + If `apply_inverse` has an `auto_regularization` block, `regularization_strength` is chosen from the data after the transfer function is computed, and the resolved settings are written alongside the config. @@ -58,6 +64,12 @@ def _reconstruct_cli( settings = utils.yaml_to_model(config_filepath, ReconstructionSettings) + # Resolve z_focus_offset: 'auto' before anything consumes the offset (the + # optimizer needs a numeric init and the transfer function a numeric + # offset), and freeze the resolved value into a sibling config. + if _has_auto_z_focus_offset(settings): + config_filepath = _run_auto_z_focus_offset(settings, input_position_dirpaths[0], config_filepath) + # Check for optimizable parameters and run optimization if needed if has_optimizable_params(settings): config_filepath = _run_optimization(settings, input_position_dirpaths[0], config_filepath) @@ -111,6 +123,55 @@ def _reconstruct_cli( ) +def _has_auto_z_focus_offset(settings) -> bool: + """Whether the phase block asks for an automatic focus-offset search. + + 3D reconstructions do not use a focus offset, so 'auto' is left for the + API to warn about and drop. + """ + from waveorder.api import phase + + return ( + settings.phase is not None + and settings.reconstruction_dimension == 2 + and phase._z_focus_offset_is_auto(settings.phase.transfer_function) + ) + + +def _run_auto_z_focus_offset(settings, input_position_dirpath, config_filepath): + """Resolve z_focus_offset from the data, then freeze it into a sibling config. + + Returns the path of the resolved config, which the rest of the pipeline uses. + """ + # Deferred imports for the same reason as above + from iohub.ngff import open_ome_zarr + + from waveorder.api import phase + from waveorder.cli.printing import echo_headline + from waveorder.cli.utils import resolve_time_indices + from waveorder.io import utils + + echo_headline("\nResolving z_focus_offset with a per-slice focus search...") + + input_dataset = open_ome_zarr(input_position_dirpath, layout="fov", mode="r") + try: + t_idx = resolve_time_indices(settings.time_indices, input_dataset.data.shape[0])[0] + czyx = input_dataset.to_xarray().isel(t=t_idx).sel(c=settings.input_channel_names) + finally: + input_dataset.close() + + settings.phase = phase.resolve_z_focus_offset(czyx, settings.phase, device=settings.device) + + offset = settings.phase.transfer_function.z_focus_offset + typer.echo(typer.style(f" z_focus_offset = {float(offset):g} (slice units), t={t_idx}", fg="green")) + + resolved_path = config_filepath.parent / (config_filepath.stem + "_autofocus.yml") + utils.model_to_yaml(settings, resolved_path) + print(f"Auto-focus settings saved to {resolved_path}") + + return resolved_path + + def _run_optimization(settings, input_position_dirpath, config_filepath): """Run parameter optimization before standard reconstruction.""" # Deferred imports for the same reason as above diff --git a/waveorder/models/isotropic_thin_3d.py b/waveorder/models/isotropic_thin_3d.py index 0f958413..3db61d2a 100644 --- a/waveorder/models/isotropic_thin_3d.py +++ b/waveorder/models/isotropic_thin_3d.py @@ -229,9 +229,10 @@ def calculate_singular_system( phase_2d_to_3d_transfer_function : Tensor Transfer function for phase, same shape as absorption TF use_svd : bool - If True (default), use torch.linalg.svd for best reconstruction - accuracy. Set to False for a norm-based decomposition that - supports gradient backpropagation on all devices + If True (default), use torch.linalg.svd. Set to False for an + equivalent closed-form decomposition of the 2x2 Gram matrix with + epsilon-smoothed branch points, which supports gradient + backpropagation on all devices Returns ------- @@ -257,11 +258,10 @@ def calculate_singular_system( if use_svd: # Full SVD over the (s=2, Z) matrix at each spatial frequency. - # Captures cross-channel coupling between absorption and phase, - # giving the best reconstruction accuracy. - # Note: torch.linalg.svd backward fails with complex tensors on - # some GPU types due to singular vector phase ambiguity - # (svd_backward: e^{i phi} error). Use use_svd=False when + # Note: torch.linalg.svd backward produces non-finite gradients + # where singular values are degenerate (e.g. both zero beyond the + # band limit) and fails outright with complex tensors on some GPU + # types (svd_backward: e^{i phi} error). Use use_svd=False when # gradients are needed. BVyVxsZ = sfYX.permute(0, 3, 4, 1, 2) Up, Sp, Vhp = torch.linalg.svd(BVyVxsZ, full_matrices=False) @@ -269,17 +269,67 @@ def calculate_singular_system( S = Sp.permute(0, 3, 1, 2) # (B, s, Vy, Vx) Vh = Vhp.permute(0, 3, 4, 1, 2) # (B, s, Z, Vy, Vx) else: - # Norm-based decomposition: assumes absorption and phase channels - # are independent (U=identity). Less accurate than full SVD but - # supports gradient backpropagation on all devices. - # Per-channel norms: S[b, k] = norm(H[b, k, :]) - S = torch.sqrt(torch.clamp(torch.sum(torch.abs(sfYX) ** 2, dim=2), min=1e-12)) # (B, s=2, Vy, Vx) - # Normalized rows: Vh[b, k, z] = H[b, k, z] / S[b, k] - Vh = sfYX / (S[:, :, None] + 1e-12) # (B, s=2, Z, Vy, Vx) - # U = identity (each channel reconstructs independently) - U = torch.zeros(B, s, s, Vy, Vx, dtype=sfYX.dtype, device=sfYX.device) - for i in range(s): - U[:, i, i] = 1.0 + # Closed-form singular system from the 2x2 Gram matrix G = H H^†. + # This is the same decomposition the SVD branch computes (up to + # the phase gauge, which the reconstruction is invariant to), but + # with epsilon-smoothed branch points so gradients stay finite + # where the singular values are degenerate or zero, where + # torch.linalg.svd backward does not. Works on all devices. + absorption_row = sfYX[:, 0] # (B, Z, Vy, Vx) + phase_row = sfYX[:, 1] + # G = [[a, b], [conj(b), c]] at each spatial frequency. The Gram + # entries square the transfer function, so accumulate and solve the + # (Z-free) eigenproblem in double precision to keep the + # decomposition tight in float32 overall. + a = (absorption_row.abs() ** 2).sum(dim=1, dtype=torch.float64) # (B, Vy, Vx) + c = (phase_row.abs() ** 2).sum(dim=1, dtype=torch.float64) + b = (absorption_row * phase_row.conj()).sum(dim=1, dtype=torch.complex128) + + mean = (a + c) / 2 + # Smooth the eigenvalue gap with a relative epsilon (so healthy + # frequencies are perturbed by <1e-9 relative) plus a tiny absolute + # floor that keeps sqrt backward finite where H is exactly zero. + tiny = 1e-30 + half_gap = torch.sqrt(((a - c) / 2) ** 2 + b.abs() ** 2 + (1e-9 * mean) ** 2 + tiny) + lam1 = mean + half_gap + lam2 = torch.clamp(mean - half_gap, min=0.0) + S = torch.stack((torch.sqrt(lam1 + tiny), torch.sqrt(lam2 + tiny)), dim=1).to( + absorption_row.real.dtype + ) # (B, 2, Vy, Vx) + + # Leading eigenvector of G, taking whichever closed form is better + # conditioned at each frequency: [b, lam1 - a] or [lam1 - c, conj(b)]. + # Where H is ~zero both forms vanish and the normalized vector decays + # to zero along with S, so no explicit degenerate branch is needed. + u_x, u_y = b, (lam1 - a).to(b.dtype) + w_x, w_y = (lam1 - c).to(b.dtype), b.conj() + use_u = (u_x.abs() ** 2 + u_y.abs() ** 2) >= (w_x.abs() ** 2 + w_y.abs() ** 2) + v1_x = torch.where(use_u, u_x, w_x) + v1_y = torch.where(use_u, u_y, w_y) + norm1 = torch.sqrt(v1_x.abs() ** 2 + v1_y.abs() ** 2 + tiny) + v1_x = v1_x / norm1 + v1_y = v1_y / norm1 + + v1_x = v1_x.to(sfYX.dtype) + v1_y = v1_y.to(sfYX.dtype) + + # Columns of U: v1 and its orthonormal complement [-conj(v1_y), conj(v1_x)] + U = torch.stack( + ( + torch.stack((v1_x, -v1_y.conj()), dim=1), + torch.stack((v1_y, v1_x.conj()), dim=1), + ), + dim=1, + ) # (B, s, s, Vy, Vx) + + # Vh = S^-1 U^† H; rows have unit norm where S > 0 and vanish with S. + Vh = torch.stack( + ( + (v1_x.conj()[:, None] * absorption_row + v1_y.conj()[:, None] * phase_row) / S[:, 0, None], + (-v1_y[:, None] * absorption_row + v1_x[:, None] * phase_row) / S[:, 1, None], + ), + dim=1, + ) # (B, s, Z, Vy, Vx) if not batched: U = U.squeeze(0) @@ -547,8 +597,10 @@ def reconstruct( tilt_angle_azimuth=tilt_angle_azimuth, pupil_steepness=pupil_steepness, ) - # Use norm-based decomposition when gradients are needed (optimization), - # full SVD otherwise (better accuracy for final reconstruction) + # Both decompositions produce the same reconstruction, but + # torch.linalg.svd backward is non-finite at degenerate singular values + # (ubiquitous beyond the band limit) and unsupported for complex inputs + # on some GPUs, so gradient mode uses the closed-form Gram decomposition. needs_grad = absorption_tf.requires_grad or phase_tf.requires_grad singular_system = calculate_singular_system(absorption_tf, phase_tf, use_svd=not needs_grad) return apply_inverse_transfer_function( diff --git a/waveorder/optim/_types.py b/waveorder/optim/_types.py index 7d4e9111..7331653c 100644 --- a/waveorder/optim/_types.py +++ b/waveorder/optim/_types.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any, Union +from typing import Any, Literal, Union from pydantic import BaseModel, model_validator @@ -13,6 +13,10 @@ class OptimizableFloat(BaseModel): Accepts a bare float/int (backwards compatible, not optimizable) or a mapping with `init`/`initial_value` and `lr`/`learning_rate`. + ``initial_value: "auto"`` marks the value to be estimated from the data + before any numeric use; only fields whose resolution is implemented + (currently phase ``z_focus_offset``) accept it. + Examples -------- >>> OptimizableFloat(init=0.0, lr=0.01) @@ -21,7 +25,7 @@ class OptimizableFloat(BaseModel): OptimizableFloat(initial_value=3.0, learning_rate=0.5) """ - initial_value: float = 0.0 + initial_value: Union[float, Literal["auto"]] = 0.0 learning_rate: float = 0.01 @model_validator(mode="before") @@ -40,6 +44,11 @@ def _normalize_aliases(cls, data: Any) -> Any: @property def value(self) -> float: + if self.initial_value == "auto": + raise ValueError( + "initial_value: 'auto' must be resolved from the data before use; " + "'auto' is currently supported only for phase z_focus_offset" + ) return self.initial_value @property @@ -51,7 +60,7 @@ def is_optimizable(self) -> bool: return self.learning_rate > 0 def __float__(self) -> float: - return self.initial_value + return self.value def __repr__(self) -> str: return f"OptimizableFloat(initial_value={self.initial_value}, learning_rate={self.learning_rate})" @@ -85,9 +94,9 @@ def extract_optimizable_params(obj: Any, prefix: str = "") -> tuple[dict[str, tu if isinstance(obj, OptimizableFloat): key = prefix if obj.is_optimizable: - optimizable[key] = (obj.initial_value, obj.learning_rate) + optimizable[key] = (obj.value, obj.learning_rate) else: - fixed[key] = obj.initial_value + fixed[key] = obj.value return optimizable, fixed if isinstance(obj, BaseModel): diff --git a/waveorder/optim/optimize.py b/waveorder/optim/optimize.py index 588902be..88fe2472 100644 --- a/waveorder/optim/optimize.py +++ b/waveorder/optim/optimize.py @@ -287,9 +287,18 @@ def closure(): if torch.isnan(loss) or torch.isnan(recon).any(): raise ValueError("NaN in reconstruction or loss") + # These parameters produced a finite reconstruction, so they + # are safe to roll back to. Checkpointing after the step + # would save unvalidated values: a degenerate SVD gradient + # can step parameters to non-finite values, and rolling back + # to those crashes the final reconstruction. + last_good = {name: t.detach().clone() for name, t in param_tensors.items()} + if use_gradients: loss.backward() optimizer.step() + if any(not torch.isfinite(t).all() for t in param_tensors.values()): + raise ValueError("non-finite parameters after optimizer step") except (RuntimeError, ValueError): # SVD convergence failure or NaN — revert to last good params @@ -298,8 +307,10 @@ def closure(): t.copy_(last_good[name]) break - # Checkpoint good params - last_good = {name: t.detach().clone() for name, t in param_tensors.items()} + # Checkpoint post-step parameters (the lbfgs closure validates them + # during the step) only when they are finite. + if method == "lbfgs" and all(torch.isfinite(t).all() for t in param_tensors.values()): + last_good = {name: t.detach().clone() for name, t in param_tensors.items()} wall_times.append(time.monotonic() - t_start) loss_val = loss.item()