Merge pull request #185 from astro-informatics/features/pytorch_precompute_transforms

add pytorch support for precompute transform

Author: CosmoMatt

.github/workflows/docs.yml

@@ -29,7 +29,7 @@ jobs:
         pip install jaxlib
         pip install -r requirements/requirements-core.txt
         pip install -r requirements/requirements-docs.txt
-        pip install .
+        pip install .\[torch\]
   
     - name: Build Documentation
       run: |

.github/workflows/tests.yml

@@ -30,7 +30,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install -r requirements/requirements-tests.txt
         pip install -r requirements/requirements-core.txt
-        pip install .
+        pip install .\[torch\]
     
     - name: Run tests
       run: |

.pip_readme.rst

@@ -27,6 +27,10 @@ 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.
+
 Documentation
 =============
 Read the full documentation `here <https://astro-informatics.github.io/s2fft/>`_.

README.md

@@ -22,6 +22,10 @@ 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.
+
 ## Algorithms :zap:
 
 `S2FFT` leverages new algorithmic structures that can he highly
@@ -123,6 +127,9 @@ f = fft.wigner.inverse_jax(flmn, L, N)
 
 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.
+
 ## Benchmarking :hourglass_flowing_sand:
 
 We benchmarked the spherical harmonic and Wigner transforms implemented

docs/api/precompute_transforms/index.rst

@@ -11,17 +11,21 @@ Precompute Functions
    * - Function Name
      - Description
    * - :func:`~s2fft.precompute_transforms.spherical.inverse`
-     - Wrapper function around NumPy/JAX inverse methods
+     - Wrapper function around NumPy/JAX/Torch inverse methods
    * - :func:`~s2fft.precompute_transforms.spherical.inverse_transform` 
      - Inverse spherical harmonic transform (NumPy)
    * - :func:`~s2fft.precompute_transforms.spherical.inverse_transform_jax`
      - Inverse spherical harmonic transform (JAX)
+   * - :func:`~s2fft.precompute_transforms.spherical.inverse_transform_torch`
+     - Inverse spherical harmonic transform (Torch)
    * - :func:`~s2fft.precompute_transforms.spherical.forward` 
-     - Wrapper function around NumPy/JAX forward methods
+     - Wrapper function around NumPy/JAX/Torch forward methods
    * - :func:`~s2fft.precompute_transforms.spherical.forward_transform`
      - Forward spherical harmonic transform (NumPy)
    * - :func:`~s2fft.precompute_transforms.spherical.forward_transform_jax`
      - Forward spherical harmonic transform (JAX)
+   * - :func:`~s2fft.precompute_transforms.spherical.forward_transform_torch`
+     - Forward spherical harmonic transform (Torch)
 
 .. list-table:: Wigner transforms.
    :widths: 25 25
@@ -30,17 +34,21 @@ Precompute Functions
    * - Function Name
      - Description
    * - :func:`~s2fft.precompute_transforms.wigner.inverse`
-     - Wrapper function around NumPy/JAX inverse methods
+     - Wrapper function around NumPy/JAX/Torch inverse methods
    * - :func:`~s2fft.precompute_transforms.wigner.inverse_transform`
      - Inverse Wigner transform (NumPy)
    * - :func:`~s2fft.precompute_transforms.wigner.inverse_transform_jax`
      - Inverse Wigner transform (JAX)
+   * - :func:`~s2fft.precompute_transforms.wigner.inverse_transform_torch`
+     - Inverse Wigner transform (Torch)
    * - :func:`~s2fft.precompute_transforms.wigner.forward`
-     - Wrapper function around NumPy/JAX forward methods
+     - Wrapper function around NumPy/JAX/Torch forward methods
    * - :func:`~s2fft.precompute_transforms.wigner.forward_transform`
      - Forward Wigner transform (NumPy)
    * - :func:`~s2fft.precompute_transforms.wigner.forward_transform_jax`
      - Forward Wigner transform (JAX)
+   * - :func:`~s2fft.precompute_transforms.wigner.forward_transform_torch`
+     - Forward Wigner transform (Torch)
 
 .. list-table:: Constructing Kernels for precompute transforms.
    :widths: 25 25

docs/api/utility/index.rst

@@ -16,20 +16,28 @@ Utility Functions
      - Computes the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing (NumPy).
    * - :func:`~s2fft.utils.healpix_ffts.healpix_ifft_jax`
      - Computes the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing (JAX).
+   * - :func:`~s2fft.utils.healpix_ffts.healpix_ifft_torch`
+     - Computes the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing (Torch).
    * - :func:`~s2fft.utils.healpix_ffts.healpix_fft`
      - Wrapper function for the Forward Fast Fourier Transform with spectral back-projection in the polar regions to manually enforce Fourier periodicity.
    * - :func:`~s2fft.utils.healpix_ffts.healpix_fft_numpy`
      - Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions (NumPy).
    * - :func:`~s2fft.utils.healpix_ffts.healpix_fft_jax`
-     - Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions (NumPy).
+     - Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions (JAX).
+   * - :func:`~s2fft.utils.healpix_ffts.healpix_fft_torch`
+     - Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions (Torch).
    * - :func:`~s2fft.utils.healpix_ffts.spectral_folding`
      - Folds higher frequency Fourier coefficients back onto lower frequency coefficients (NumPy).
    * - :func:`~s2fft.utils.healpix_ffts.spectral_folding_jax`
      - Folds higher frequency Fourier coefficients back onto lower frequency coefficients (JAX).
+   * - :func:`~s2fft.utils.healpix_ffts.spectral_folding_torch`
+     - Folds higher frequency Fourier coefficients back onto lower frequency coefficients (Torch).
    * - :func:`~s2fft.utils.healpix_ffts.spectral_periodic_extension`
      - Extends lower frequency Fourier coefficients onto higher frequency coefficients (NumPy).
    * - :func:`~s2fft.utils.healpix_ffts.spectral_periodic_extension_jax`
      - Extends lower frequency Fourier coefficients onto higher frequency coefficients (JAX).
+   * - :func:`~s2fft.utils.healpix_ffts.spectral_periodic_extension_torch`
+     - Extends lower frequency Fourier coefficients onto higher frequency coefficients (Torch).
 
 
 .. list-table:: Quadrature functions.
@@ -61,8 +69,9 @@ Utility Functions
 
 .. note::
 
-      JAX versions of these functions share an almost identical function trace and 
-      are simply accessed by the sub-module :func:`~s2fft.utils.quadrature_jax`.
+      JAX and Torch versions of these functions share an almost identical function trace and 
+      are simply accessed by the sub-modules :func:`~s2fft.utils.quadrature_jax` and 
+      :func:`~s2fft.utils.quadrature_torch` respectively.
 
 .. list-table:: Periodic resampling functions
    :widths: 25 25
@@ -102,8 +111,9 @@ Utility Functions
 
 .. note::
 
-      JAX versions of these functions share an almost identical function trace and 
-      are simply accessed by the sub-module :func:`~s2fft.utils.resampling_jax`.
+      JAX and Torch versions of these functions share an almost identical function trace and 
+      are simply accessed by the sub-modules :func:`~s2fft.utils.resampling_jax` and 
+      :func:`~s2fft.utils.resampling_torch` respectively.
 
 .. list-table:: Rotation functions
    :widths: 25 25
@@ -124,8 +134,10 @@ Utility Functions
    signal_generator
    resampling
    resampling_jax
+   resampling_torch
    quadrature
    quadrature_jax
+   quadrature_torch
    healpix_ffts
    utils
    rotation

docs/api/utility/quadrature_torch.rst

@@ -0,0 +1,7 @@
+:html_theme.sidebar_secondary.remove:
+
+**************************
+quadrature_torch
+**************************
+.. automodule:: s2fft.utils.quadrature_torch
+   :members: 

docs/api/utility/resampling_torch.rst

@@ -0,0 +1,7 @@
+:html_theme.sidebar_secondary.remove:
+
+**************************
+resampling_torch
+**************************
+.. automodule:: s2fft.utils.resampling_torch
+   :members: 

docs/conf.py

@@ -25,9 +25,9 @@
 author = "Matthew Price, Jason McEwen, Matthew Graham, Sofia Miñano, Devaraj Gopinathan"
 
 # The short X.Y version
-version = "1.0.1"
+version = "1.0.2"
 # The full version, including alpha/beta/rc tags
-release = "1.0.1"
+release = "1.0.2"
 
 
 # -- General configuration ---------------------------------------------------
@@ -112,12 +112,12 @@
         #     "icon": "fa-brands fa-youtube fa-2x",
         #     "type": "fontawesome",
         # },
-        # {
-        #     "name": "PyPi",
-        #     "url": "https://github.com/astro-informatics/s2fft/",
-        #     "icon": "_static/pypi.png",
-        #     "type": "local",
-        # },
+        {
+            "name": "PyPi",
+            "url": "https://pypi.org/project/s2fft/",
+            "icon": "_static/pypi.png",
+            "type": "local",
+        },
         {
             "name": "GitHub",
             "url": "https://github.com/astro-informatics/s2fft/",

docs/index.rst

@@ -10,6 +10,10 @@ transforms (for both real and complex signals), with support for adjoint transfo
 where needed, and comes with different optimisations (precompute or not) that one
 may select depending on available resources and desired angular resolution :math:`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.
+
 Algorithms |:zap:|
 -------------------