diff --git a/.pip_readme.rst b/.pip_readme.rst index c4a7837e..0e796e03 100644 --- a/.pip_readme.rst +++ b/.pip_readme.rst @@ -13,11 +13,11 @@ .. image:: https://colab.research.google.com/assets/colab-badge.svg :target: https://colab.research.google.com/drive/1YmJ2ljsF8HBvhPmD4hrYPlyAKc4WPUgq?usp=sharing -Differentiable and accelerated spherical transforms with JAX +Differentiable and accelerated spherical transforms ================================================================================================================= -`S2FFT` is a JAX package for computing Fourier transforms on the sphere -and rotation group. It leverages autodiff to provide differentiable +`S2FFT` is a Python package for computing Fourier transforms on the sphere +and rotation group in JAX and PyTorch. It leverages autodiff to provide differentiable transforms, which are also deployable on hardware accelerators (e.g. GPUs and TPUs). diff --git a/README.md b/README.md index c6b08d59..5e86d4c6 100644 --- a/README.md +++ b/README.md @@ -9,12 +9,12 @@ -# Differentiable and accelerated spherical transforms with JAX +# Differentiable and accelerated spherical transforms -`S2FFT` is a JAX package for computing Fourier transforms on the sphere -and rotation group [(Price & McEwen 2023)](https://arxiv.org/abs/2311.14670). It leverages autodiff to provide differentiable -transforms, which are also deployable on hardware accelerators -(e.g. GPUs and TPUs). +`S2FFT` is a Python package for computing Fourier transforms on the sphere +and rotation group [(Price & McEwen 2023)](https://arxiv.org/abs/2311.14670) Using +JAX or 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 @@ -83,21 +83,33 @@ 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 +``` bash +pip install s2fft +``` +This will install all core functionality which includes JAX support. To install `S2FFT` +with PyTorch support run the following + +``` bash +pip install s2fft[torch] +``` + +Alternatively, the `S2FFT` package may be installed directly from GitHub by cloning this +repository and then running + ``` bash pip install . ``` -from the root directory of the repository. Unit tests can then be -executed to ensure the installation was successful by running +from the root directory of the repository. To enable PyTorch support you will need to run ``` bash -pytest tests/ +pip install .[torch] ``` -Alternatively, the `S2FFT` package may be installed directly from PyPi by running +Unit tests can then be executed to ensure the installation was successful by running ``` bash -pip install s2fft +pytest tests/ ``` > [!NOTE] diff --git a/docs/index.rst b/docs/index.rst index 9271b3e0..9277de02 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,8 +1,9 @@ Differentiable and accelerated spherical transforms =================================================== -``S2FFT`` is a JAX package for computing Fourier transforms on the sphere and rotation -group `(Price & McEwen 2023) `_. It leverages autodiff to provide differentiable transforms, which are also +``S2FFT`` is a Python package for computing Fourier transforms on the sphere and rotation +group `(Price & McEwen 2023) `_ in 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 @@ -65,12 +66,17 @@ key `_): - - - - - - + + + + + + + + + + +
Matt Price
Matt Price
💻 👀 🤔		Jason McEwen
Jason McEwen
💻 👀 🤔		Matt Graham
Matt Graham
💻 👀		sfmig
sfmig
💻 👀		Devaraj Gopinathan
Devaraj Gopinathan
💻		Francois Lanusse
Francois Lanusse
💻 🐛		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
🐛
diff --git a/docs/user_guide/install.rst b/docs/user_guide/install.rst index 6f83d3f2..10d15695 100644 --- a/docs/user_guide/install.rst +++ b/docs/user_guide/install.rst @@ -11,13 +11,19 @@ from PyPi, an online python package manager. Quick install (PyPi) -------------------- -The simplest way to pick up ``S2FFT`` is to install it directly from PyPi by running +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. +To install the PyTorch functionality you will need to install the subpackage by running + +.. code-block:: bash + + pip install s2fft[torch] Install from source (GitHub) ---------------------------- @@ -38,8 +44,20 @@ and pip installing locally cd s2fft pip install . -from the root directory of the repository. Unit tests can then be executed to ensure the -installation was successful by running +from the root directory of the repository. To install the Pytorch support you will need to +install the subpackage by running + +.. code-block:: bash + + pip install .[torch] + +which, depending on operating system, can sometimes be + +.. code-block:: bash + + pip install .\[torch\] + +Unit tests can then be executed to ensure the installation was successful by running .. code-block:: bash