s2fft 1.1.1

Differentiable and accelerated spherical transforms with JAX

Differentiable and accelerated spherical transforms

S2FFT is a Python package for computing Fourier transforms on the sphere and rotation group using JAX and 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 $L$.

As of version 1.0.2 S2FFT also provides PyTorch implementations of underlying precompute transforms. In future releases this support will be extended to our on-the-fly algorithms.

As of version 1.1.0 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 currently this C/C++ to JAX interoperability is currently limited to CPU, however for many applications this is desirable due to memory constraints.

Documentation

Read the full documentation here.

Attribution

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        = "2023",
    eprint      = "arXiv:2311.14670"
}

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"
}

License

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 LICENSE file.