JAX-WAVE turns inverse wave design into one forward pass: 45% lower error than FISTA inside 5 ms

Inverse wave design moves from repeated solving to a single network forward pass: given a target terminal field and a known wave-speed map, the network outputs an initial pressure that respects the pressure bounds, without rerunning numerical inversion for each target and medium. The prior route alternated forward propagation with gradient-based control updates for each target-medium pair, using adjoint methods for gradients, until a preset accuracy or computation limit was reached; every design started from scratch. On an independent 2,048-case confirmation set, the width-64 model reaches 25.24% error on unseen geometric targets and 4.77% on unseen reachable targets, close to its errors on known materials.

Training needs no pre-optimized control labels because the physics supplies the supervision: the predicted initial pressure is fed through a differentiable finite-difference time-domain propagator, and the relative terminal-field mismatch is returned as a gradient. Conventional supervised learning would first solve the constrained inverse problem for every training example to obtain a label, at a cost that grows with the number of examples. Training uses 65,536 target and speed-map pairs with a relative terminal-field loss, differentiating through the network and all propagation steps; validation runs every 1,000 updates and selects the checkpoint by validation error.

Pressure-sign symmetry is built into the network: linear propagation, a quadratic control penalty and symmetric amplitude bounds imply that reversing the target reverses the optimal control, and a shared-weight two-sign projection makes every prediction satisfy that relation. Random sign augmentation exposes the network to the same physical relation but leaves it free to violate it, and a plain U-Net carries no such constraint. In a paired comparison matching capacity, initialization and minibatch sequences, the sign projection lowers unseen geometric error by 9.0% and reachable error by 28.6%, improving all three seeds with no change in trainable parameter count.

Where a design is needed quickly, the trained map carries wave physics learned offline straight to a new query, at the cost of two network evaluations per query. Numerical inversion completes only three or four FISTA updates inside the same 5 ms median budget. Across 96 heterogeneous cases, JAX-WAVE reaches 18.79% mean field error within the 5 ms median budget against 34.18% for FISTA, a paired 95% interval of [40.4%,49.4%]; at a 10 ms budget FISTA catches up to 18.71%.

The demonstrated scope is a discrete wave model on the unit square with a fixed terminal time, zero boundary values and zero initial derivative, over layered and elliptical-inclusion media with speeds inside the training range; the online benchmark uses smooth geometric targets and low- to mid-frequency reachable targets at moderate speed contrast. The authors propose replacing distributed pressure control with calibrated actuator responses to connect the formulation to source-plane or transducer-array design, but that step is not implemented here.

Geometric targets retain large errors in both approaches, 14.15% for the numerical reference and 25.37% for the network, so a shape specified independently of propagation is harder to reproduce than a target generated by an admissible control; widening from 64 to 128 lowers geometric error only from 24.22% to 23.42% while costing about 2.5 times as much per optimizer update. Smooth media and six-inclusion media both fall outside the training families, and smooth media produce larger errors, so generalization depends on the spatial organization of the medium and not only its speed range. The training cost, estimated on an eight-case workload, needs roughly 141,000 queries to amortize, and that estimate excludes data preparation, compilation and model selection.