diff --git a/.pip_readme.rst b/.pip_readme.rst deleted file mode 100644 index c1bc11b3..00000000 --- a/.pip_readme.rst +++ /dev/null @@ -1,103 +0,0 @@ -.. image:: https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml/badge.svg?branch=main - :target: https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml -.. image:: https://img.shields.io/badge/GitHub-s2fft-brightgreen.svg?style=flat - :target: https://github.com/astro-informatics/s2fft -.. image:: https://codecov.io/gh/astro-informatics/s2fft/branch/main/graph/badge.svg?token=7QYAFAAWLE - :target: https://codecov.io/gh/astro-informatics/s2fft -.. image:: https://img.shields.io/badge/License-MIT-yellow.svg - :target: https://opensource.org/licenses/MIT -.. image:: http://img.shields.io/badge/arXiv-2311.14670-orange.svg?style=flat - :target: https://arxiv.org/abs/2311.14670 -.. image:: https://img.shields.io/badge/code%20style-black-000000.svg - :target: https://github.com/psf/black -.. image:: https://colab.research.google.com/assets/colab-badge.svg - :target: https://colab.research.google.com/github/astro-informatics/s2fft/blob/main/notebooksspherical_harmonic_transform.ipynb - -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: - -.. code-block:: - - @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: - -.. code-block:: - - @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" - } - - -.. code-block:: - - @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. diff --git a/MANIFEST.in b/MANIFEST.in index 1f1b5611..2e363a26 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -2,7 +2,6 @@ graft s2fft graft tests graft docs -include .pip_readme.rst include README.rst include pytest.ini include LICENCE.txt diff --git a/README.md b/README.md index 165f9af2..454a692e 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,21 @@ +
+ +s2fft logo - schematic representation of a tiled sphere + +# S2FFT: differentiable and accelerated spherical transforms + [![Tests status](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml) [![Linting status](https://github.com/astro-informatics/s2fft/actions/workflows/linting.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/linting.yml) [![Documentation status](https://github.com/astro-informatics/s2fft/actions/workflows/docs.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/docs.yml) [![Codecov](https://codecov.io/gh/astro-informatics/s2fft/branch/main/graph/badge.svg?token=7QYAFAAWLE)](https://codecov.io/gh/astro-informatics/s2fft) [![MIT License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) [![PyPI package](https://badge.fury.io/py/s2fft.svg)](https://badge.fury.io/py/s2fft) -[![arXiv](http://img.shields.io/badge/arXiv-2311.14670-orange.svg?style=flat)](https://arxiv.org/abs/2311.14670) -[![All Contributors](https://img.shields.io/badge/all_contributors-11-orange.svg?style=flat-square)](#contributors-) - +[![arXiv](http://img.shields.io/badge/arXiv-2311.14670-orange.svg?style=flat)](https://arxiv.org/abs/2311.14670) +![All Contributors](https://img.shields.io/github/all-contributors/astro-informatics/s2fft?color=ee8449&style=flat-square) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/astro-informatics/s2fft/blob/main/notebooks/spherical_harmonic_transform.ipynb) [![Linter](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) - - -# Differentiable and accelerated spherical transforms +
`S2FFT` is a Python package for computing Fourier transforms on the sphere and rotation group [(Price & McEwen 2024)](https://arxiv.org/abs/2311.14670) using @@ -25,21 +28,7 @@ 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$. -> [!IMPORTANT] -> HEALPix long JIT compile time fixed for CPU! Fix for GPU coming soon. - -> [!TIP] -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. - -> [!TIP] -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. - -## Algorithms :zap: +## Algorithms ⚑ `S2FFT` leverages new algorithmic structures that can he highly parallelised and distributed, and so map very well onto the architecture @@ -48,22 +37,27 @@ algorithms are based on new Wigner-d recursions that are stable to high angular resolution $L$. The diagram below illustrates the recursions (for further details see [Price & McEwen 2024]((https://arxiv.org/abs/2311.14670))). -![image](./docs/assets/figures/Wigner_recursion_legend_darkmode.png) +
+Schematic of Wigner recursions +
+ 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 L ~ 1024. +with an equally extreme memory overhead, which is infeasible at $L \sim 1024$. Alternatively, the real polar-d functions may calculated recursively, computing only a portion of the projection at a time, hence incurring negligible memory overhead at the cost of slightly slower execution. The diagram below illustrates the separable spherical harmonic transform (for further details see [Price & McEwen 2024]((https://arxiv.org/abs/2311.14670))). -![image](./docs/assets/figures/sax_schematic_legend_darkmode.png) +
+Schematic of forward and inverse spherical harmonic transforms +
-## Sampling :earth_africa: +## Sampling 🌍 The structure of the algorithms implemented in `S2FFT` can support any isolatitude sampling scheme. A number of sampling schemes are currently @@ -86,87 +80,133 @@ 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. -

+
+Visualization of spherical sampling schemes +
> [!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](https://astro-informatics.github.io/s2fft/tutorials/spherical_harmonic/JAX_HEALPix_backend.html)). Fix for GPU execution is coming soon. +> 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)). -## Installation :computer: +## Installation πŸ’» -The Python dependencies for the `S2FFT` package are listed in the file -`requirements/requirements-core.txt` and will be automatically installed -into the active python environment by [pip](https://pypi.org) when running +The latest release of `S2FFT` published [on PyPI](https://pypi.org/project/s2fft/) can be installed by running -``` bash +```bash pip install s2fft ``` -This will install all core functionality which includes JAX support (including PyTorch support). -Alternatively, the `S2FFT` package may be installed directly from GitHub by cloning this -repository and then running +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](https://docs.jax.dev/en/latest/installation.html#installation) +and then install `S2FFT` as above. + +Alternatively, the latest development version of `S2FFT` may be installed directly from GitHub by running + +```bash +pip install git+https://github.com/astro-informatics/s2fft +``` + +## Tests 🚦 + +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](https://setuptools.pypa.io/en/latest/userguide/development_mode.html)) +with the extra test dependencies by running from the root of the repository + +```bash +pip install -e ".[tests]" +``` + +To run the tests, run from the root of the repository -``` bash -pip install . +```bash +pytest ``` -from the root directory of the repository. +## Documentation πŸ“– -Unit tests can then be executed to ensure the installation was successful by first installing the test requirements and then running pytest +Documentation for the released version is available [here](https://astro-informatics.github.io/s2fft/). +To install the documentation dependencies, clone the repository and install the package (in [editable mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html)) +with the extra documentation dependencies by running from the root of the repository -``` bash -pip install -r requirements/requirements-tests.txt -pytest tests/ +```bash +pip install -e ".[docs]" ``` -Documentation for the released version is available [here](https://astro-informatics.github.io/s2fft/). To build the documentation locally run +To build the documentation, run from the root of the repository -``` bash -pip install -r requirements/requirements-docs.txt +```bash cd docs make html open _build/html/index.html ``` -> [!NOTE] -> For plotting functionality which can be found throughout our various notebooks, one must install the requirements which can be found in `requirements/requirements-plotting.txt`. +## Notebooks πŸ““ + +A series of tutorial notebooks are included in the `notebooks` directory +and rendered [in the documentation](https://astro-informatics.github.io/s2fft/tutorials/index.html). + +To install the dependencies required to run the notebooks locally, clone the repository and install the package (in [editable mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html)) +with the extra documentation and plotting dependencies by running from the root of the repository + +```bash +pip install -e ".[docs,plotting]" +``` + +To run the notebooks in Jupyter Lab, run from the root of the repository + +```bash +jupyter lab +``` -## Usage :rocket: +## Usage πŸš€ To import and use `S2FFT` is as simple follows: For a signal on the sphere -``` python +```python +import s2fft + +# Define sampled signal to transform and harmonic bandlimit +f = ... +L = ... # Compute harmonic coefficients -flm = s2fft.forward_jax(f, L) +flm = s2fft.forward(f, L, method="jax") # Map back to pixel-space signal -f = s2fft.inverse_jax(flm, L) +f = s2fft.inverse(flm, L, method="jax") ``` For a signal on the rotation group -``` python +```python +import s2fft + +# Define sampled signal to transform and harmonic and azimuthal bandlimits +f = ... +L = ... +N = ... # Compute Wigner coefficients -flmn = s2fft.wigner.forward_jax(f, L, N) +flmn = s2fft.wigner.forward(f, L, N, method="jax") # Map back to pixel-space signal -f = fft.wigner.inverse_jax(flmn, L, N) +f = fft.wigner.inverse_jax(flmn, L, N, method="jax") ``` 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). > [!NOTE] -> 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. +> We also provide PyTorch support for the precompute version of our transforms, as demonstrated in the [_Torch frontend_ tutorial notebook](https://astro-informatics.github.io/s2fft/tutorials/torch_frontend/torch_frontend.html). -## C/C++ JAX Frontends for SSHT/HEALPix :bulb: +## SSHT & HEALPix wrappers πŸ’‘ `S2FFT` also provides JAX support for existing C/C++ packages, specifically [`HEALPix`](https://healpix.jpl.nasa.gov) and [`SSHT`](https://github.com/astro-informatics/ssht). This works -by wrapping python bindings with custom JAX frontends. Note that this C/C++ to JAX interoperability is currently limited to CPU. +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: ``` python # Forward SSHT spherical harmonic transform -flm = s2fft.forward(f, L, sampling=["mw"], method="jax_ssht") +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") @@ -179,31 +219,9 @@ applications! For further details on usage see the associated [notebooks](https://astro-informatics.github.io/s2fft/tutorials/spherical_harmonic/JAX_SSHT_backend.html). - +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. ## Contributors ✨ @@ -240,7 +258,7 @@ We encourage contributions from any interested developers. A simple first addition could be adding support for more spherical sampling patterns! -## Attribution :books: +## 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: @@ -289,7 +307,7 @@ code builds: } ``` -## License :memo: +## License πŸ“ We provide this code under an MIT open-source licence with the hope that it will be of use to a wider community. @@ -297,13 +315,13 @@ 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`](LICENCE.txt) file. +details see the [`LICENCE.txt`](https://github.com/astro-informatics/s2fft/blob/main/LICENCE.txt) file. -The file [`lib/include/kernel_helpers.h`](lib/include/kernel_helpers.h) is adapted from +The file [`lib/include/kernel_helpers.h`](https://github.com/astro-informatics/s2fft/blob/main/lib/include/kernel_helpers.h) is adapted from [code](https://github.com/dfm/extending-jax/blob/c33869665236877a2ae281f3f5dbff579e8f5b00/lib/kernel_helpers.h) in [a tutorial on extending JAX](https://github.com/dfm/extending-jax) by [Dan Foreman-Mackey](https://github.com/dfm) and licensed under a [MIT license](https://github.com/dfm/extending-jax/blob/371dca93c6405368fa8e71690afd3968d75f4bac/LICENSE). -The file [`lib/include/kernel_nanobind_helpers.h`](lib/include/kernel_nanobind_helpers.h) +The file [`lib/include/kernel_nanobind_helpers.h`](https://github.com/astro-informatics/s2fft/blob/main/lib/include/kernel_nanobind_helpers.h) is adapted from [code](https://github.com/jax-ml/jax/blob/3d389a7fb440c412d95a1f70ffb91d58408247d0/jaxlib/kernel_nanobind_helpers.h) by the [JAX](https://github.com/jax-ml/jax) authors and licensed under a [Apache-2.0 license](https://github.com/jax-ml/jax/blob/3d389a7fb440c412d95a1f70ffb91d58408247d0/LICENSE). diff --git a/docs/_static/css/custom.css b/docs/_static/css/custom.css index 32f8edc1..9fffade8 100644 --- a/docs/_static/css/custom.css +++ b/docs/_static/css/custom.css @@ -11,7 +11,7 @@ .bd-sidebar-secondary { flex-grow: 1; - max-width: 180px; + max-width: 220px; } i.fa-youtube:before { @@ -155,4 +155,4 @@ section#examples li.toctree-l1>a { /* So that Tutorials/Reference are not clickable in the sidebar */ li.toctree-l2.has-children>a { pointer-events: none; -} \ No newline at end of file +} diff --git a/docs/conf.py b/docs/conf.py index f7b09146..cd26fb92 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -13,6 +13,7 @@ import os import sys +from importlib.metadata import version as get_version sys.path.insert(0, os.path.abspath("..")) @@ -23,10 +24,8 @@ copyright = "2023, Matthew Price and Jason McEwen" author = "Matthew Price, Jason McEwen, Matthew Graham, Sofia MiΓ±ano, Devaraj Gopinathan" -# The short X.Y version -version = "1.1.1" -# The full version, including alpha/beta/rc tags -release = "1.1.1" +release = get_version("s2fft") +version = ".".join(release.split(".")[:2]) # -- General configuration --------------------------------------------------- @@ -53,6 +52,7 @@ "sphinxcontrib.texfigure", "sphinx.ext.autosectionlabel", "sphinxemoji.sphinxemoji", + "sphinx_mdinclude", ] nbsphinx_execute = "never" @@ -138,21 +138,6 @@ ], } -html_sidebars = { - # "tutorials/*": [ - # "indices.html", - # "navbar-nav.html", - # ], - "user_guide/*": [ - "indices.html", - "navbar-nav.html", - ], - # "background/*": [ - # "indices.html", - # "navbar-nav.html", - # ], -} - html_static_path = ["_static"] html_css_files = [ "css/custom.css", diff --git a/docs/index.rst b/docs/index.rst index 1ad9fb2c..f797350f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,182 +1,9 @@ -Differentiable and accelerated spherical transforms -=================================================== - -``S2FFT`` is a Python package for computing Fourier transforms on the sphere and rotation -group `(Price & McEwen 2023) `_ using JAX and PyTorch. -It leverages autodiff to provide differentiable transforms, which are also -deployable on modern 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 :math:`L`. - -.. important:: - HEALPix long JIT compile time fixed for CPU! Fix for GPU coming soon. - -.. tip:: - 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. - -.. tip:: - 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. - -Algorithms |:zap:| -------------------- - -``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 :math:`L`. The diagram below illustrates the -recursions (for further details see Price & McEwen 2023). - -.. image:: ./assets/figures/Wigner_recursion_github_docs.png - -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 L ~ 1024. -Alternatively, the real polar-d functions may calculated recursively, -computing only a portion of the projection at a time, hence incurring -negligible memory overhead at the cost of slightly slower execution. The -diagram below illustrates the separable spherical harmonic transform. - -.. image:: ./assets/figures/sax_schematic_github_docs.png - -.. 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 `_). Fix for GPU execution is coming soon. - -Sampling |:earth_africa:| ------------------------------------ - -The structure of the algorithms implemented in ``S2FFT`` can support any isolattitude 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. - -.. image:: ./assets/figures/spherical_sampling.png - :width: 900 - :align: center - -Contributors ✨ ------------------------------------ - -Thanks goes to these wonderful people (`emoji -key `_): - -.. raw:: html - - - - - - - - - - - - - - - - - - - - -
Matt Price
Matt Price
πŸ’» πŸ‘€ πŸ€”		Jason McEwen
Jason McEwen
πŸ’» πŸ‘€ πŸ€”		Matt Graham
Matt Graham
πŸ’» πŸ‘€		sfmig
sfmig
πŸ’» πŸ‘€		Devaraj Gopinathan
Devaraj Gopinathan
πŸ’»		Francois Lanusse
Francois Lanusse
πŸ’» πŸ›		Ikko Eltociear Ashimine
Ikko Eltociear Ashimine
πŸ“–
Kevin Mulder
Kevin Mulder
πŸ›		Philipp Misof
Philipp Misof
πŸ›		Elis Roberts
Elis Roberts
πŸ› πŸ“–		Wassim KABALAN
Wassim KABALAN
πŸ’» πŸ‘€ ⚠️
- - -We encourage contributions from any interested developers. A simple -first addition could be adding support for more spherical sampling -patterns! - -Attribution |:books:| ------------------- - -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: - -.. code-block:: - - @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: - -.. code-block:: - - @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" - } - -.. code-block:: - - @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 |:memo:| ----------------- - -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. - -.. bibliography:: - :notcited: - :list: bullet - -.. * :ref:`modindex` - -.. toctree:: - :hidden: - :maxdepth: 2 - :caption: User Guide - - user_guide/install +.. mdinclude:: ../README.md .. toctree:: :hidden: :maxdepth: 3 - :caption: Interactive Tutorials + :caption: Tutorials tutorials/index @@ -186,4 +13,3 @@ the LICENSE file. :caption: API api/index - diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index 20702c99..08f89bcf 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -1,57 +1,42 @@ :html_theme.sidebar_secondary.remove: -************************** -Notebooks -************************** -A series of tutorial notebooks which go through the absolute base level application of -``S2FFT`` apis. Post alpha release we will add examples for more involved applications, -in the time being feel free to contact contributors for advice! At a high-level the -``S2FFT`` package is structured such that the 2 primary transforms, the Wigner and -spherical harmonic transforms, can easily be accessed. +***************************** +Tutorials +***************************** + +This section contains a series of tutorial notebooks which go through some of the +key features of the ``S2FFT`` package. + +At a high-level the ``S2FFT`` package is structured such that the two primary transforms, +the Wigner and spherical harmonic transforms, can easily be accessed. Core usage |:rocket:| ------------------ -To import and use ``S2FFT`` is as simple follows: +----------------------------- + +To import and use ``S2FFT`` is as simple follows: +-------------------------------------------------------+------------------------------------------------------------+ |For a signal on the sphere |For a signal on the rotation group | | | | |.. code-block:: Python |.. code-block:: Python | | | | +| import s2fft | import s2fft | +| | | +| # Specify sampled signal and harmonic bandlimit | # Define sampled signal, harmonic & azimuthal bandlimits | +| f = ... | f = ... | +| L = ... | L, N = ... | +| | | | # Compute harmonic coefficients | # Compute Wigner coefficients | -| flm = s2fft.forward_jax(f, L) | flmn = s2fft.wigner.forward_jax(f, L, N) | +| flm = s2fft.forward(f, L, method="jax") | flmn = s2fft.wigner.forward(f, L, N, method="jax") | | | | | # Map back to pixel-space signal | # Map back to pixel-space signal | -| f = s2fft.inverse_jax(flm, L) | f = s2fft.wigner.inverse_jax(flmn, L, N) | +| f = s2fft.inverse(flm, L, method="jax") | f = s2fft.wigner.inverse(flmn, L, N, method="jax") | +-------------------------------------------------------+------------------------------------------------------------+ -C/C++ backend usage |:bulb:| ------------------ -``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. - -For example, one may call these alternate backends for the spherical harmonic transform by: - -.. code-block:: python - - # 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 signal processing -applications! - - .. toctree:: :hidden: :maxdepth: 3 - :caption: Jupyter Notebooks + :caption: Tutorials spherical_harmonic/spherical_harmonic_transform.nblink wigner/wigner_transform.nblink diff --git a/docs/user_guide/install.rst b/docs/user_guide/install.rst deleted file mode 100644 index 7dad3ef5..00000000 --- a/docs/user_guide/install.rst +++ /dev/null @@ -1,82 +0,0 @@ -:html_theme.sidebar_secondary.remove: - -.. _install: - -Installation -========================= -There are two primary ways to install ``S2FFT``. One can either build the project from -the most recent GitHub source, which comes with the added benefit of being able to -locally execute the unit testing. Alternately, one may simply install the package directly -from PyPi, an online python package manager. - -Quick install (PyPi) --------------------- -The simplest way to pick up ``S2FFT`` is to install the core JAX functionality directly -from PyPi by running - -.. code-block:: bash - - pip install s2fft - -after which ``S2FFT`` may be imported and run as outlined in the associated notebooks and collab tutorials. -This will include PyTorch functionality. - -Install from source (GitHub) ----------------------------- - -When installing from source we recommend working within an existing conda environment, or creating a fresh conda environment to avoid any dependency conflicts, - -.. code-block:: bash - - conda create -n "env_name" python>=3.9 - conda activate "env_name" - -Once within a fresh environment ``S2FFT`` may be installed by cloning the GitHub repository -and pip installing locally - -.. code-block:: bash - - git clone https://github.com/astro-informatics/s2fft - cd s2fft - pip install . - -from the root directory of the repository. This will include PyTorch functionality. - -Unit tests can then be executed to ensure the installation was successful by first installing the test requirements and then running pytest - -.. code-block:: bash - - pip install -r requirements/requirements-tests.txt - pytest tests/ - -Documentation for the released version is available `here `_. To build the documentation locally run - -.. code-block:: bash - - pip install -r requirements/requirements-docs.txt - cd docs - make html - open _build/html/index.html - - - -Installing JAX for NVIDIA GPUs ------------------------------- -We include both ``jax`` and ``jaxlib`` as dependencies in ``requirements/requirements-core.txt`` -however to get things running on GPUs can be a bit more involved. We strongly recommend -this installation `guide `_ provided by -Google. To summarise you will first need to install NVIDIA drivers for -`CUDA `_ and `CuDNN `_, -following which a pre-built CUDA-compatible wheels shoulld be installed by running - -.. code-block:: bash - - pip install --upgrade pip - - # Wheels only built for linux - pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html - - # Wheels built for many machine architectures - pip install "jax[cuda11_cudnn86]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html - -where the versions of CUDA and CuDNN should match those you have installed on the machine. diff --git a/pyproject.toml b/pyproject.toml index 6a2c310b..10a77fae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dynamic = [ keywords = [ ] name = "s2fft" -readme = ".pip_readme.rst" +readme = "README.md" requires-python = ">=3.8" license.file = "LICENCE.txt" urls.homepage = "https://github.com/astro-informatics/s2fft" @@ -61,6 +61,7 @@ docs = [ "sphinx-copybutton", "sphinxemoji", "sphinx_rtd_theme", + "sphinx_mdinclude", "ipython>=7.16.1", "jupyter>=1.0.0", ]