Data Assimilation and Inverse Problems with Differentiable Solvers in JAX

Value of the Information & Strength of the Argument

The video provides valuable practical insights into applying differentiable programming to scientific computing, specifically for data assimilation. It demonstrates the use of JAX’s automatic differentiation to solve inverse problems, highlighting both the potential and the challenges. The argumentation is solid: the presenter explains the mathematical background (gradient computation as a product of Jacobians), shows empirical results, and connects them to theoretical concepts like Lyapunov exponents. The tutorial is well-structured, with clear code examples and explanations of why certain approaches work or fail.

Scientific Rigor, Source Quality, Title Accuracy

The scientific rigor is high: the presenter is knowledgeable and provides references to relevant libraries (Exponax, Chaotax, Optax, Equinox) and a GitHub repository with the code. The title accurately reflects the content. The video is a tutorial, so it does not present original research but rather demonstrates established techniques. The sources cited are appropriate and directly related to the content.

The title accurately reflects the content: the video focuses on data assimilation and inverse problems using differentiable solvers in JAX.

8/10

The video is a hands-on tutorial by an expert (likely the developer of Exponax and Chaotax) demonstrating practical techniques for data assimilation using differentiable solvers in JAX. It provides clear explanations, code walkthroughs, and references to relevant libraries and concepts. The approach is methodical and reproducible, with links to the code repository.

• 00:00Intro: what data assimilation means here
• 01:04Colab setup; installing Exponax, Optax and Chaotax
• 02:38Rebuilding the Kolmogorov flow stepper and warming up
• 03:45Creating a synthetic target by rolling forward
• 04:44The problem: recover the initial condition from a later state
• 05:43Working out the gradient: a product of stepper Jacobians
• 06:13Code walkthrough: unroll steps, noise scale, Adam
• 08:09Loss, optimality error and the update loop
• 09:15First results — and why this problem is hard
• 09:45Where regularization would help
• 10:17Raising the unroll length: 2, 4, 10, 15
• 12:38Where the optimization breaks down
• 14:58Measuring the Lyapunov exponent with Chaotax
• 16:09Lyapunov time ≈ 25 steps, and what that means for BPTT
• 17:15Second task: inferring the diffusivity from a trajectory
• 18:23Building the stepper as a function of the diffusivity
• 19:29Results: trajectory aggregation is far better posed
• 20:02Why long unrolls stay stable here
• 22:46Trading compute for memory with equinox.filter_checkpoint
• 24:01Watching GPU memory with nvitop
• 25:10Things to try, and what's coming next

• Exponax documentation — Library used for solving PDEs with differentiable solvers.
• Chaotax documentation — Library used for chaos analysis, specifically for computing Lyapunov exponents.
• Optax documentation — Optimization library used for Adam optimizer.
• Equinox documentation — Library for neural networks and modules in JAX, used for checkpointing.
• GitHub repository: hybridization-in-jax — Repository containing all material for the workshop.
• Lecture: The Hybridization of Solvers and Deep Learning — Lecture that provides background on the hybridization of solvers and deep learning.
• Part 0: Exponax and JAX for Kolmogorov flow — Previous part of the series introducing Exponax and JAX for Kolmogorov flow.
• Part 2: Neural-hybrid correctors with solver-in-the-loop — Next part of the series on neural-hybrid correctors.

• Exponax documentation — The library used in the video, which provides differentiable PDE solvers.
• Chaotax documentation — The library used for Lyapunov exponent computation.