WaterLily.jl: A differentiable and backend-agnostic Julia solver for incompressible viscous flow around dynamic bodies

Abstract: Integrating computational fluid dynamics (CFD) solvers into optimization and machine-learning frameworks is hampered by the rigidity of classic computational languages and the slow performance of more flexible high-level languages. In this work, we introduce WaterLily.jl: an open-source incompressible viscous flow solver written in the Julia language. An immersed boundary method is used to enforce the effect of solid boundaries on flow past complex geometries with arbitrary motions. The small code base is multidimensional, multiplatform and backend-agnostic, ie. it supports serial and multithreaded CPU execution, and GPUs of different vendors. Additionally, the pure-Julia implementation allows the solver to be fully differentiable using automatic differentiation. The computational cost per time step and grid point remains constant with increasing grid size on CPU backends, and we measure up to two orders of magnitude speed-up on a supercomputer GPU compared to serial CPU execution. This leads to comparable performance with low-level CFD solvers written in C and Fortran on research-scale problems, opening up exciting possible future applications on the cutting edge of machine-learning research.

• The paper introduces WaterLily.jl, a differentiable and backend-agnostic CFD solver that integrates automatic differentiation for ML optimization.
• It employs a finite-volume approach with immersed boundary methods to deliver second-order accurate simulations of incompressible viscous flows.
• Benchmark results show up to 200x GPU acceleration and robust validation against classic CFD cases, underscoring its efficiency for complex simulations.

Overview of WaterLily.jl: A Julia-Based Differentiable CFD Solver

The paper introduces WaterLily.jl, a computational fluid dynamics (CFD) solver written in the Julia programming language, designed to simulate incompressible viscous flow and dynamically interacting bodies. This solver distinguishes itself by being both differentiable and backend-agnostic, addressing significant challenges in integrating CFD with optimization and ML frameworks. Traditional methods often face issues with the rigidity of low-level programming languages and the inefficiencies of high-level languages. WaterLily.jl leverages Julia’s dynamic type system, allowing seamless integration with ML frameworks while maintaining computational performance on various hardware architectures, including CPUs and GPUs.

Core Features and Numerical Methods

WaterLily.jl is characterized by its concise, multidimensional codebase that supports both CPU and GPU execution. The solver efficiently enforces the effects of solid boundaries using the immersed boundary method. It is particularly notable for its differentiability across the entire codebase, enabled through automatic differentiation (AD), offering significant advantages in optimization problems and ML integration.

The solver utilizes a finite-volume approach on a uniform Cartesian grid, implementing second-order accurate numerical schemes. The momentum equation governing fluid flow is integrated in time with an explicit predictor-corrector update, ensuring incompressibility through a pressure projection scheme. Handling of immersed boundaries is facilitated by the boundary data immersion method (BDIM), ensuring accuracy and flexibility across different dimensional simulations.

Performance and Validation

Benchmarking results demonstrate WaterLily.jl’s efficiency, showcasing up to 200x speed-up using CUDA-enabled GPUs, achieving comparable performance to that of low-level solvers like those written in C and Fortran. Profiling indicates that the pressure-solving component is the most time-consuming, a common bottleneck in CFD solvers. This efficiency in GPU performance highlights the solver's capacity for handling large-scale simulations and its adaptability to various computational environments.

The validation against classic cases like the Taylor-Green vortex and flow past a sphere reveals its accuracy and robustness. The solver's results closely match direct numerical simulation (DNS) references, confirming its applicability to complex CFD problems at various scales.

Implications and Future Directions

The introduction of WaterLily.jl holds significant implications for research at the intersection of CFD and ML. Its differentiability and backend-agnostic design allow researchers to easily integrate advanced optimization techniques and ML models into CFD simulations. This capability could streamline workflows in optimization tasks and accelerate data-driven methods development.

Looking ahead, potential developments include exploring distributed-memory parallelization for simulations requiring extensive computational resources and expanding its scope to multiphase flow simulations. Enhancements in solver performance and new test case validations would further solidify its applicability across broader CFD roles in both academic and industrial settings.

In summary, WaterLily.jl presents a versatile and efficient tool for the modern CFD landscape, adeptly bridging gaps between high-performance computing and emerging ML techniques. Its innovative construction and continued improvements will likely foster broader applications and stimulate research developments in CFD and associated fields.