S2FFT
is a Python package for computing Fourier transforms on the sphere
and rotation group (Price & McEwen 2024) using
JAX or PyTorch. It leverages autodiff to provide differentiable transforms, which are
also deployable on hardware accelerators (e.g. GPUs and TPUs).
More specifically, S2FFT
provides support for spin spherical harmonic
and Wigner transforms (for both real and complex signals), with support
for adjoint transformations where needed, and comes with different
optimisations (precompute or not) that one may select depending on
available resources and desired angular resolution
S2FFT
leverages new algorithmic structures that can he highly
parallelised and distributed, and so map very well onto the architecture
of hardware accelerators (i.e. GPUs and TPUs). In particular, these
algorithms are based on new Wigner-d recursions that are stable to high
angular resolution
With this recursion to hand, the spherical harmonic coefficients of an
isolatitudinally sampled map may be computed as a two step process. First,
a 1D Fourier transform over longitude, for each latitudinal ring. Second,
a projection onto the real polar-d functions. One may precompute and store
all real polar-d functions for extreme acceleration, however this comes
with an equally extreme memory overhead, which is infeasible at
The structure of the algorithms implemented in S2FFT
can support any
isolatitude sampling scheme. A number of sampling schemes are currently
supported.
The equiangular sampling schemes of McEwen & Wiaux (2012), Driscoll & Healy (1995) and Gauss-Legendre (1986) are supported, which exhibit associated sampling theorems and so harmonic transforms can be computed to machine precision. Note that the McEwen & Wiaux sampling theorem reduces the Nyquist rate on the sphere by a factor of two compared to the Driscoll & Healy approach, halving the number of spherical samples required.
The popular HEALPix sampling scheme (Gorski et al. 2005) is also supported. The HEALPix sampling does not exhibit a sampling theorem and so the corresponding harmonic transforms do not achieve machine precision but exhibit some error. However, the HEALPix sampling provides pixels of equal areas, which has many practical advantages.
Note
For algorithmic reasons JIT compilation of HEALPix transforms can become slow at high bandlimits, due to XLA unfolding of loops which currently cannot be avoided. After compiling HEALPix transforms should execute with the efficiency outlined in the associated paper, therefore this additional time overhead need only be incurred once. We are aware of this issue and are working to fix it. A fix for CPU execution has now been implemented (see example notebook).
The latest release of S2FFT
published on PyPI can be installed by running
pip install s2fft
This will install S2FFT
's dependencies including JAX if not already installed.
As by default installing JAX from PyPI will use a CPU-only build,
if you wish to install JAX with GPU or TPU support,
you should first follow the relevant installation instructions in JAX's documentation
and then install S2FFT
as above.
Alternatively, the latest development version of S2FFT
may be installed directly from GitHub by running
pip install git+https://github.com/astro-informatics/s2fft
A pytest
test suite for the package is included in the tests
directory.
To install the test dependencies, clone the repository and install the package (in editable mode)
with the extra test dependencies by running from the root of the repository
pip install -e ".[tests]"
To run the tests, run from the root of the repository
pytest
Documentation for the released version is available here. To install the documentation dependencies, clone the repository and install the package (in editable mode) with the extra documentation dependencies by running from the root of the repository
pip install -e ".[docs]"
To build the documentation, run from the root of the repository
cd docs
make html
open _build/html/index.html
A series of tutorial notebooks are included in the notebooks
directory
and rendered in the documentation.
To install the dependencies required to run the notebooks locally, clone the repository and install the package (in editable mode) with the extra documentation and plotting dependencies by running from the root of the repository
pip install -e ".[docs,plotting]"
To run the notebooks in Jupyter Lab, run from the root of the repository
jupyter lab
To import and use S2FFT
is as simple follows:
For a signal on the sphere
import s2fft
# Define sampled signal to transform and harmonic bandlimit
f = ...
L = ...
# Compute harmonic coefficients
flm = s2fft.forward(f, L, method="jax")
# Map back to pixel-space signal
f = s2fft.inverse(flm, L, method="jax")
For a signal on the rotation group
import s2fft
# Define sampled signal to transform and harmonic and azimuthal bandlimits
f = ...
L = ...
N = ...
# Compute Wigner coefficients
flmn = s2fft.wigner.forward(f, L, N, method="jax")
# Map back to pixel-space signal
f = fft.wigner.inverse_jax(flmn, L, N, method="jax")
For further details on usage see the documentation and associated notebooks.
Note
We also provide PyTorch support for the precompute version of our transforms, as demonstrated in the Torch frontend tutorial notebook.
S2FFT
also provides JAX support for existing C/C++ packages, specifically HEALPix
and SSHT
. This works
by wrapping Python bindings with custom JAX frontends. Note that this C/C++ to JAX interoperability is currently limited to CPU.
For example, one may call these alternate backends for the spherical harmonic transform by:
# Forward SSHT spherical harmonic transform
flm = s2fft.forward(f, L, sampling="mw", method="jax_ssht")
# Forward HEALPix spherical harmonic transform
flm = s2fft.forward(f, L, nside=nside, sampling="healpix", method="jax_healpy")
All of these JAX frontends supports out of the box reverse mode automatic differentiation,
and under the hood is simply linking to the C/C++ packages you are familiar with. In this
way S2fft
enhances existing packages with gradient functionality for modern scientific computing or machine learning
applications!
For further details on usage see the associated notebooks.
A suite of benchmark functions for both the on-the-fly and precompute versions of the spherical harmonic and Wigner transforms are available in the benchmarks
directory, along with utilities for running the benchmarks and plotting the results.
Thanks goes to these wonderful people (emoji key):
We encourage contributions from any interested developers. A simple first addition could be adding support for more spherical sampling patterns!
Should this code be used in any way, we kindly request that the following article is referenced. A BibTeX entry for this reference may look like:
@article{price:s2fft,
author = "Matthew A. Price and Jason D. McEwen",
title = "Differentiable and accelerated spherical harmonic and Wigner transforms",
journal = "Journal of Computational Physics",
year = "2024",
volume = "510",
pages = "113109",
eprint = "arXiv:2311.14670",
doi = "10.1016/j.jcp.2024.113109"
}
You might also like to consider citing our related papers on which this code builds:
@article{mcewen:fssht,
author = "Jason D. McEwen and Yves Wiaux",
title = "A novel sampling theorem on the sphere",
journal = "IEEE Trans. Sig. Proc.",
year = "2011",
volume = "59",
number = "12",
pages = "5876--5887",
eprint = "arXiv:1110.6298",
doi = "10.1109/TSP.2011.2166394"
}
@article{mcewen:so3,
author = "Jason D. McEwen and Martin B{\"u}ttner and Boris ~Leistedt and Hiranya V. Peiris and Yves Wiaux",
title = "A novel sampling theorem on the rotation group",
journal = "IEEE Sig. Proc. Let.",
year = "2015",
volume = "22",
number = "12",
pages = "2425--2429",
eprint = "arXiv:1508.03101",
doi = "10.1109/LSP.2015.2490676"
}
We provide this code under an MIT open-source licence with the hope that it will be of use to a wider community.
Copyright 2023 Matthew Price, Jason McEwen and contributors.
S2FFT
is free software made available under the MIT License. For
details see the LICENCE.txt
file.
The file lib/include/kernel_helpers.h
is adapted from
code in a tutorial on extending JAX by
Dan Foreman-Mackey and licensed under a MIT license.
The file lib/include/kernel_nanobind_helpers.h
is adapted from code
by the JAX authors
and licensed under a Apache-2.0 license.