Introduction to Exponax and JAX for Kolmogorov Flow

Value of the Information & Strength of the Argument

The video provides valuable hands-on guidance for using Exponax, a specialized library for differentiable PDE simulations. The argumentation is clear and logical, walking through the setup and execution of a Kolmogorov flow simulation. The presenter explains the rationale behind design choices, such as the stream-function-vorticity form and the use of JIT compilation, and demonstrates the performance benefits. The tutorial is well-structured and encourages experimentation, which enhances its educational value.

Scientific Rigor, Source Quality, Title Accuracy

The video is scientifically rigorous, as it is based on the author’s own library and references the accompanying GitHub repository. The sources cited are the repository and related videos in the series, which are directly relevant. The title accurately reflects the content, and the tutorial is consistent with standard practices in computational fluid dynamics and JAX. The video does not include any external sources beyond the provided links, but the information is reliable due to the author’s expertise.

The title accurately reflects the content, which introduces Exponax and JAX for simulating Kolmogorov flow.

8/10

The video is a hands-on tutorial by the author of the Exponax library, providing clear explanations and reproducible code. The methods are standard in computational fluid dynamics and JAX, and the content aligns with the provided GitHub repository.

• 00:00What this three-part series covers
• 00:47Repo, notebook, and the Colab extension for VS Code
• 02:10Installing Exponax on the remote runtime
• 02:41Conventions: 2D domain, Cartesian grid, (C, N, N) states
• 03:54Periodic boundaries and the redundant boundary point
• 04:20Stream-function-vorticity form: one channel instead of two
• 04:59An initial condition from pixel-wise Gaussian noise
• 05:23The stepper: an Equinox module you call like a network
• 06:23One step, then an autoregressive for loop
• 07:10Why jit matters: 1000 steps in 4 seconds
• 07:38Reading a vorticity field
• 08:03exponax.repeat: pushing the loop into the compiler
• 08:44exponax.rollout and animating the trajectory
• 09:59How fast this actually is
• 10:18Try it: diffusivity and the effective Reynolds number
• 11:13Try it: other initial conditions
• 11:46Coming up next: data assimilation

• GitHub repository: hybridization-in-jax — The repository contains all the material for the practical session, including the notebook used in the video.
• Lecture: The Hybridization of Solvers and Deep Learning — The lecture that provides theoretical background for the series.
• Part 1: Data assimilation with differentiable solvers — The next video in the series, which builds on the concepts introduced here.
• Part 2: Neural-hybrid correctors with solver-in-the-loop — The final video in the series, which extends the tutorial.

• Exponax documentation — Official documentation for Exponax, which would provide additional details on the library's features.