Skip to content

refactor!: waveorder as torch modules and models - #597

Draft
srivarra wants to merge 27 commits into
wo-4.0from
wo-4.0-nn-modules
Draft

refactor!: waveorder as torch modules and models#597
srivarra wants to merge 27 commits into
wo-4.0from
wo-4.0-nn-modules

Conversation

@srivarra

@srivarra srivarra commented Nov 14, 2026

Copy link
Copy Markdown
Contributor

waveorder becomes a set of torch modules, with one model per contrast. A model calibrates its optics, reconstructs, and simulates, and it can be distilled into a fast neural network. Most wo commands and wo-4.0 configs work as before.

import waveorder as wo
from waveorder.nn import Detection, Free, Gaussian, Illumination, Tikhonov

model = wo.models.Phase3D(
    Illumination(wavelength=0.532, numerical_aperture=Free(0.45, bounds=(0.3, 0.7))),
    Detection(wavelength=0.532, numerical_aperture=0.9),
    index_of_refraction_media=1.33,
    pixel_size=(0.4, 0.15),
    object_prior=Tikhonov(weight=4e-4),
    noise=Gaussian(sigma=1e-3),
)
model.calibrate(data)             # estimates the free illumination NA from a (C, Z, Y, X) stack
result = model.reconstruct(data)

What changes for users

  • Reconstructions match wo-4.0 on the regression suite and run in about half the time.
  • wo rec reads single-position stores, reconstructs large volumes in tiles with --tile, spreads time points over GPUs with --device cuda:0,cuda:1, and reuses a stored transfer function with -t tf.zarr.
  • wo distill indexes reconstructions and trains a network on them through Lightning's CLI.

Breaking changes

  • Python 3.12 to 3.14.
  • The 3.x Python API is gone. waveorder.focus.focus_from_transverse_band remains and warns.
  • A 3.x config must go through wo config migrate, which replaces wo migrate-config.
  • wo tile-stitch is now wo rec --tile.
  • Not in yet: birefringence, wo view, and the napari plugin behind wo interactive.

Status

I tested the reconstructions on neuromast and zebrafish data. I'm working on a cleanup pass.

@srivarra
refactor!: remove the 3.x code and its tests
1721f1a 
The package is rebuilt on torch modules under src/waveorder in the commits
that follow. This commit removes the old top-level waveorder/ package (the
functional models, the api and cli layers, the calibration and optim
modules, the napari plugin and GUI) and the tests written against it.

BREAKING CHANGE: the 3.x Python API and the napari plugin are gone. The new
API is waveorder.models; README.md lists what each 3.x name becomes.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
build: src layout, Python 3.12 to 3.14, new dependencies
22726a5 
The package moves to src/waveorder. Python is capped below 3.15 and torch
is 2.15 or newer. New required dependencies: lightning (pinned, its restart
behaviour is version-specific), timm, tensordict, jsonargparse and pyarrow.
wandb and the ONNX export stay optional extras. CI runs ruff and ty; the
clang-format step for the removed C sources is gone, and ruff skips docs/,
which still describes the 3.x API.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(nn): operators, optics and parameters as torch modules
c946edb 
waveorder 3.x computed transfer functions and inverses in free functions, one set per contrast. These are the
shared parts every imaging system and solver is built from: Fourier-diagonal operators, pupils and apertures,
named physical parameters with bounds, sampling, and the errors and registry they report through.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(nn): phase and fluorescence imaging systems
d0aaaec 
Each contrast is one module that turns its optical parameters into a transfer function on any sampling. Transfer
functions are built in slabs of planes and cropped to the data grid as they go, so a build needs about 4 times a
tile's memory instead of 40 and stays under cuFFT's 2^31 element limit. Results match wo-4.0 to rounding.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(nn): object priors, noise models and solvers
7d49e71 
A prior, a noise model and a solver are separate parts, so a new one works with every contrast. Covers the
closed-form Wiener solve, FISTA, ADMM and Richardson-Lucy with gradient consensus.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(nn): scores and search stages for calibration
45b9c72 
Calibration is a search over free optical parameters scored on the data: grid and gradient stages, the evidence and
contrast scores, automatic regularization weight, and bead calibration.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(nn): restoration networks and the network solver
6b66e52 
Networks that stand in for a solver after distillation: the U-Net family, the defocus-kernel stem, and the solver
that runs a trained bundle on data with the model's sampling checks.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(testing): contract checks for new parts
c826e6d 
waveorder.testing checks that a new system, solver, prior or noise model obeys the rules the solvers and models rely on
(adjoint, units, gradients, device and dtype), so parts written outside the package can be verified the same way.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(io): positions and plates through iohub, transfer-function store…
786aca6 
…s, bundles

All OME-Zarr access goes through iohub. A single-position store is accepted as input and gives a single-position
output. Transfer functions are stored in tf.zarr through tensordict so a fixed configuration computes them once.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(tiling): plan, run and merge overlapping tiles
0f55505 
Volumes too large for one solve are reconstructed tile by tile with cross-faded overlaps, reading ahead and
writing behind the GPU.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(config): version 4 schema, migration from 3.x, model building
046eac7 
Configs written by wo-4.0 load unchanged. wo config migrate turns a 3.x file into a version 4 file, and
solver.padding.z accepts auto, which picks the padding from the optics.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(data): parquet index of fields and the patch loader
084e5c7 
Training reads recorded reconstructions through an index with named splits held in the index, so hold-out plates
and wells stay fixed as data is added. The loader gives each DDP rank a disjoint shard and resumes mid-epoch.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(train): Lightning module, data module and callbacks
96139f9 
Distillation runs on Lightning's CLI, so configs name classes and a stopped run resumes exactly from its
checkpoint, on one GPU or many nodes.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(models)!: Phase and Fluorescence models with reconstruct, calibr…
9b4f9fe 
…ate, simulate, train

waveorder.models is the user API: one model per contrast with the same methods, built from a config or in Python.
Reconstruction runs over positions and time points on several GPUs, in tiles when a volume is large.

BREAKING CHANGE: the 3.x Python API is removed. waveorder.focus.focus_from_transverse_band stays and warns.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
feat(cli): wo commands on the new models
02fc356 
wo rec, compute-tf, apply-inv-tf, calibrate-tf, sim, config, distill and bm call waveorder.models, with the flags
of wo-4.0.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
test(bench): regression suite on the new models, training examples
e5f2cc2 
The 19 regression cases run through the new models and compare against wo-4.0. Adds example training configs and
the two-node resume test.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
docs: README and CONTRIBUTING for the new API
fa19eaf 
Examples in both files are run by the test suite.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
fix(cli): keep the --init and --unique-id flags of wo-4.0
344d2c2 
wo calibrate-tf --init and wo rec --unique-id exist in wo-4.0 and were missing here, so commands written for
wo-4.0 failed with an unknown option. --init is a second name of --start. --unique-id is accepted and ignored.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
fix(config): accept the calibration keys wo-4.0 writes
b86b33d 
Files written by wo-4.0 carry convergence_patience, use_gradients and log_dir under solver.calibration.advanced,
and stopping_tolerance for Richardson-Lucy. Loading refused them. The first three now load with a warning that
they have no effect here, and stopping_tolerance is another name for tolerance.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
docs(models): trim docstrings to a summary and parameter sections
e5541d9 
Public methods of the models now document every parameter in one or two lines. History and design notes are gone.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
fix(nn): stop gradient-consensus Richardson-Lucy once every voxel is …
6371928 
…frozen

wo-4.0 ends the iteration when no voxel has a nonzero step. Here the loop went on drawing new photon halves,
which unfroze voxels and moved the estimate away from wo-4.0's. A solver can now report that its state is a
fixed point, and Richardson-Lucy does so at a full freeze.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
fix(nn): store the Hermitian part of cropped transfer functions
b87a13b 
A transfer function cropped from a finer grid is not Hermitian on the Nyquist frequencies of an even axis when
the illumination is tilted, the pupil is aberrated, or z is oversampled. wo-4.0 applies the real part of the
filtered result, which uses the Hermitian part there. Storing that part makes the forward model the same real
operator as wo-4.0's on every grid.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
fix(nn): zero-pad when the z padding reaches the stack depth
5cebaa3 
wo-4.0 mirrors edge slices when the padding is smaller than the stack and pads with zeros otherwise. Here the
second case raised, so a wo-4.0 config with a deep padding on a shallow stack did not run.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
chore(bench): drop the Intel node constraint from the regression job
4581a8b 
The job no longer pins a CPU vendor, so it starts on any free GPU node.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@talonchandler

talonchandler commented Nov 14, 2026

Copy link
Copy Markdown
Collaborator

Exciting stuff! I just took a very brief first pass through the filetrees:

  • the core text of the main README can be left alone...feel free to update the configs and code to fit the new description
  • I'll suggest leaving out the distillation and training code from this PR for now
  • I'm seeing some committed .zarr files in /io/data. Intended?
  • napari widgets entirely deleted? or am I missing the move? The pieces can be lower-priority for testing, but I'm not sure if we'll delete them completely.

@srivarra

srivarra commented Nov 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

@talonchandler

the core text of the main README can be left alone...feel free to update the configs and code to fit the new description

Yup will do

I'll suggest leaving out the distillation and training code from this PR for now

👍 makes sense to me, I'll put it on it's own branch and in waveorder-design.

I'm seeing some committed .zarr files in /io/data. Intended?

nope, will remove them.

napari widgets entirely deleted? or am I missing the move? The pieces can be lower-priority for testing, but I'm not sure if we'll delete them completely.

They'll be added back in a future commit.

@srivarra
fix(models): write a new output in the OME-Zarr version of its input
321760c 
wo-4.0 keeps the version of the input. Here every new output was written as version 0.5, so a pipeline reading
0.4 stores downstream broke. The version iohub reads from the first input now goes to iohub when the output is
created, for plates and for single-position stores.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
fix(config): load every version 4 file that wo-4.0 accepts
0363c3d 
Compared key by key with wo-4.0's schema, this branch refused values that wo-4.0 accepts, among them a file
with both an optimizer and search stages, which wo-4.0's own migration writes. Each now loads with wo-4.0's
meaning, or with a warning at the key where the value has no effect here.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
@srivarra
test: catch layer violations by attribute and strings in type aliases
8803a67 
The layering test missed a lower layer reaching waveorder.models through an attribute after import waveorder,
and the annotation test did not look inside type statements.

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@talonchandler talonchandler Awaiting requested review from talonchandler talonchandler will be requested when the pull request is marked ready for review talonchandler is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@srivarra @talonchandler