You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
# Differentiable and accelerated spherical transforms with JAX
12
+
# Differentiable and accelerated spherical transforms
13
13
14
-
`S2FFT` is a JAX package for computing Fourier transforms on the sphere
15
-
and rotation group [(Price & McEwen 2023)](https://arxiv.org/abs/2311.14670). It leverages autodiff to provide differentiable
16
-
transforms, which are also deployable on hardware accelerators
17
-
(e.g. GPUs and TPUs).
14
+
`S2FFT` is a Python package for computing Fourier transforms on the sphere
15
+
and rotation group [(Price & McEwen 2023)](https://arxiv.org/abs/2311.14670) using
16
+
JAX or PyTorch. It leverages autodiff to provide differentiable transforms, which are
17
+
also deployable on hardware accelerators (e.g. GPUs and TPUs).
18
18
19
19
More specifically, `S2FFT` provides support for spin spherical harmonic
20
20
and Wigner transforms (for both real and complex signals), with support
21
21
for adjoint transformations where needed, and comes with different
22
22
optimisations (precompute or not) that one may select depending on
23
23
available resources and desired angular resolution $L$.
24
24
25
+
As of version 1.0.2 `S2FFT` also provides PyTorch implementations of underlying
26
+
precompute transforms. In future releases this support will be extended to our
27
+
on-the-fly algorithms.
28
+
25
29
## Algorithms :zap:
26
30
27
31
`S2FFT` leverages new algorithmic structures that can he highly
@@ -79,21 +83,33 @@ The Python dependencies for the `S2FFT` package are listed in the file
79
83
`requirements/requirements-core.txt` and will be automatically installed
80
84
into the active python environment by [pip](https://pypi.org) when running
81
85
86
+
```bash
87
+
pip install s2fft
88
+
```
89
+
This will install all core functionality which includes JAX support. To install `S2FFT`
90
+
with PyTorch support run
91
+
92
+
```bash
93
+
pip install s2fft[torch]
94
+
```
95
+
96
+
Alternatively, the `S2FFT` package may be installed directly from GitHub by cloning this
97
+
repository and then running
98
+
82
99
```bash
83
100
pip install .
84
101
```
85
102
86
-
from the root directory of the repository. Unit tests can then be
87
-
executed to ensure the installation was successful by running
103
+
from the root directory of the repository. To enable PyTorch support you will need to run
88
104
89
105
```bash
90
-
pytest tests/
106
+
pip install .[torch]
91
107
```
92
108
93
-
Alternatively, the `S2FFT` package may be installed directly from PyPi by running
109
+
Unit tests can then be executed to ensure the installation was successful by running
94
110
95
111
```bash
96
-
pip install s2fft
112
+
pytest tests/
97
113
```
98
114
99
115
> [!NOTE]
@@ -123,6 +139,9 @@ f = fft.wigner.inverse_jax(flmn, L, N)
123
139
124
140
For further details on usage see the [documentation](https://astro-informatics.github.io/s2fft/) and associated [notebooks](https://astro-informatics.github.io/s2fft/tutorials/spherical_harmonic/spherical_harmonic_transform.html).
125
141
142
+
> [!NOTE]
143
+
> We also provide PyTorch support for the precompute version of our transforms. These are called through forward/inverse_torch(). Full PyTorch support will be provided in future releases.
144
+
126
145
## Benchmarking :hourglass_flowing_sand:
127
146
128
147
We benchmarked the spherical harmonic and Wigner transforms implemented
0 commit comments