High performance Python for direct numerical simulations of turbulent flows (1602.03638v1)

Abstract: Direct Numerical Simulations (DNS) of the Navier Stokes equations is an invaluable research tool in fluid dynamics. Still, there are few publicly available research codes and, due to the heavy number crunching implied, available codes are usually written in low-level languages such as C/C++ or Fortran. In this paper we describe a pure scientific Python pseudo-spectral DNS code that nearly matches the performance of C++ for thousands of processors and billions of unknowns. We also describe a version optimized through Cython, that is found to match the speed of C++. The solvers are written from scratch in Python, both the mesh, the MPI domain decomposition, and the temporal integrators. The solvers have been verified and benchmarked on the Shaheen supercomputer at the KAUST supercomputing laboratory, and we are able to show very good scaling up to several thousand cores. A very important part of the implementation is the mesh decomposition (we implement both slab and pencil decompositions) and 3D parallel Fast Fourier Transforms (FFT). The mesh decomposition and FFT routines have been implemented in Python using serial FFT routines (either NumPy, pyFFTW or any other serial FFT module), NumPy array manipulations and with MPI communications handled by MPI for Python (mpi4py). We show how we are able to execute a 3D parallel FFT in Python for a slab mesh decomposition using 4 lines of compact Python code, for which the parallel performance on Shaheen is found to be slightly better than similar routines provided through the FFTW library. For a pencil mesh decomposition 7 lines of code is required to execute a transform.

• The paper demonstrates that a fully Python-based DNS code, optimized with libraries like NumPy, pyFFTW, and MPI, can achieve performance comparable to traditional C++ implementations.
• It employs a pseudo-spectral method to solve the Navier-Stokes equations accurately on massively parallel systems, validated by a Taylor Green vortex benchmark.
• The results imply that optimizing Python with tools such as Cython, Numba, and Numexpr can lower barriers in CFD research and foster integration with AI techniques.

High Performance Python for DNS of Turbulent Flows

The paper presents a high-performance scientific Python implementation for Direct Numerical Simulations (DNS) of turbulent flows using a pseudo-spectral method. Traditionally, DNS codes are developed in low-level languages like C++ or Fortran to optimize computational efficiency, particularly given the demanding nature of simulating the Navier-Stokes equations on massively parallel computing platforms. However, the authors demonstrate that Python, a high-level language known for ease of use, can achieve performance comparable to those implementations through intelligent structuring of the code and use of appropriate libraries.

Key Features and Results

1. Python Implementation:
   • The DNS code is implemented entirely in Python—covering mesh creation, domain decomposition, and numerical integration routines with pseudo-spectral methods. The implementation includes both slab and pencil decompositions for parallel processing.
   • The paper employs NumPy for array manipulations and MPI for Python for parallel processing. For FFT operations, pyFFTW interfaces are used to optimize performance.
2. Pseudo-Spectral DNS:
   • This approach solves the Navier-Stokes equations in spectral space using Fourier transform methods, which are highly accurate for resolving turbulence. The primary computational load involves executing 3D FFTs, a task efficiently managed via MPI communications.
3. Performance:
   • The paper features comprehensive benchmarking executed on the Shaheen supercomputer at KAUST. The Python code shows promising scaling up to several thousand cores, with a demonstration that the optimized Python version matches the speed of a comparable C++ DNS code.
   • The authors used a Taylor Green vortex as a test case, which serves as a standard benchmark for validating DNS implementations numerically. The results show excellent agreement with reference data.
4. Optimization Techniques:
   • Performance is enhanced using Cython and other Python optimization tools (e.g., Numexpr, Numba) to minimize the overhead associated with Python's high-level abstractions. By translating critical parts of the code into low-level operations, they demonstrate significant performance gains.
5. Implications for Computational Fluid Dynamics (CFD):
   • The paper suggests that Python, when properly optimized, offers a viable alternative for developing DNS codes while maintaining user-friendly interfaces and rapid prototyping capabilities.
   • The approach could attract a broader audience in the CFD community by lowering the barriers of entry to DNS research, without sacrificing performance.

Future Developments

The successful demonstration of Python for DNS has potential future implications for AI-driven model development, as effective high-level programming interfaces could accelerate the integration of AI with existing simulation methodologies. Continued optimization of Python's computational capabilities will reinforce its role in high-performance computing, facilitating advancements in CFD research methodologies and perhaps encouraging cross-pollination between simulation and data-driven techniques enabled by AI.

In conclusion, the findings of the paper represent a significant advancement in computational approaches for fluid dynamics simulations, suggesting more research is warranted in leveraging high-level languages for high-performance computing tasks.