JAX-WAVE: A Differentiable Physics-Supervised Framework for Inverse Wave Design in Heterogeneous Media

Honghui Du, Binyao Guo, Yushu Pan

Designing a prescribed wave field in a heterogeneous medium requires an initial pressure distribution that compensates for the medium's propagation response. Repeated numerical inversion solves this task afresh for each target and environment. We introduce JAX-WAVE for amortized inverse wave design: a shared neural network maps desired terminal fields and known wave-speed maps directly to bounded initial pressures. Training through a differentiable finite-difference time-domain solver uses the physical design objective without requiring optimized control labels. Linear propagation and symmetric pressure bounds imply an exact sign-reversal relation for the regularized optimum. A shared-weight projection enforces this relation, preserves feasibility, and satisfies an exact identity for the paired-sign objective. We characterize the discrete propagator and its weighted adjoint to connect learning, numerical inversion, and optimality assessment. On an independent confirmation set, three trained models achieve unseen-material errors of 25.24% for geometric targets and 4.77% for reachable targets. Matched-capacity comparisons show improved reconstruction with the sign parameterization.…