|
1 | | -Differentiable and accelerated spherical transforms |
2 | | -=================================================== |
3 | | - |
4 | | -``S2FFT`` is a Python package for computing Fourier transforms on the sphere and rotation |
5 | | -group `(Price & McEwen 2023) <https://arxiv.org/abs/2311.14670>`_ using JAX and PyTorch. |
6 | | -It leverages autodiff to provide differentiable transforms, which are also |
7 | | -deployable on modern hardware accelerators (e.g. GPUs and TPUs). |
8 | | - |
9 | | -More specifically, ``S2FFT`` provides support for spin spherical harmonic and Wigner |
10 | | -transforms (for both real and complex signals), with support for adjoint transformations |
11 | | -where needed, and comes with different optimisations (precompute or not) that one |
12 | | -may select depending on available resources and desired angular resolution :math:`L`. |
13 | | - |
14 | | -.. important:: |
15 | | - HEALPix long JIT compile time fixed for CPU! Fix for GPU coming soon. |
16 | | - |
17 | | -.. tip:: |
18 | | - As of version 1.0.2 ``S2FFT`` also provides PyTorch implementations of underlying |
19 | | - precompute transforms. In future releases this support will be extended to our |
20 | | - on-the-fly algorithms. |
21 | | - |
22 | | -.. tip:: |
23 | | - As of version 1.1.0 ``S2FFT`` also provides JAX support for existing C/C++ packages, |
24 | | - specifically ``HEALPix`` and ``SSHT``. This works by wrapping python bindings with custom |
25 | | - JAX frontends. Note that currently this C/C++ to JAX interoperability is currently |
26 | | - limited to CPU. |
27 | | - |
28 | | -Algorithms |:zap:| |
29 | | -------------------- |
30 | | - |
31 | | -``S2FFT`` leverages new algorithmic structures that can he highly parallelised and |
32 | | -distributed, and so map very well onto the architecture of hardware accelerators (i.e. |
33 | | -GPUs and TPUs). In particular, these algorithms are based on new Wigner-d recursions |
34 | | -that are stable to high angular resolution :math:`L`. The diagram below illustrates the |
35 | | -recursions (for further details see Price & McEwen 2023). |
36 | | - |
37 | | -.. image:: ./assets/figures/Wigner_recursion_github_docs.png |
38 | | - |
39 | | -With this recursion to hand, the spherical harmonic coefficients of an |
40 | | -isolatitudinally sampled map may be computed as a two step process. First, |
41 | | -a 1D Fourier transform over longitude, for each latitudinal ring. Second, |
42 | | -a projection onto the real polar-d functions. One may precompute and store |
43 | | -all real polar-d functions for extreme acceleration, however this comes |
44 | | -with an equally extreme memory overhead, which is infeasible at L ~ 1024. |
45 | | -Alternatively, the real polar-d functions may calculated recursively, |
46 | | -computing only a portion of the projection at a time, hence incurring |
47 | | -negligible memory overhead at the cost of slightly slower execution. The |
48 | | -diagram below illustrates the separable spherical harmonic transform. |
49 | | - |
50 | | -.. image:: ./assets/figures/sax_schematic_github_docs.png |
51 | | - |
52 | | -.. note:: |
53 | | - 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 <https://astro-informatics.github.io/s2fft/tutorials/spherical_harmonic/JAX_HEALPix_backend.html>`_). Fix for GPU execution is coming soon. |
54 | | - |
55 | | -Sampling |:earth_africa:| |
56 | | ------------------------------------ |
57 | | - |
58 | | -The structure of the algorithms implemented in ``S2FFT`` can support any isolattitude sampling scheme. A number of sampling schemes are currently supported. |
59 | | - |
60 | | -The equiangular sampling schemes of `McEwen & Wiaux (2012) <https://arxiv.org/abs/1110.6298>`_, |
61 | | -`Driscoll & Healy (1995) <https://www.sciencedirect.com/science/article/pii/S0196885884710086>`_, and `Gauss-Legendre (1986) <https://link.springer.com/article/10.1007/BF02519350>`_ 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. |
62 | | - |
63 | | -The popular `HEALPix <https://healpix.jpl.nasa.gov>`_ sampling scheme (`Gorski et al. 2005 <https://arxiv.org/abs/astro-ph/0409513>`_) 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. |
64 | | - |
65 | | -.. image:: ./assets/figures/spherical_sampling.png |
66 | | - :width: 900 |
67 | | - :align: center |
68 | | - |
69 | | -Contributors ✨ |
70 | | ------------------------------------ |
71 | | - |
72 | | -Thanks goes to these wonderful people (`emoji |
73 | | -key <https://allcontributors.org/docs/en/emoji-key>`_): |
74 | | - |
75 | | -.. raw:: html |
76 | | - |
77 | | - <embed> |
78 | | - <table> |
79 | | - <tbody> |
80 | | - <tr> |
81 | | - <td align="center" valign="top" width="14.28%"><a href="https://cosmomatt.github.io"><img src="https://avatars.githubusercontent.com/u/32554533?v=4?s=100" width="100px;" alt="Matt Price"/><br /><sub><b>Matt Price</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=CosmoMatt" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3ACosmoMatt" title="Reviewed Pull Requests">👀</a> <a href="#ideas-CosmoMatt" title="Ideas, Planning, & Feedback">🤔</a></td> |
82 | | - <td align="center" valign="top" width="14.28%"><a href="http://www.jasonmcewen.org"><img src="https://avatars.githubusercontent.com/u/3181701?v=4?s=100" width="100px;" alt="Jason McEwen "/><br /><sub><b>Jason McEwen </b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=jasonmcewen" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3Ajasonmcewen" title="Reviewed Pull Requests">👀</a> <a href="#ideas-jasonmcewen" title="Ideas, Planning, & Feedback">🤔</a></td> |
83 | | - <td align="center" valign="top" width="14.28%"><a href="http://matt-graham.github.io"><img src="https://avatars.githubusercontent.com/u/6746980?v=4?s=100" width="100px;" alt="Matt Graham"/><br /><sub><b>Matt Graham</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=matt-graham" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3Amatt-graham" title="Reviewed Pull Requests">👀</a></td> |
84 | | - <td align="center" valign="top" width="14.28%"><a href="https://sfmig.github.io/"><img src="https://avatars.githubusercontent.com/u/33267254?v=4?s=100" width="100px;" alt="sfmig"/><br /><sub><b>sfmig</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=sfmig" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3Asfmig" title="Reviewed Pull Requests">👀</a></td> |
85 | | - <td align="center" valign="top" width="14.28%"><a href="https://github.com/Devaraj-G"><img src="https://avatars.githubusercontent.com/u/36169767?v=4?s=100" width="100px;" alt="Devaraj Gopinathan"/><br /><sub><b>Devaraj Gopinathan</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=Devaraj-G" title="Code">💻</a></td> |
86 | | - <td align="center" valign="top" width="14.28%"><a href="http://flanusse.net"><img src="https://avatars.githubusercontent.com/u/861591?v=4?s=100" width="100px;" alt="Francois Lanusse"/><br /><sub><b>Francois Lanusse</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=EiffL" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/issues?q=author%3AEiffL" title="Bug reports">🐛</a></td> |
87 | | - <td align="center" valign="top" width="14.28%"><a href="https://github.com/eltociear"><img src="https://avatars.githubusercontent.com/u/22633385?v=4?s=100" width="100px;" alt="Ikko Eltociear Ashimine"/><br /><sub><b>Ikko Eltociear Ashimine</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=eltociear" title="Documentation">📖</a></td> |
88 | | - </tr> |
89 | | - <tr> |
90 | | - <td align="center" valign="top" width="14.28%"><a href="https://github.com/kmulderdas"><img src="https://avatars.githubusercontent.com/u/33317219?v=4?s=100" width="100px;" alt="Kevin Mulder"/><br /><sub><b>Kevin Mulder</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/issues?q=author%3Akmulderdas" title="Bug reports">🐛</a></td> |
91 | | - <td align="center" valign="top" width="14.28%"><a href="https://github.com/PhilippMisofCH"><img src="https://avatars.githubusercontent.com/u/142883157?v=4?s=100" width="100px;" alt="Philipp Misof"/><br /><sub><b>Philipp Misof</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/issues?q=author%3APhilippMisofCH" title="Bug reports">🐛</a></td> |
92 | | - <td align="center" valign="top" width="14.28%"><a href="https://github.com/ElisR"><img src="https://avatars.githubusercontent.com/u/19764906?v=4?s=100" width="100px;" alt="Elis Roberts"/><br /><sub><b>Elis Roberts</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/issues?q=author%3AElisR" title="Bug reports">🐛</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=ElisR" title="Documentation">📖</a></td> |
93 | | - <td align="center" valign="top" width="14.28%"><a href="https://github.com/ASKabalan"><img src="https://avatars.githubusercontent.com/u/83787080?v=4?s=100" width="100px;" alt="Wassim KABALAN"/><br /><sub><b>Wassim KABALAN</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=ASKabalan" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3AASKabalan" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=ASKabalan" title="Tests">⚠️</a></td> |
94 | | - </tr> |
95 | | - </tbody> |
96 | | - </table> |
97 | | - </embed> |
98 | | - |
99 | | -We encourage contributions from any interested developers. A simple |
100 | | -first addition could be adding support for more spherical sampling |
101 | | -patterns! |
102 | | - |
103 | | -Attribution |:books:| |
104 | | ------------------- |
105 | | - |
106 | | -Should this code be used in any way, we kindly request that the following |
107 | | -article is referenced. A BibTeX entry for this reference may look like: |
108 | | - |
109 | | -.. code-block:: |
110 | | -
|
111 | | - @article{price:s2fft, |
112 | | - author = "Matthew A. Price and Jason D. McEwen", |
113 | | - title = "Differentiable and accelerated spherical harmonic and Wigner transforms", |
114 | | - journal = "Journal of Computational Physics", |
115 | | - year = "2024", |
116 | | - volume = "510", |
117 | | - pages = "113109", |
118 | | - eprint = "arXiv:2311.14670", |
119 | | - doi = "10.1016/j.jcp.2024.113109" |
120 | | - } |
121 | | -
|
122 | | -You might also like to consider citing our related papers on which this code builds: |
123 | | - |
124 | | -.. code-block:: |
125 | | -
|
126 | | - @article{mcewen:fssht, |
127 | | - author = "Jason D. McEwen and Yves Wiaux", |
128 | | - title = "A novel sampling theorem on the sphere", |
129 | | - journal = "IEEE Trans. Sig. Proc.", |
130 | | - year = "2011", |
131 | | - volume = "59", |
132 | | - number = "12", |
133 | | - pages = "5876--5887", |
134 | | - eprint = "arXiv:1110.6298", |
135 | | - doi = "10.1109/TSP.2011.2166394" |
136 | | - } |
137 | | -
|
138 | | -.. code-block:: |
139 | | - |
140 | | - @article{mcewen:so3, |
141 | | - author = "Jason D. McEwen and Martin B{\"u}ttner and Boris ~Leistedt and Hiranya V. Peiris and Yves Wiaux", |
142 | | - title = "A novel sampling theorem on the rotation group", |
143 | | - journal = "IEEE Sig. Proc. Let.", |
144 | | - year = "2015", |
145 | | - volume = "22", |
146 | | - number = "12", |
147 | | - pages = "2425--2429", |
148 | | - eprint = "arXiv:1508.03101", |
149 | | - doi = "10.1109/LSP.2015.2490676" |
150 | | - } |
151 | | -
|
152 | | -License |:memo:| |
153 | | ----------------- |
154 | | - |
155 | | -We provide this code under an MIT open-source licence with the hope that it will be of use |
156 | | -to a wider community. |
157 | | - |
158 | | -Copyright 2023 Matthew Price, Jason McEwen and contributors. |
159 | | - |
160 | | -``S2FFT`` is free software made available under the MIT License. For details see |
161 | | -the LICENSE file. |
| 1 | +.. mdinclude:: ../README.md |
162 | 2 |
|
163 | 3 | .. bibliography:: |
164 | 4 | :notcited: |
|
0 commit comments