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

TL;DR

FFTArray addresses the challenge of general, high-performance discretization of Fourier transforms for pseudo-spectral methods by introducing a general discretized Fourier transform framework (gdFT/gdIFT) that works on arbitrarily placed multi-dimensional grids. It decouples physics solvers from FFT discretization, enabling direct translation from textbook equations to code and seamless dimensional broadcasting. The library implements Dimension and Array classes to manage coordinate grids, spaces (position vs frequency), and lazy phase-factor application, with support for multiple backends via the Python Array API and GPU acceleration. Practical demonstrations focus on ultracold atomic systems, while the framework emphasizes modularity, performance, and adaptability to diverse boundary conditions and coordinate choices. Overall, FFTArray provides a flexible, backend-agnostic foundation for scalable, maintainable pseudo-spectral computations in multi-dimensions.

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.

Figures (1)

• Figure 1: Conceptual overview of FFTArray's role in scientific software architecture. Traditional software couples the definition of the physics problem, the Fourier-based differential equation solver and the discretization of Fourier transforms with Fast Fourier Transforms (FFTs) within a monolithic framework. These components interact with each other via problem-specific internal APIs, creating complex interdependent codebases. FFTArray decouples the implementation of discretized Fourier transforms from the solver and physics problem. This architectural simplification enables researchers to focus on core physics and solver logic without managing low-level FFT implementation details, resulting in more maintainable and reusable scientific code.