Introduction to S2FFT: Efficient Spherical Transforms
Overview
The S2FFT package is a versatile Python library designed for performing Fourier transforms on spherical data and for processing signals on rotational groups. Developed by Price & McEwen in 2023, this project capitalizes on the computational capabilities of JAX and PyTorch to offer transforms that are not only differentiable but also suitable for execution on powerful hardware accelerators like GPUs and TPUs.
Core Capabilities
S2FFT provides tools for computing spin spherical harmonic and Wigner transforms for both real and complex signals. A distinguishing feature of this package is its support for adjoint transformations and the ability to choose from different optimizations depending on resources and desired angular resolution ($L$).
Recent Advances
One of the notable improvements in version 1.0.2 is the introduction of PyTorch implementations for underlying precompute transforms, with future updates promising expanded support for on-the-fly algorithms. With version 1.1.0, JAX support remains robust, even extending to certain C/C++ packages such as HEALPix and SSHT, albeit with current interoperability chiefly on CPUs.
Innovative Algorithms
At the core of S2FFT are novel algorithmic structures that efficiently exploit the architectural strengths of GPUs and TPUs. These new methods, particularly the Wigner-d recursions, offer stability and high performance even at high angular resolutions. This allows for computations that can be precomputed for enhanced speed, though this comes with a tradeoff in memory usage.
Sampling Schemes
The package supports various spherical sampling schemes, including the equiangular sampling methods devised by McEwen & Wiaux, Driscoll & Healy, and Gauss-Legendre. These methods enable harmonic transforms with high precision. Additionally, it accommodates the widely used HEALPix scheme, offering equal-area pixelation despite some precision compromises.
Installation and Usage
Installing S2FFT is straightforward. It comes with dependencies that are easily set up using pip, ensuring a seamless integration into any active Python environment. From straightforward installation to executing unit tests and building documentation locally, the package is designed to deliver both ease of use and comprehensive functionality.
A typical usage example involves computing harmonic coefficients from a signal on the sphere or rotation group and then reverting to pixel-space signals. This is achieved using simple function calls for forward and inverse transformations, whether with JAX or PyTorch.
Collaborative Development
S2FFT is an open project and actively invites contributions from developers. Whether it's enhancing functionality by supporting more sampling patterns or improving compatibility, the project encourages collaborative efforts.
Acknowledgment
If you use S2FFT in your research or projects, citations to relevant papers and the software itself are recommended. These works not only underpin the technical foundations of the package but also represent the innovation and community efforts driving the project forward.
In conclusion, S2FFT is a key tool for those needing efficient, high-resolution spherical analysis, with the added benefits of state-of-the-art algorithmic support and collaborative openness.