# A solver-in-the-loop framework for end-to-end differentiable coastal hydrodynamics

arXiv: https://arxiv.org/abs/2604.07129 · HTML: https://mlsanity.com/paper/2604.07129/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 55 references that are in the map.
- First posted on arXiv 8 Apr 2026.

## Builds on (most cited first, 13 of 13)

- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [When and why PINNs fail to train: A neural tangent kernel perspective](https://mlsanity.com/paper/2007.14527/) (2020)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [Neural General Circulation Models for Weather and Climate](https://mlsanity.com/paper/2311.07222/) (2023)
- [Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers](https://mlsanity.com/paper/2007.00016/) (2020)
- [Learning to Control PDEs with Differentiable Physics](https://mlsanity.com/paper/2001.07457/) (2020)
- [Physics-informed neural networks for the shallow-water equations on the sphere](https://mlsanity.com/paper/2104.00615/) (2021)
- [Differentiable Programming for Differential Equations: A Review](https://mlsanity.com/paper/2406.09699/) (2024)
- [Learning Free-Surface Flow with Physics-Informed Neural Networks](https://mlsanity.com/paper/2111.09705/) (2021)
- [Control of flow behavior in complex fluids using automatic differentiation](https://mlsanity.com/paper/2403.06257/) (2024)
- [Learning vertical coordinates via automatic differentiation of a dynamical core](https://mlsanity.com/paper/2512.17877/) (2025)

## Abstract (arXiv)

Numerical simulation of wave propagation and run-up is a cornerstone of coastal engineering and tsunami hazard assessment. However, applying these forward models to inverse problems, such as bathymetry estimation, source inversion, and structural optimization, remains notoriously difficult due to the rigidity and high computational cost of deriving discrete adjoints. In this paper, we introduce AegirJAX, a fully differentiable hydrodynamic solver based on the depth-integrated, non-hydrostatic shallow-water equations. By implementing the solver entirely within a reverse-mode automatic differentiation framework, AegirJAX treats the time-marching physics loop as a continuous computational graph. We demonstrate the framework's versatility across a suite of scientific machine learning tasks: (1) discovering regime-specific neural corrections for model misspecifications in highly dispersive wave propagation; (2) performing continuous topology optimization for breakwater design; (3) training recurrent neural networks in-the-loop for active wave cancellation; and (4) inverting hidden bathymetry and submarine landslide kinematics directly from downstream sensor data. The proposed differentiable paradigm fundamentally blurs the line between forward simulation and inverse optimization, offering a unified, end-to-end framework for coastal hydrodynamics.