Merge branch 'main' into mmg/update-test-workflow

Author: matt-graham

.github/workflows/tests.yml

@@ -54,5 +54,9 @@ jobs:
     
     - name: Run tests
       run: |
-        pytest -v --cov-report term --cov=s2fft --cov-config=.coveragerc 
-        codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702
+        pytest -v --cov-report=xml --cov=s2fft --cov-config=.coveragerc
+  
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v3
+      env:
+        CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

README.md

@@ -5,10 +5,9 @@
 [![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-BADGE:START - Do not remove or modify this section -->
-[![All Contributors](https://img.shields.io/badge/all_contributors-11-orange.svg?style=flat-square)](#contributors-)
-<!-- ALL-CONTRIBUTORS-BADGE:END --> 
+[![All Contributors](https://img.shields.io/badge/all_contributors-11-orange.svg?style=flat-square)](#contributors-)<!-- ALL-CONTRIBUTORS-BADGE:END --> 
 [![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)
-<!-- [![image](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) -->
+[![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)
 
 <img align="left" height="85" width="98" src="./docs/assets/sax_logo.png">
 

docs/api/precompute_transforms/custom_ops.rst

@@ -0,0 +1,7 @@
+:html_theme.sidebar_secondary.remove:
+
+**************************
+Custom Operations
+**************************
+.. automodule:: s2fft.precompute_transforms.custom_ops
+   :members: 

docs/api/precompute_transforms/fourier_wigner.rst

@@ -0,0 +1,7 @@
+:html_theme.sidebar_secondary.remove:
+
+**************************
+Fourier-Wigner Transform
+**************************
+.. automodule:: s2fft.precompute_transforms.fourier_wigner
+   :members: 

docs/api/precompute_transforms/index.rst

@@ -50,6 +50,40 @@ Precompute Functions
    * - :func:`~s2fft.precompute_transforms.wigner.forward_transform_torch`
      - Forward Wigner transform (Torch)
 
+.. list-table:: Fourier-Wigner transforms.
+   :widths: 25 25
+   :header-rows: 1
+
+   * - Function Name
+     - Description
+   * - :func:`~s2fft.precompute_transforms.fourier_wigner.inverse_transform`
+     - Inverse Wigner transform with Fourier method (NumPy)
+   * - :func:`~s2fft.precompute_transforms.fourier_wigner.inverse_transform_jax`
+     - Inverse Wigner transform with Fourier method (JAX)
+   * - :func:`~s2fft.precompute_transforms.fourier_wigner.forward_transform`
+     - Forward Wigner transform with Fourier method (NumPy)
+   * - :func:`~s2fft.precompute_transforms.fourier_wigner.forward_transform_jax`
+     - Forward Wigner transform with Fourier method (JAX)
+    
+.. list-table:: Custom Operations
+   :widths: 25 25
+   :header-rows: 1
+
+   * - Function Name
+     - Description
+   * - :func:`~s2fft.precompute_transforms.custom_ops.wigner_subset_to_s2`
+     - Transforms an arbitrary subset of Wigner coefficients onto a subset of spin signals on the sphere.
+   * - :func:`~s2fft.precompute_transforms.custom_ops.wigner_subset_to_s2_jax`
+     - Transforms an arbitrary subset of Wigner coefficients onto a subset of spin signals on the sphere (JAX).
+   * - :func:`~s2fft.precompute_transforms.custom_ops.so3_to_wigner_subset`
+     - Transforms a signal on the rotation group to an arbitrary subset of its Wigner coefficients.
+   * - :func:`~s2fft.precompute_transforms.custom_ops.so3_to_wigner_subset_jax`
+     - Transforms a signal on the rotation group to an arbitrary subset of its Wigner coefficients (JAX).
+   * - :func:`~s2fft.precompute_transforms.custom_ops.s2_to_wigner_subset`
+     - Transforms from a collection of arbitrary spin signals on the sphere to the corresponding collection of their harmonic coefficients.
+   * - :func:`~s2fft.precompute_transforms.custom_ops.s2_to_wigner_subset_jax`
+     - Transforms from a collection of arbitrary spin signals on the sphere to the corresponding collection of their harmonic coefficients (JAX).
+
 .. list-table:: Constructing Kernels for precompute transforms.
    :widths: 25 25
    :header-rows: 1
@@ -60,6 +94,14 @@ Precompute Functions
      - Builds a kernel including quadrature weights and Wigner-D coefficients for spherical harmonic transform.
    * - :func:`~s2fft.precompute_transforms.construct.wigner_kernel`
      - Builds a kernel including quadrature weights and Wigner-D coefficients for Wigner transform.
+   * - :func:`~s2fft.precompute_transforms.construct.spin_spherical_kernel_jax`
+     - Builds a kernel including quadrature weights and Wigner-D coefficients for spherical harmonic transform (JAX).
+   * - :func:`~s2fft.precompute_transforms.construct.wigner_kernel_jax`
+     - Builds a kernel including quadrature weights and Wigner-D coefficients for Wigner transform (JAX).
+   * - :func:`~s2fft.precompute_transforms.construct.fourier_wigner_kernel`
+     - Builds a kernel including quadrature weights and Fourier coefficienfs of Wigner d-functions
+   * - :func:`~s2fft.precompute_transforms.construct.fourier_wigner_kernel_jax`
+     - Builds a kernel including quadrature weights and Fourier coefficienfs of Wigner d-functions (JAX).
    * - :func:`~s2fft.precompute_transforms.construct.healpix_phase_shifts`
      - Builds a vector of corresponding phase shifts for each HEALPix latitudinal ring.
 
@@ -69,6 +111,9 @@ Precompute Functions
    :caption: Precompute Transforms
 
    construct
+   alt_construct
    spin_spherical 
    wigner
+   fourier_wigner
+   custom_ops
 

docs/index.rst

@@ -78,17 +78,19 @@ key <https://allcontributors.org/docs/en/emoji-key>`_):
         <table>
         <tbody>
             <tr>
-                <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>
+            <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>
+            <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>
+            <td align="center" valign="top" width="14.28%"><a href="https://github.com/ElisR"><img src="https://avatars.githubusercontent.com/u/19764906?v=4?s=100" width="100px;" alt="Elis Roberts"/><br /><sub><b>Elis Roberts</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/issues?q=author%3AElisR" title="Bug reports">🐛</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=ElisR" title="Documentation">📖</a></td>
+            <td align="center" valign="top" width="14.28%"><a href="https://github.com/ASKabalan"><img src="https://avatars.githubusercontent.com/u/83787080?v=4?s=100" width="100px;" alt="Wassim KABALAN"/><br /><sub><b>Wassim KABALAN</b></sub></a><br /><a href="https://github.com/astro-informatics/s2fft/commits?author=ASKabalan" title="Code">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3AASKabalan" title="Reviewed Pull Requests">👀</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=ASKabalan" title="Tests">⚠️</a></td>
             </tr>
         </tbody>
         </table>

pyproject.toml

@@ -73,7 +73,6 @@ plotting = [
 tests = [
     "pytest",
     "pytest-cov",
-    "codecov",
     "so3",
 ]
 

s2fft/precompute_transforms/__init__.py

@@ -1 +1 @@
-from . import construct, spherical, wigner
+from . import construct, custom_ops, fourier_wigner, spherical, wigner