Skip to content

Commit c15a3e8

Browse files
matt-grahammatt-graham
committed
Merge branch 'main' into dg/benchmarks
2 parents 54ded23 + a796dfe commit c15a3e8

File tree

146 files changed

+9655
-3374
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

146 files changed

+9655
-3374
lines changed

‎.all-contributorsrc‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
{
2+
"files": [
3+
"README.md"
4+
],
5+
"imageSize": 100,
6+
"commit": false,
7+
"commitConvention": "angular",
8+
"contributors": [
9+
{
10+
"login": "CosmoMatt",
11+
"name": "Matt Price",
12+
"avatar_url": "https://avatars.githubusercontent.com/u/32554533?v=4",
13+
"profile": "https://cosmomatt.github.io",
14+
"contributions": [
15+
"code",
16+
"review",
17+
"ideas"
18+
]
19+
},
20+
{
21+
"login": "jasonmcewen",
22+
"name": "Jason McEwen ",
23+
"avatar_url": "https://avatars.githubusercontent.com/u/3181701?v=4",
24+
"profile": "http://www.jasonmcewen.org",
25+
"contributions": [
26+
"code",
27+
"review",
28+
"ideas"
29+
]
30+
},
31+
{
32+
"login": "matt-graham",
33+
"name": "Matt Graham",
34+
"avatar_url": "https://avatars.githubusercontent.com/u/6746980?v=4",
35+
"profile": "http://matt-graham.github.io",
36+
"contributions": [
37+
"code",
38+
"review"
39+
]
40+
},
41+
{
42+
"login": "sfmig",
43+
"name": "sfmig",
44+
"avatar_url": "https://avatars.githubusercontent.com/u/33267254?v=4",
45+
"profile": "https://sfmig.github.io/",
46+
"contributions": [
47+
"code",
48+
"review"
49+
]
50+
},
51+
{
52+
"login": "Devaraj-G",
53+
"name": "Devaraj Gopinathan",
54+
"avatar_url": "https://avatars.githubusercontent.com/u/36169767?v=4",
55+
"profile": "https://github.com/Devaraj-G",
56+
"contributions": [
57+
"code"
58+
]
59+
},
60+
{
61+
"login": "EiffL",
62+
"name": "Francois Lanusse",
63+
"avatar_url": "https://avatars.githubusercontent.com/u/861591?v=4",
64+
"profile": "http://flanusse.net",
65+
"contributions": [
66+
"code",
67+
"bug"
68+
]
69+
}
70+
],
71+
"contributorsPerLine": 7,
72+
"skipCi": true,
73+
"repoType": "github",
74+
"repoHost": "https://github.com",
75+
"projectName": "s2fft",
76+
"projectOwner": "astro-informatics"
77+
}

‎.coveragerc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,4 +14,5 @@ show_missing = true
1414
precision = 2
1515
omit =
1616
*test_*
17-
*__init__*
17+
*__init__*
18+
*plotting*

‎.github/workflows/docs.yml‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ jobs:
1313
runs-on: ubuntu-latest
1414
strategy:
1515
matrix:
16-
python-version: [3.8]
16+
python-version: [3.9]
1717

1818
steps:
1919
- name: Checkout Source
@@ -28,11 +28,14 @@ jobs:
2828
run: |
2929
sudo apt install pandoc
3030
python -m pip install --upgrade pip
31-
python -m pip install tox
31+
pip install jaxlib
32+
pip install -r requirements/requirements-core.txt
33+
pip install -r requirements/requirements-docs.txt
34+
pip install .
3235
3336
- name: Build Documentation
3437
run: |
35-
tox -e docs
38+
cd docs && make html
3639
3740
- name: Deploy
3841
if: github.ref == 'refs/heads/main'

‎.github/workflows/tests.yml‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ jobs:
1414
runs-on: ubuntu-latest
1515
strategy:
1616
matrix:
17-
python-version: [3.8]
17+
python-version: [3.9]
1818

1919
steps:
2020
- name: Checkout Source
@@ -28,9 +28,11 @@ jobs:
2828
- name: Install dependencies
2929
run: |
3030
python -m pip install --upgrade pip
31-
python -m pip install tox
31+
pip install -r requirements/requirements-tests.txt
32+
pip install -r requirements/requirements-core.txt
33+
pip install .
3234
3335
- name: Run tests
3436
run: |
35-
tox -e py38
36-
tox -e codecov -- --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702
37+
pytest --cov-report term --cov=s2wav --cov-config=.coveragerc
38+
codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702

‎.pip_readme.rst‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,12 @@ A BibTeX entry for <project-name> is:
4747
License
4848
=======
4949

50-
``S2FFT`` is released under the MIT license (see `LICENSE.txt <https://github.com/astro-informatics/s2fft/blob/main/LICENCE.txt>`_).
50+
``S2FFT`` is released under the MIT license (see
51+
`LICENSE.txt <https://github.com/astro-informatics/s2fft/blob/main/LICENCE.txt>`_).
5152

5253
.. code-block::
5354
54-
S2fft
55+
S2FFT
5556
Copyright (C) 2022 Author names & contributors
5657
57-
This program is released under the MIT license (see `LICENSE.txt <https://github.com/astro-informatics/s2fft/blob/main/LICENCE.txt>`_).
58+
This program is released under the MIT license (see `LICENSE.txt`).

‎README.md‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
[![image](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml)
2+
[![image](https://readthedocs.org/projects/ansicolortags/badge/?version=latest)](https://astro-informatics.github.io/s2fft)
3+
[![image](https://codecov.io/gh/astro-informatics/s2fft/branch/main/graph/badge.svg?token=7QYAFAAWLE)](https://codecov.io/gh/astro-informatics/s2fft)
4+
[![image](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
5+
[![image](http://img.shields.io/badge/arXiv-xxxx.xxxxx-orange.svg?style=flat)](https://arxiv.org/abs/xxxx.xxxxx)<!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
6+
[![All Contributors](https://img.shields.io/badge/all_contributors-6-orange.svg?style=flat-square)](#contributors-) <!-- ALL-CONTRIBUTORS-BADGE:END -->
7+
[![image](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
8+
9+
<img align="left" height="85" width="98" src="./docs/assets/sax_logo.png">
10+
11+
# Differentiable and accelerated spherical transforms with JAX
12+
13+
`S2FFT` is a JAX package for computing Fourier transforms on the sphere
14+
and rotation group. It leverages autodiff to provide differentiable
15+
transforms, which are also deployable on modern hardware accelerators
16+
(e.g. GPUs and TPUs).
17+
18+
More specifically, `S2FFT` provides support for spin spherical harmonic
19+
and Wigner transforms (for both real and complex signals), with support
20+
for adjoint transformations where needed, and comes with different
21+
optimisations (precompute or not) that one may select depending on
22+
available resources and desired angular resolution $L$.
23+
24+
## Algorithms :zap:
25+
26+
`S2FFT` leverages new algorithmic structures that can he highly
27+
parallelised and distributed, and so map very well onto the architecture
28+
of hardware accelerators (i.e. GPUs and TPUs). In particular, these
29+
algorithms are based on new Wigner-d recursions that are stable to high
30+
angular resolution $L$. The diagram below illustrates the recursions
31+
(for further details see Price & McEwen, in prep.).
32+
33+
![image](./docs/assets/figures/schematic.png)
34+
35+
## Sampling :earth_africa:
36+
37+
The structure of the algorithms implemented in `S2FFT` can support any
38+
isolattitude sampling scheme. A number of sampling schemes are currently
39+
supported.
40+
41+
The equiangular sampling schemes of [McEwen & Wiaux
42+
(2012)](https://arxiv.org/abs/1110.6298) and [Driscoll & Healy
43+
(1995)](https://www.sciencedirect.com/science/article/pii/S0196885884710086)
44+
are supported, which exhibit associated sampling theorems and so
45+
harmonic transforms can be computed to machine precision. Note that the
46+
McEwen & Wiaux sampling theorem reduces the Nyquist rate on the sphere
47+
by a factor of two compared to the Driscoll & Healy approach, halving
48+
the number of spherical samples required.
49+
50+
The popular [HEALPix](https://healpix.jpl.nasa.gov) sampling scheme
51+
([Gorski et al. 2005](https://arxiv.org/abs/astro-ph/0409513)) is also
52+
supported. The HEALPix sampling does not exhibit a sampling theorem and
53+
so the corresponding harmonic transforms do not achieve machine
54+
precision but exhibit some error. However, the HEALPix sampling provides
55+
pixels of equal areas, which has many practical advantages.
56+
57+
<p align="center"><img src="./docs/assets/figures/spherical_sampling.png" width="500"></p>
58+
59+
## Installation :computer:
60+
61+
The Python dependencies for the `S2FFT` package are listed in the file
62+
`requirements/requirements-core.txt` and will be automatically installed
63+
into the active python environment by [pip](https://pypi.org) when running
64+
65+
``` bash
66+
pip install .
67+
```
68+
69+
from the root directory of the repository. Unit tests can then be
70+
executed to ensure the installation was successful by running
71+
72+
``` bash
73+
pytest tests/ # for pytest
74+
tox -e py38 # for tox
75+
```
76+
77+
In the very near future one will be able to install `S2FFT` directly
78+
from [PyPi](https://pypi.org) by `pip install s2fft` but this is not yet
79+
supported. Note that to run `JAX` on NVIDIA GPUs you will need to follow
80+
the [guide](https://github.com/google/jax#installation) outlined by
81+
Google.
82+
83+
Note: For plotting functionality which can be found throughout our various notebooks, one
84+
must install the requirements which can be found in `requirements/requirements-plotting.txt`.
85+
86+
## Usage :rocket:
87+
88+
To import and use `S2FFT` is as simple follows:
89+
90+
For a signal on the sphere
91+
92+
``` python
93+
# Compute harmonic coefficients
94+
flm = s2fft.forward_jax(f, L)
95+
# Map back to pixel-space signal
96+
f = s2fft.inverse_jax(flm, L)
97+
```
98+
99+
For a signal on the rotation group
100+
101+
``` python
102+
# Compute Wigner coefficients
103+
flmn = s2fft.wigner.forward_jax(f, L, N)
104+
# Map back to pixel-space signal
105+
f = fft.wigner.inverse_jax(flmn, L, N)
106+
```
107+
108+
## Benchmarking :hourglass_flowing_sand:
109+
110+
We benchmarked the spherical harmonic and Wigner transforms implemented
111+
in `S2FFT` against the C implementations in the
112+
[SSHT](https://github.com/astro-informatics/ssht) pacakge.
113+
114+
A brief summary is shown in the table below for the recursion (left) and
115+
precompute (right) algorithms, with `S2FFT` running on GPUs (for further
116+
details see Price & McEwen, in prep.). Note that our compute time is
117+
agnostic to spin number (which is not the case for many other methods
118+
that scale linearly with spin).
119+
120+
| L | Wall-Time | Speed-up | Error | Wall-Time | Speed-up | Error | Memory |
121+
|------|-----------|----------|----------|-----------|----------|----------|---------|
122+
| 64 | 3.6 ms | 0.88 | 1.81E-15 | 52.4 μs | 60.5 | 1.67E-15 | 4.2 MB |
123+
| 128 | 7.26 ms | 1.80 | 3.32E-15 | 162 μs | 80.5 | 3.64E-15 | 33 MB |
124+
| 256 | 17.3 ms | 6.32 | 6.66E-15 | 669 μs | 163 | 6.74E-15 | 268 MB |
125+
| 512 | 58.3 ms | 11.4 | 1.43E-14 | 3.6 ms | 184 | 1.37E-14 | 2.14 GB |
126+
| 1024 | 194 ms | 32.9 | 2.69E-14 | 32.6 ms | 195 | 2.47E-14 | 17.1 GB |
127+
| 2048 | 1.44 s | 49.7 | 5.17E-14 | N/A | N/A | N/A | N/A |
128+
| 4096 | 8.48 s | 133.9 | 1.06E-13 | N/A | N/A | N/A | N/A |
129+
| 8192 | 82 s | 110.8 | 2.14E-13 | N/A | N/A | N/A | N/A |
130+
131+
where the left hand results are for the recursive based algorithm and the right hand side are
132+
our precompute implementation.
133+
134+
## Contributors ✨
135+
136+
Thanks goes to these wonderful people ([emoji
137+
key](https://allcontributors.org/docs/en/emoji-key)):
138+
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
139+
<!-- prettier-ignore-start -->
140+
<!-- markdownlint-disable -->
141+
<table>
142+
<tbody>
143+
<tr>
144+
<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>
145+
<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>
146+
<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>
147+
<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>
148+
<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>
149+
<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>
150+
</tr>
151+
</tbody>
152+
</table>
153+
154+
<!-- markdownlint-restore -->
155+
<!-- prettier-ignore-end -->
156+
157+
<!-- ALL-CONTRIBUTORS-LIST:END -->
158+
We encourage contributions from any interested developers. A simple
159+
first addition could be adding support for more spherical sampling
160+
patterns!
161+
162+
## Attribution :books:
163+
164+
Should this code be used in any way, we kindly request that the following article is
165+
referenced. A BibTeX entry for this reference may look like:
166+
167+
```
168+
@article{price:s2fft,
169+
AUTHOR = "Matthew A. Price and Jason D. McEwen and Contributors",
170+
TITLE = "TBA",
171+
YEAR = "2023",
172+
EPRINT = "arXiv:0000.00000"
173+
}
174+
```
175+
176+
You might also like to consider citing our related papers on which this
177+
code builds:
178+
179+
```
180+
@article{mcewen:fssht,
181+
AUTHOR = "Jason D. McEwen and Yves Wiaux",
182+
TITLE = "A novel sampling theorem on the sphere",
183+
JOURNAL = "IEEE Trans. Sig. Proc.",
184+
YEAR = "2011",
185+
VOLUME = "59",
186+
NUMBER = "12",
187+
PAGES = "5876--5887",
188+
EPRINT = "arXiv:1110.6298",
189+
DOI = "10.1109/TSP.2011.2166394"
190+
}
191+
```
192+
193+
```
194+
@article{mcewen:so3,
195+
AUTHOR = "Jason D. McEwen and Martin B{\"u}ttner and Boris ~Leistedt and Hiranya V. Peiris and Yves Wiaux",
196+
TITLE = "A novel sampling theorem on the rotation group",
197+
JOURNAL = "IEEE Sig. Proc. Let.",
198+
YEAR = "2015",
199+
VOLUME = "22",
200+
NUMBER = "12",
201+
PAGES = "2425--2429",
202+
EPRINT = "arXiv:1508.03101",
203+
DOI = "10.1109/LSP.2015.2490676"
204+
}
205+
```
206+
207+
## License :memo:
208+
209+
We provide this code under an MIT open-source licence with the hope that
210+
it will be of use to a wider community.
211+
212+
Copyright 2023 Matthew Price, Jason McEwen and contributors.
213+
214+
`S2FFT` is free software made available under the MIT License. For
215+
details see the LICENSE file.

0 commit comments

Comments
 (0)