A differentiable software suite for accelerated simulation of turbulent flows

This paper introduces IncompressibleNavierStokes.jl, an open-source Julia package that enables efficient, differentiable, and high-resolution simulation of incompressible turbulent flows on both CPUs and GPUs, facilitating the training of neural network closure models within large-eddy simulations.

The Gist

Imagine you are trying to simulate a storm, a whirlpool in a bathtub, or the airflow over a car wing. To do this accurately, you need to solve a massive, complex set of math equations called the Navier-Stokes equations. Think of these equations as the "laws of physics" for how fluids (like water or air) move.

For decades, solving these equations has been like trying to solve a giant jigsaw puzzle where the pieces keep changing shape, and you need a supercomputer to do it. If you want to see the tiny details (like individual swirling eddies), you need a puzzle with billions of pieces. This usually requires expensive, specialized supercomputers and takes forever.

This paper introduces a new tool called IncompressibleNavierStokes.jl (let's call it INS.jl for short). It's a free software package written in a programming language called Julia that makes simulating these fluid flows faster, cheaper, and smarter.

Here is the breakdown of what makes it special, using some everyday analogies:

1. The "Universal Translator" (Hardware Agnostic)

Usually, if you write code for a standard computer (CPU), it runs slow on a graphics card (GPU), and vice versa. You often have to write two different versions of the same code.

• The Analogy: Imagine a chef who can cook a meal perfectly whether they are using a gas stove, an electric stove, or a campfire, without changing the recipe.
• How INS.jl does it: The software writes the "recipe" (the math code) once. It then automatically translates that recipe to run efficiently on either a standard computer processor or a powerful graphics card (GPU). This means researchers don't have to be coding wizards to get high performance; they just write the logic, and the software handles the rest.

2. The "Memory Saver" (Running on a Single Card)

Simulating a storm at high resolution usually requires a supercomputer with hundreds of gigabytes of memory. Most people only have a single graphics card in their computer (like an NVIDIA RTX 4090 or H100), which has limited memory.

• The Analogy: Imagine trying to fit a massive library of books into a single backpack. Usually, you'd need a truck. But this software is like a magic compression spell: it folds the books so tightly that they fit in the backpack without losing a single page.
• How it works: The software is incredibly efficient. It reuses memory space instead of constantly asking for new space, and it calculates things "on the fly" rather than storing everything at once. This allows it to run Direct Numerical Simulations (DNS)—the most accurate type of simulation—on a single consumer-grade graphics card, something that was previously thought impossible for such high resolutions.

3. The "Self-Correcting Tutor" (Differentiability & AI)

This is the most exciting part. Usually, if you want to train an Artificial Intelligence (AI) to predict fluid behavior, you have to run the simulation, get the result, stop the simulation, and then run a separate program to teach the AI. It's like driving a car, stopping, getting out, and then trying to teach yourself how to drive better based on a memory of the trip.

• The Analogy: Imagine driving a car that has a "rewind" button. You drive, make a mistake, hit rewind, and the car instantly tells you exactly which pedal you pressed too hard and how to fix it for next time.
• How it works: INS.jl is differentiable. This means the software can run the simulation and, in the exact same moment, calculate how to improve the AI model based on the result. It allows researchers to embed a neural network (an AI) directly inside the fluid simulation and train it in real-time. This is a huge leap forward for creating "smart" simulations that learn from data.

4. The "Lego Builder" (Easy to Extend)

Scientific software is often like a black box: you put data in, and you get results out, but you can't easily change the inside.

• The Analogy: Think of INS.jl as a set of high-quality Lego bricks. If you want to add a new feature—like simulating heat (for weather), adding random wind gusts, or plugging in a new AI model—you just snap a new brick onto the existing structure. You don't have to rebuild the whole tower.
• How it works: The code is designed so that adding new physics or new AI models is as simple as writing a small function. This makes it easy for scientists to experiment and innovate without getting bogged down in complex coding.

5. The "Quality Control" (Reliability)

In science, if you can't reproduce someone else's results, the discovery isn't trusted.

• The Analogy: Imagine a restaurant that not only serves great food but also publishes its exact recipe, lists every ingredient, and invites customers to come in and taste-test the dish to prove it's authentic.
• How it works: The authors didn't just write the code; they built a full "factory" around it. They use automated tests (like a robot taste-tester) to check every line of code. They archive the exact version of the code used for the paper so anyone can download it and get the exact same results. They even validated their results against real-world data (simulating a turbulent channel flow) and showed it matches reference data perfectly.

The Bottom Line

IncompressibleNavierStokes.jl is a game-changer because it democratizes high-end fluid simulation.

1. It's Fast: It runs on standard hardware (even single GPUs).
2. It's Smart: It lets you train AI models directly inside the simulation.
3. It's Open: It's free, easy to use, and built to be trusted.

It turns the "supercomputer" into something that can fit on a researcher's desk, opening the door for more scientists to explore the chaotic beauty of turbulent flows and use AI to understand them better.

Technical Summary

1. Problem Statement

The simulation of turbulent flows using Direct Numerical Simulation (DNS) and Large-Eddy Simulation (LES) faces three primary challenges:

• Computational Cost: High-resolution DNS requires massive memory and compute power, often limiting the use of consumer-grade hardware.
• Differentiability: Integrating Machine Learning (ML) closure models into CFD solvers typically requires complex, error-prone coupling between different programming languages (e.g., Python for ML and C++/Fortran for CFD), hindering end-to-end training.
• Memory Constraints: Storing sparse matrices for discrete operators in traditional solvers consumes excessive memory, preventing high-resolution simulations on single GPUs.

2. Methodology

The authors present IncompressibleNavierStokes.jl (INS.jl), an open-source Julia package designed to address these issues through a unified, differentiable, and hardware-agnostic approach.

Numerical Framework

• Discretization: The solver uses second-order staggered finite volume discretization on Cartesian grids (uniform and non-uniform).
• Equations: It solves the incompressible Navier–Stokes equations. The pressure is stored at cell centers, while velocity components are stored at cell faces (Harlow–Welch scheme), naturally coupling pressure and velocity without artificial stabilization.
• Time Integration: The system is solved using explicit time-stepping (Runge–Kutta methods) coupled with a pressure projection step to enforce the divergence-free constraint. Three pressure solvers are available: Spectral (FFT-based), Direct (LDLT factorization), and Iterative (Conjugate Gradient).
• Boundary Conditions: Supports periodic, Dirichlet (no-slip), symmetric, and outlet conditions via ghost volumes.

Software Architecture & Differentiability

• Matrix-Free Kernels: Instead of storing sparse matrices, the package uses hardware-agnostic kernels compiled via KernelAbstractions.jl. The same source code compiles for multi-threaded CPUs or GPUs (NVIDIA).
• Differentiable Programming: The core innovation is the implementation of hand-written adjoint kernels for all discrete operators. These are registered via ChainRules.jl, enabling reverse-mode automatic differentiation (via Zygote.jl) through the entire solver.
• Memory Optimization: To fit high-resolution DNS on a single GPU (e.g., NVIDIA H100), the package employs:
  1. Array Reuse: Mutating operators overwrite pre-allocated buffers rather than allocating new memory.
  2. Low-Storage RK: Using methods like Wray3 that require minimal storage registers.
  3. Lazy Tensor Computation: Computing convective tensor components on-the-fly rather than assembling full tensors.
• Extensibility: New physics (e.g., temperature, stochastic forcing) or closure models can be added by defining new right-hand side terms without modifying the core solver.

Neural Network Integration

The package supports a-posteriori training of neural network closure models. The neural network (built with Lux.jl) is embedded directly into the LES solver. The entire time-integration loop is unrolled, and gradients are backpropagated through the solver to update the neural network weights, eliminating the need for external coupling frameworks.

3. Key Contributions

1. First Differentiable CFD Solver for Turbulence: Provides a single environment (Julia) for DNS data generation, closure model training, and LES evaluation, fully differentiable via hand-written adjoints.
2. Hardware Agnosticism: A single codebase compiles efficiently for both CPUs and GPUs, removing the need for language interoperability (e.g., Python-C++ bridges).
3. High-Resolution Single-GPU DNS: Achieved double-precision DNS at resolutions up to $840^3$ on a single NVIDIA H100 GPU through aggressive memory optimization.
4. Robust Software Engineering: Implements professional software practices including automated testing, continuous integration (GitHub Actions), literate programming for documentation, and Zenodo archiving for reproducibility.

4. Results

The authors validated the solver through several benchmarks:

• Analytical Validation (2D Taylor-Green Vortex):

  • Confirmed second-order spatial convergence for both uniform and non-uniform grids.
  • Demonstrated that Float64 provides reliable convergence, while Float32 is viable for moderate resolutions but suffers from rounding errors at very high resolutions. Float16 was found unstable for spatial discretization in this context.

• Turbulent Channel Flow (DNS):

  • Simulated a channel flow at friction Reynolds number $Re_\tau = 180$.
  • Results (mean velocity profile, RMS fluctuations, and higher-order moments) showed excellent agreement with reference DNS data from Vreman and Kuerten.
  • Minor discrepancies in higher-order moments were attributed to the second-order spatial scheme compared to the fourth-order reference.

• Large-Eddy Simulation (LES):

  • Tested on a coarser grid ($128 \times 64 \times 64$) using five eddy-viscosity models (No-model, Smagorinsky, WALE, QR, Vreman).
  • Smagorinsky was overly dissipative near walls (non-vanishing viscosity).
  • WALE, QR, and Vreman models correctly decayed to zero at the wall and produced results closer to the DNS reference, validating the implementation of classical closure models.

5. Significance

• Accelerated ML for CFD: By enabling end-to-end differentiability, INS.jl lowers the barrier for training data-driven turbulence models, a critical step toward "AI-enhanced" fluid dynamics.
• Democratization of High-Performance Computing: The ability to run $840^3$ double-precision DNS on a single consumer/datacenter GPU makes high-fidelity turbulence research accessible without requiring supercomputing clusters.
• Reproducibility: The rigorous software engineering practices (version control, automated testing, DOI archiving) set a new standard for scientific code reproducibility in computational fluid dynamics.
• Future-Proofing: The modular design and support for mixed-precision strategies position the software to leverage future GPU architectures that prioritize lower-precision throughput.

In summary, INS.jl represents a significant leap forward in computational fluid dynamics software, bridging the gap between traditional high-fidelity simulation and modern differentiable machine learning workflows.