minor tweak to documentation

astro-informatics · Feb 27, 2024 · 0f3128c · 0f3128c
1 parent f17ac17
commit 0f3128c
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 24 deletions.
diff --git 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
@@ -9,12 +9,12 @@
 
 <img align="left" height="85" width="98" src="./docs/assets/sax_logo.png">
 
-# 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
@@ -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) <https://arxiv.org/abs/2311.14670>`_.  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) <https://arxiv.org/abs/2311.14670>`_ 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 <https://allcontributors.org/docs/en/emoji-key>`_):
         <table>
         <tbody>
             <tr>
-            <td align="center" valign="top" width="16.66%"><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>
-            <td align="center" valign="top" width="16.66%"><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>
-            <td align="center" valign="top" width="16.66%"><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>
-            <td align="center" valign="top" width="16.66%"><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>
-            <td align="center" valign="top" width="16.66%"><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>
-            <td align="center" valign="top" width="16.66%"><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>
+                <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>
+                <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>
+                <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>
+                <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>
+                <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>
+                <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>
+                <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>
+            </tr>
+            <tr>
+                <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>
+                <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>
             </tr>
         </tbody>
         </table>

diff --git 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