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

TL;DR

WaterLily.jl is a versatile, differentiable Julia-based CFD solver that efficiently handles complex geometries and motions, supporting multi-platform execution and enabling advanced integration with machine learning workflows.

Contribution

The paper introduces WaterLily.jl, a novel, open-source, differentiable CFD solver in Julia supporting multi-platform backends and complex geometries, with significant performance improvements.

Findings

Supports GPU acceleration with up to 100x speed-up

Fully differentiable using automatic differentiation

Comparable performance to low-level CFD codes

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.