Automatic differentiation frameworks for photoacoustic tomography

Photoacoustic tomography (PAT) reconstructs the initial pressure distribution from acoustic signals measured on the boundary. We develop differentiable reconstruction methods in JAX in which the forward acoustic model, implemented with the jwave wave solver, is composed with the reconstruction objective in a single automatically differentiated program. Gradients are therefore obtained by reverse-mode automatic differentiation (AD) through jwave and the optimizable representation, without a manually implemented adjoint of the forward operator. Synthetic phantoms are used with additive noise being applied to the simulated measurements.

We compare two parameterizations of the initial pressure inside this framework. The first optimizes a discrete pixelized image directly. The second uses a coordinate-based multilayer perceptron (MLP) that maps spatial coordinates to pressure, representing the field as a continuous function. Both are trained by minimizing a sensor fidelity loss on simulated noisy data together with total variation on the reconstructed pressure. For the MLP, we further compare purely data-driven training and physics-informed training incorporating a collocation-based acoustic wave-equation residual. To balance data and PDE terms, we use paired normalization: a root-mean-square (RMS)–normalized sensor misfit (errors scaled by the RMS amplitude of the measurements) together with a PDE residual scaled by a mesh- and sound-speed–dependent operator norm (geometric scaling).

Optimization uses GPU-accelerated training in JAX with exponential learning-rate decay. Hyperparameters—including learning rate, total variation regularization weight, and, when applicable, the PDE collocation weight—are selected using Optuna through a two-stage procedure consisting of quasi-Monte Carlo exploration with pruning followed by Tree-structured Parzen Estimator refinement within a narrowed search region without pruning. The best-performing trial is subsequently trained for an extended sequence of optimization iterations.

Using synthetic phantom benchmarks with noisy measurement data, we evaluate relative reconstruction error together with qualitative agreement between reconstructed and reference pressure fields for both the pixel-based and coordinate-based approaches. The study also examines how continuous-function parameterization and PDE supervision influence convergence behavior, artifact formation, robustness to noise, and final reconstruction accuracy under identical sensor geometries and noise conditions.

The proposed framework shows how JAX-based AD and the jwave forward model can jointly support discrete image optimization, coordinate-based neural fields, and physics-informed objectives in one computational pipeline for PAT. Automated hyperparameter optimization enables systematic and reproducible comparison between these strategies. Future work will investigate alternative neural parameterizations, and more realistic acquisition settings and noise models.