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

This paper introduces AegirJAX, a fully differentiable hydrodynamic solver built on reverse-mode automatic differentiation that unifies forward simulation and inverse optimization to enable end-to-end training for diverse coastal engineering tasks, including bathymetry estimation, breakwater design, and active wave control.

The Gist

Imagine you are trying to predict how a giant wave will crash against a coastline, or perhaps you are an engineer trying to design the perfect breakwall to protect a harbor. Traditionally, scientists have used complex computer programs (solvers) to simulate these events. These programs are like highly precise, rigid robots. They follow strict mathematical rules to calculate how water moves.

The problem? These robots are terrible at answering "What if?" questions.

• "What if the ocean floor is shaped differently?"
• "What if we move this breakwall 5 meters to the left?"
• "What if we could generate a counter-wave to cancel out a tsunami?"

To answer these questions, the old way required scientists to manually rewrite the robot's instruction manual (mathematically deriving "adjoint equations") every time they wanted to change a variable. It was like trying to fix a car engine by taking it apart, rebuilding it from scratch, and hoping you didn't break anything. It was slow, expensive, and prone to human error.

The New Solution: AegirJAX (The "Self-Healing" Simulator)

The authors of this paper, Elsa Cardoso-Bihlo and Alex Bihlo, have built a new tool called AegirJAX. Think of AegirJAX not as a rigid robot, but as a flexible, self-aware video game engine.

Here is the core magic: AegirJAX is "differentiable."

In simple terms, this means the simulator can look at its own calculations and say, "If I had changed this one tiny thing at the start, the result at the end would have been different by exactly this much." It keeps a perfect mental map of every single step it took, allowing it to work backward instantly.

How It Works: The "Solver-in-the-Loop"

Instead of treating the physics simulation and the learning process as two separate things, AegirJAX combines them into one continuous loop. Imagine a chef tasting a soup while cooking it.

• Old Way: The chef cooks the soup, serves it, gets a critique, then goes back to the kitchen and tries to guess what ingredient to change for the next batch.
• AegirJAX Way: The chef has a magical tongue that tells them exactly how much salt to add or remove while the soup is still in the pot to make it perfect.

What Can AegirJAX Do? (The Four Superpowers)

The paper demonstrates four amazing things this new system can do:

1. The "Smart Tutor" (Fixing Broken Physics)

Sometimes, the math used to simulate waves is too simple (like a low-resolution video). It misses tiny, fast ripples.

• The Analogy: Imagine a student trying to solve a math problem but they keep making the same small mistake because they don't understand a specific concept.
• The Fix: AegirJAX attaches a small "neural network" (a mini-AI brain) to the simulator. As the simulation runs, the AI watches the errors. If the simulation misses a ripple, the AI learns to inject a tiny "nudge" of force to fix it. It learns the missing physics on the fly, turning a blurry video into HD.

2. The "Architect" (Designing Perfect Breakwaters)

Engineers usually design breakwalls by trial and error or using simple rules.

• The Analogy: Imagine trying to find the best shape for a dam to stop water, but you can only move one brick at a time and wait a week to see if it works.
• The Fix: With AegirJAX, the computer can "dream" up thousands of shapes instantly. It treats the shape of the breakwall as a variable it can tweak. It simulates the wave, sees how much energy gets through, and instantly calculates exactly how to reshape the wall to block the most energy. It can even design a wall that looks like a smooth, curved hill rather than a blocky wall, because it found the mathematically perfect shape.

3. The "Conductor" (Active Wave Cancellation)

Can we stop a tsunami before it hits the shore?

• The Analogy: Imagine a noise-canceling headphone for the ocean. If a loud sound comes from the left, the headphones generate a sound from the right to cancel it out.
• The Fix: AegirJAX trains a "conductor" (a neural network) to watch incoming waves and instantly command a wave-maker to create a "destructive" wave that cancels the incoming one. Because the simulator is differentiable, the conductor learns exactly when and how hard to push the wave-maker to create the perfect cancellation, reducing the harbor's energy by over 97%.

4. The "Detective" (Solving the Mystery of the Ocean Floor)

Often, we don't know what the ocean floor looks like, or we don't know where a landslide started underwater. We only see the waves hitting the shore.

• The Analogy: You hear a sound in a dark room and try to guess where the person is standing.
• The Fix: AegirJAX works backward. It takes the wave data from the shore and asks, "What shape of the ocean floor would cause these exact waves?" It adjusts its guess of the ocean floor until the simulation matches the real data perfectly. It can even figure out how fast a submarine landslide was moving, just by looking at the waves it created.

Why This Matters

This paper is a game-changer because it blurs the line between simulating nature and optimizing it.

• Before: Simulation and Optimization were separate, slow, and difficult.
• Now: They are one seamless process.

The authors admit that the AI needs a lot of data to learn general rules (it's currently very good at specific scenarios but needs more practice to be a universal expert). However, by combining the rigor of physics with the flexibility of machine learning, AegirJAX offers a powerful new way to protect our coasts, design better structures, and understand the ocean without needing to spend years deriving complex math equations by hand.

In short: They built a simulator that doesn't just watch the movie; it can edit the script, redesign the set, and rewrite the ending, all in real-time.

Technical Summary

1. Problem Statement

Coastal engineering and tsunami hazard assessment rely heavily on numerical solvers for wave propagation and run-up (typically based on shallow-water or Boussinesq equations). While forward simulations are mature, applying these models to inverse problems (e.g., bathymetry estimation, source inversion, structural optimization) and active control remains a significant bottleneck.

• Limitations of Traditional Methods: Traditional approaches use the adjoint-state method to compute gradients. Deriving discrete adjoints for complex, non-linear, non-hydrostatic solvers is labor-intensive, error-prone, and rigid. Any change in physics or boundary conditions requires re-deriving the adjoint equations.
• Differentiability Issues: Standard numerical treatments for discontinuities (e.g., boolean wet–dry flags for moving shorelines) and iterative solvers (e.g., ILU preconditioning) break the flow of gradients, making gradient-based optimization mathematically ill-posed near coastlines.
• Limitations of Pure ML: Purely data-driven models (PINNs, neural operators) often struggle with multi-scale coastal phenomena (dispersion, breaking waves) due to spectral bias and the difficulty of balancing loss terms.

2. Methodology: AegirJAX

The authors introduce AegirJAX, a fully differentiable, depth-integrated, non-hydrostatic hydrodynamic solver built entirely within the JAX framework (reverse-mode automatic differentiation).

Governing Equations

The solver is based on the depth-integrated non-hydrostatic shallow-water equations (Stelling & Zijlema; Yamazaki et al.), decomposing pressure into hydrostatic and non-hydrostatic components. It solves for free-surface elevation ($\zeta$), depth-averaged velocity ($U$), and non-hydrostatic pressure ($q$).

Key Algorithmic Modifications for Differentiability

To ensure end-to-end differentiability, the authors replaced non-differentiable components with continuous approximations:

1. Differentiable Wet–Dry Interface: Instead of discrete boolean flags, a scaled sigmoid function ($M = \sigma(S \cdot (D - h_{min}))$) is used to create a continuous "wet fraction" mask. This allows gradients to flow through the moving shoreline, enabling optimization of inundation metrics.
2. Neural Forcing Integration: A dedicated fractional step injects data-driven momentum corrections ($a_{nn}$) from a neural network. These corrections are interpolated to velocity nodes and modulated by the continuous wet–dry mask to ensure stability.
3. Differentiable Linear Solvers:
   • 1D: Uses an exact, parallelizable tridiagonal solver.
   • 2D: Replaces sequential ILU preconditioning with a fully parallel Jacobi (diagonal) preconditioner for the Bi-CGSTAB algorithm, ensuring compatibility with GPU acceleration and automatic differentiation.
4. Vectorized Boundary Conditions: All spatial gradients and flux limiters use whole-array slicing and padding operations, eliminating conditional control flow loops that break computational graphs.

3. Key Contributions

The paper demonstrates the versatility of AegirJAX across four distinct scientific machine learning domains:

A. Model Discovery and Physical Correction

• Goal: Correct model misspecifications (e.g., missing dispersive physics in hydrostatic models) using neural networks.
• Implementation: A neural network predicts a conservative pseudo-pressure potential or acceleration field based on physical features (Froude number, Laplacian of surface elevation).
• Result: The hybrid solver successfully recovers high-frequency super-harmonics and phase shifts in the Beji–Battjes and Monai Valley benchmarks, reducing Mean Squared Error (MSE) by 85–94% compared to the baseline hydrostatic solver, without needing complex multi-layer physics.

B. Inverse Design and Topological Optimization

• Goal: Optimize coastal structures (breakwaters) directly against transient wave forcing.
• Implementation:
  • Rigid Body: Uses a signed distance field with a sigmoid function to make gate position/rotation differentiable.
  • Topology: Uses a latent logit field transformed via sigmoid and Gaussian smoothing (inspired by SIMP) to optimize material distribution. A cubic penalty enforces binary material states (wall vs. no wall).
• Result: The optimizer autonomously discovers optimal gate placements and curved offshore berms that minimize harbor inundation, treating structural parameters as trainable variables.

C. Active Control (Policy-in-the-Loop)

• Goal: Train a neural policy for active wave cancellation (destructive interference).
• Implementation: A Gated Recurrent Unit (GRU) policy reads sensor data and commands an actuator. The policy is trained via back-propagation through the unrolled fluid simulation using exact analytical gradients.
• Result: The policy learns to anticipate wave arrival and phase, reducing harbor energy by >97% in zero-shot tests on unseen wave amplitudes, outperforming sample-inefficient model-free reinforcement learning.

D. Source Inversion and Parameter Estimation

• Goal: Infer hidden physical fields (bathymetry, submarine landslide kinematics) from downstream sensor data.
• Implementation:
  • Bathymetry: Parameterized as a Multi-Layer Perceptron (MLP) mapping coordinates to depth.
  • Landslide: Jointly inverts for the spatial shape (MLP with Fourier features) and kinematic parameters (acceleration profile) of a moving mass.
  • Tracer Source: Localizes a chemical spill source by back-propagating gradients through coupled advection-diffusion equations.
• Result: Successfully recovers complex submerged topography and dynamic landslide parameters, demonstrating the ability to solve ill-posed inverse problems without manual adjoint derivation.

4. Results and Performance

• Accuracy: In the Conical Island benchmark, the neural-augmented solver achieved 74.4% MSE reduction on an unseen intermediate wave case, proving it learns generalized sub-grid physics rather than just memorizing trajectories.
• Efficiency: The framework solves high-dimensional optimization problems (e.g., topology optimization, source inversion) at a fraction of the cost of traditional adjoint methods by avoiding manual derivation.
• Robustness: The differentiable wet–dry interface allows for stable optimization near the shoreline, a region where traditional methods fail due to gradient discontinuities.

5. Significance and Limitations

Significance:

• Unified Framework: AegirJAX blurs the line between forward simulation and inverse optimization, offering a single, end-to-end differentiable paradigm for coastal engineering.
• Scientific ML Advancement: It demonstrates that "solver-in-the-loop" approaches can retain rigorous physical conservation laws while leveraging the flexibility of neural networks for correction and control.
• New Capabilities: Enables tasks previously considered intractable, such as continuous topology optimization for breakwaters and active control of wave energy converters using exact gradients.

Limitations:

• Data Scarcity: The neural corrections learned in the Beji–Battjes and Monai Valley experiments were regime-specific. Because training data was sparse (single events), the networks acted as curve-fitters rather than learning universal physical laws.
• Generalization: The framework requires diverse, densely sampled datasets to learn truly generalizable sub-grid parameterizations.
• Ill-Posed Regions: In inverse problems, areas without downstream sensor coverage (shadow zones) remain mathematically unconstrained, reflecting the physical reality of the problem rather than a solver failure.

Conclusion:
The authors argue that AegirJAX represents a promising step toward next-generation coastal hazard assessment. Future work should focus on generating large-scale, diverse experimental datasets to train universally generalizable neural corrections and potentially pre-training on high-fidelity 3D Navier-Stokes simulations.