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

FFTArray is a modular, open-source Python library built on the Array API Standard that simplifies the implementation of discretized multi-dimensional Fourier transforms for pseudo-spectral methods by automating complex grid and scaling corrections while ensuring compatibility with diverse backends like NumPy, JAX, and PyTorch.

The Gist

Imagine you are a chef trying to bake a perfect, complex cake (a physical system like a quantum particle). The recipe (the laws of physics) is written in a language of continuous, flowing ingredients (calculus and Fourier transforms). But your kitchen only has a digital scale that can only weigh ingredients in tiny, discrete scoops (discrete grids).

The problem? Translating that smooth, flowing recipe into discrete scoops is tricky. If you just start scooping randomly, your cake might collapse, or worse, it might taste like a completely different dessert because you messed up the timing and spacing.

For a long time, scientists had to build their own custom "scooping machines" for every single recipe they wanted to bake. If they wanted to change the oven size or the type of flour, they often had to rebuild the whole machine from scratch. This was slow, prone to errors, and made it hard to share recipes with other chefs.

Enter FFTArray: The Universal, Smart Scooping Machine.

This paper introduces FFTArray, a new Python tool that acts like a universal adapter for these digital scooping machines. Here is how it works, using some everyday analogies:

1. The "Magic Adapter" (The Core Idea)

Think of the Fast Fourier Transform (FFT) as a magical machine that can instantly switch a cake batter between "Position Space" (where the ingredients are physically located in the bowl) and "Frequency Space" (a secret code that tells you how the ingredients vibrate and interact).

Usually, scientists had to manually write code to tell the machine: "Hey, start scooping at -5cm, stop at +5cm, and remember to add a little extra sugar (phase factor) when you switch modes." If they forgot the sugar, the cake was ruined.

FFTArray is like a smart adapter that handles all that math for you. You just tell it, "I want to switch to Frequency Space," and it automatically knows exactly where to start scooping, how much "sugar" (mathematical correction) to add, and ensures the transition is perfect. It takes the messy, error-prone math and hides it in a neat little box.

2. The "Lego" System (Modular Design)

Imagine you are building a model with Legos. In the old days, if you wanted to build a 3D castle, you had to glue the bricks together in a specific, rigid way. If you wanted to change it to a 2D wall, you had to smash the whole thing and start over.

FFTArray treats dimensions (like width, height, and depth) like named Lego bricks.

• You can snap a "Width" brick onto a "Height" brick.
• The system automatically knows how they fit together.
• You can switch from a 1D line to a 3D cube without rewriting your code.
  It's like having a set of instructions that says, "Connect the red brick to the blue brick," regardless of whether you are building a tower or a bridge.

3. The "Lazy Chef" (Performance Optimization)

One of the coolest features is how it saves energy. Imagine you are cooking a meal and you need to add salt, then pepper, then salt again. A normal chef might add salt, stir, add pepper, stir, and then realize, "Oh wait, I added salt twice, I should have just added it once at the end."

FFTArray is a lazy but brilliant chef. It tracks what you need to do but waits to actually do it until the very last second. If it realizes that a step you planned to take (like adding a specific mathematical correction) will be cancelled out by the next step, it simply skips it entirely. This makes the cooking process (the computer calculation) incredibly fast without ruining the taste (the accuracy).

4. The "Universal Remote" (Working with Any Hardware)

In the past, if you wanted to cook on a high-tech stove (a super-fast computer GPU), you needed a specific remote control for that stove. If you wanted to switch to a standard stove (a regular CPU), you needed a different remote.

FFTArray is a universal remote. It works with the standard "remote" (Python's Array API) that controls all the major stoves in the kitchen:

• NumPy: The reliable, standard stove found in almost every kitchen.
• JAX & PyTorch: The high-performance, turbo-charged stoves used for serious cooking (and AI).

You write your recipe once, and FFTArray lets you run it on a standard stove or a turbo-charged GPU without changing a single line of code. This means scientists can use cheap, powerful graphics cards (like the ones in gaming PCs) to solve complex physics problems that used to require million-dollar supercomputers.

Why Does This Matter?

Before FFTArray, a physicist wanting to simulate how atoms interact in a new type of laser had to spend months debugging the "scooping" math. Now, they can focus on the physics (the recipe) rather than the math (the scooping).

• For the Student: It turns a confusing, error-filled math problem into a simple "plug-and-play" coding task.
• For the Researcher: It allows them to simulate massive, complex systems (like clouds of atoms) in seconds rather than days.
• For the Future: It makes it easier to discover new materials, better lasers, and understand the quantum world because the tools are finally easy to use.

In short, FFTArray is the tool that lets scientists stop worrying about the plumbing and start focusing on the water flow. It turns the complex, rigid world of Fourier transforms into a flexible, easy-to-use playground for discovery.

Technical Summary

1. Problem Statement

Partial differential equations (PDEs) governing physical systems, such as the Schrödinger equation, often lack closed-form solutions. Fourier spectral methods, utilizing Fast Fourier Transforms (FFTs), are a standard approach for approximating these solutions. However, translating continuous Fourier integrals into discrete FFTs on arbitrary coordinate grids is non-trivial due to:

• Grid Constraints: The coupling between position and frequency grid spacings (governed by the Sampling Theorem) requires careful selection of grid offsets ($x_{min}, f_{min}$) and scaling factors.
• Phase and Scaling Factors: Correct discretization requires specific coordinate-dependent phase shifts and amplitude scaling to ensure the discrete transform is the exact inverse of the continuous one.
• Software Monoliths: Existing implementations often tightly couple the physics definition, the solver algorithm, and the FFT discretization logic. This "monolithic" design leads to:
  • Lack of Generality: Code is often hard-coded for specific grid symmetries (e.g., $x_{min}=0$) or specific dimensions.
  • Code Duplication: Researchers must re-derive discretization logic for every new problem or coordinate system.
  • Complexity: Modifications to one part of the system risk breaking interconnected components, making maintenance difficult.

2. Methodology

The authors present FFTArray, a Python library designed to decouple the discretization of Fourier transforms from the physics solver. The methodology relies on three core pillars:

A. Mathematical Framework: Generalized Discretized Fourier Transform (gdFT)

FFTArray implements a general discretized Fourier transform that works on arbitrarily shifted grids ($x_{min} \neq 0$, $f_{min} \neq 0$).

• Decomposition: The library decomposes the general transform into a standard FFT (or IFFT) surrounded by pre- and post-multiplication of coordinate-dependent phase and scaling factors.
• Lazy Evaluation: To maintain performance, FFTArray employs a "lazy application" strategy. Instead of immediately applying phase factors during every intermediate step, it tracks a factors_applied flag. The library automatically determines when factors are necessary and when they can be canceled out (e.g., during a forward-backward transform pair), applying them only when strictly required for the final result. This avoids unnecessary floating-point operations and precision loss.

B. Software Architecture

FFTArray is built on the Python Array API Standard, ensuring compatibility with multiple backends (NumPy, JAX, PyTorch) and hardware accelerators (GPUs).

• Dimension Class: Encapsulates the metadata for a single dimension (grid size $N$, spacing $\Delta x$, offsets). It uses a constraint solver (Z3) to automatically calculate valid grid parameters based on user inputs (e.g., "center frequency at 0" or "total extent 10m"), ensuring the Sampling Theorem ($N \Delta x \Delta f = 1$) is always satisfied.
• Array Class: Manages multi-dimensional data. It tracks the current space (position or frequency) for each dimension.
  • Named Broadcasting: Dimensions are broadcast by name (e.g., "x", "y") rather than axis index, enabling seamless transitions from 1D to 3D simulations.
  • Implicit Transforms: Changing the space (e.g., arr.into_space("freq")) triggers the necessary FFT and lazy factor application automatically.
  • Arithmetic: Operations (add, multiply, etc.) are optimized to handle mixed factors_applied states, minimizing state transitions.

C. Performance Optimization

• Backend Agnosticism: By adhering to the Array API, FFTArray can leverage the high-performance FFT implementations of JAX and PyTorch for GPU acceleration.
• JAX Tracing: The library supports JAX's Just-In-Time (JIT) compilation and tracing, allowing entire simulation loops to be compiled into efficient GPU kernels.

3. Key Contributions

1. Generalized Discretization: A robust implementation of the gdFT that handles arbitrary grid offsets and non-symmetric domains, removing the need for users to manually derive phase factors for specific coordinate choices.
2. Modular Design: Decouples the physics problem definition from the numerical discretization, allowing researchers to write code that closely mirrors textbook equations.
3. Lazy Phase Factor Management: A novel mechanism to skip unnecessary phase/scaling operations, ensuring that the overhead of the abstraction layer is negligible compared to raw FFT calls.
4. Multi-Backend Support: Seamless integration with NumPy (CPU), JAX, and PyTorch (GPU), enabling high-performance simulations on consumer and server-grade hardware.
5. Constraint Solver: An automated system to define valid coordinate grids based on physical constraints (e.g., resolution, domain size) rather than manual parameter tuning.

4. Results and Evaluation

The authors validated FFTArray through several applications and performance benchmarks:

• Accuracy:

  • Derivatives: Numerical derivatives computed via FFT matched analytical solutions to 11 decimal places.
  • Quantum Harmonic Oscillator: Simulations of the 2D isotropic quantum harmonic oscillator using imaginary time evolution converged to the analytical ground state energy with a relative error better than $10^{-9}$.
  • Bragg Diffraction: Successfully simulated matter-wave beam splitters in both Raman-Nath and Bragg regimes, matching analytical Bessel function solutions.
  • Two-Species BEC: Successfully solved coupled Gross-Pitaevskii equations for interacting Rubidium-87 and Potassium-41 condensates.

• Performance:

  • Overhead: Benchmarks on a 4096x4096 grid showed that FFTArray introduces negligible overhead compared to "Raw FFT" implementations (directly calling xp.fft on raw arrays). The performance of "FFTArray Precomputed" (pre-calculating propagators) was nearly identical to the raw implementation.
  • Hardware Scaling:
    • GPUs vs. CPUs: GPU implementations (JAX on A100/RTX 4090) were approximately two orders of magnitude faster than CPU implementations (NumPy/JAX on AMD Epyc/Ryzen).
    • Precision: The library supports both float32 and float64. A hybrid approach (using float32 for evolution and float64 for energy evaluation) was shown to recover precision lost in pure float32 simulations with minimal cost.
  • Comparison: FFTArray outperformed the state-of-the-art TorchGPE library in the tested scenarios, partly because TorchGPE uses a larger frequency grid (increasing memory/compute load) without a proportional gain in precision for the specific test case.

5. Significance

FFTArray addresses a critical bottleneck in computational physics: the difficulty of implementing robust, high-performance spectral methods.

• Scientific Productivity: It allows physicists to focus on the core physics and solver logic rather than low-level FFT implementation details, significantly reducing development time and error rates.
• Scalability: By leveraging GPU acceleration and the Python Array API, it makes large-scale 3D simulations (e.g., $>10^9$ samples) computationally feasible on consumer hardware.
• Reusability: The modular, dimension-agnostic design promotes code reuse across different physical systems (from single-particle quantum mechanics to multi-species Bose-Einstein condensates).
• Open Source: Released under the Apache-2.0 license, it provides a foundation for the community to build specialized solvers and extend the library to other programming languages or domains.

In conclusion, FFTArray provides a rigorous, high-performance, and user-friendly framework for discretized Fourier transforms, effectively bridging the gap between theoretical mathematical formulations and efficient, scalable numerical code.