FFTArray: A Python Library for the Implementation of Discretized Multi-Dimensional Fourier Transforms

TL;DR

FFTArray is a Python library that simplifies the discretization of multi-dimensional Fourier transforms, making spectral methods more adaptable and easier to implement across various coordinate systems and boundary conditions.

Contribution

It introduces a modular, general-purpose Python library that automates Fourier transform discretization, enhancing flexibility and scalability for spectral methods in scientific computing.

Findings

Supports multiple array backends like NumPy, JAX, PyTorch

Enables straightforward definition of coordinate grids

Optimized for GPU acceleration

Abstract

Partial differential equations describing the dynamics of physical systems rarely have closed-form solutions. Fourier spectral methods, which use Fast Fourier Transforms (FFTs) to approximate solutions, are a common approach to solving these equations. However, mapping Fourier integrals to discrete FFTs is not straightforward, as the selection of the grid as well as the coordinate-dependent phase and scaling factors require special care. Moreover, most software packages that deal with this step integrate it tightly into their full-stack implementations. Such an integrated design sacrifices generality, making it difficult to adapt to new coordinate systems, boundary conditions, or problem-specific requirements. To address these challenges, we present FFTArray, a Python library that automates the general discretization of Fourier transforms. Its purpose is to reduce the barriers to developing high-performance, maintainable code for pseudo-spectral Fourier methods. Its interface enables the direct translation of textbook equations and complex research problems into code, and its modular design scales naturally to multiple dimensions. This makes the definition of valid coordinate grids straightforward, while coordinate grid specific corrections are applied with minimal impact on computational performance. Built on the Python Array API Standard, FFTArray integrates seamlessly with array backends like NumPy, JAX and PyTorch and supports Graphics Processing Unit acceleration. The code is openly available at https://github.com/QSTheory/fftarray under Apache-2.0 license.