Merge pull request #78 from astro-informatics/housekeeping/clean_code_up

Housekeeping/clean code up
astro-informatics · Apr 9, 2024 · 50622a4 · 50622a4
2 parents c5702aa + 6394602
commit 50622a4
Show file tree

Hide file tree

Showing 53 changed files with 2,996 additions and 2,109 deletions.
diff --git a/.pip_readme.rst b/.pip_readme.rst
@@ -1,34 +1,44 @@
-.. image:: https://img.shields.io/badge/GitHub-s2wav-brightgreen.svg?style=flat
-    :target: https://github.com/astro-informatics/s2wav
 .. image:: https://github.com/astro-informatics/s2wav/actions/workflows/tests.yml/badge.svg?branch=main
     :target: https://github.com/astro-informatics/s2wav/actions/workflows/tests.yml
-.. image:: https://readthedocs.org/projects/ansicolortags/badge/?version=latest
-    :target: https://astro-informatics.github.io/s2wav
-.. image:: https://codecov.io/gh/astro-informatics/s2wav/branch/main/graph/badge.svg?token=ZES6J4K3KZ 
+.. image:: https://codecov.io/gh/astro-informatics/s2wav/branch/main/graph/badge.svg?token=ZES6J4K3KZ
     :target: https://codecov.io/gh/astro-informatics/s2wav
-.. image:: https://img.shields.io/badge/License-GPL-blue.svg
-    :target: http://perso.crans.org/besson/LICENSE.html
+.. image:: https://img.shields.io/badge/License-MIT-yellow.svg
+    :target: https://opensource.org/licenses/MIT
 .. image:: http://img.shields.io/badge/arXiv-2402.01282-orange.svg?style=flat
     :target: https://arxiv.org/abs/2402.01282
+.. 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/drive/15E64EAQ7TIp2a3cCoXtnNgf7Ud9MYjVq?usp=sharing
 
-s2wav
+Differentiable and accelerated wavelet transform on the sphere
 =================================================================================================================
 
-Add some basic discussion about ``s2wav`` here.
+`S2WAV` is a python package for computing wavelet transforms on the sphere
+and rotation group, both in JAX and PyTorch. It leverages autodiff to provide differentiable
+transforms, which are also deployable on modern hardware accelerators
+(e.g. GPUs and TPUs), and can be mapped across multiple accelerators.
 
-Installation
-============
+More specifically, `S2WAV` provides support for scale-discretised
+wavelet transforms on the sphere and rotation group (for both real and
+complex signals), with support for adjoints where needed, and comes with
+a variety of different optimisations (e.g. precompute or not,
+multi-resolution algorithms) that one may select depending on available
+resources and desired angular resolution $L$. `S2WAV` is a sister package of 
+`S2FFT <https://github.com/astro-informatics/s2fft>`_, both of which are part of the `SAX` 
+project, which aims to provide comprehensive support for differentiable transforms on the 
+sphere and rotation group.
+
+As of version 1.0.0 `S2WAV` also provides partial frontend support for PyTorch. In future 
+this will be expanded to full support. Also note that this release also provides JAX support 
+for existing C spherical harmonic libraries, specifically `SSHT`. This works be wrapping 
+python bindings with custom JAX frontends. Note that currently this C to JAX interoperability 
+is limited to CPU.
 
-Add some basic installation instructions here.
-
 Documentation
 =============
 
-Link to the full documentation (when deployed).
-
-Contributors
-============
-Author names & Contributors
+Read the full documentation <here `astro-informatics.github.io/s2wav/>`_.
 
 Attribution
 ===========
@@ -43,18 +53,29 @@ A BibTeX entry for s2wav is:
         eprint      = "arXiv:2402.01282"        
     }
 
-License
-=======
+we also request that you cite the following paper 
+
+.. 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, submitted",
+        year        = "2023",
+        eprint      = "arXiv:2311.14670"        
+    }
 
-s2wav is released under the GPL-3 license (see `LICENSE.txt <https://github.com/astro-informatics/s2wav/blob/main/LICENSE.txt>`_),
+in which the core underlying algorithms for the spherical harmonic and Wigner transforms 
+are developed.
 
-.. code-block::
+License
+=======
 
-     s2wav
-     Copyright (C) 2022 Author names & contributors
+We provide this code under an MIT open-source licence with the hope that
+it will be of use to a wider community.
 
-     This program is released under the GPL-3 license (see LICENSE.txt).
+Copyright 2024 Matthew Price, Jessica Whtiney, Alicja Polanska, Jason
+McEwen and contributors.
 
-     This program is distributed in the hope that it will be useful,
-     but WITHOUT ANY WARRANTY; without even the implied warranty of
-     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+`S2WAV` is free software made available under the MIT License. For
+details see the LICENSE file.
diff --git a/README.md b/README.md
@@ -10,8 +10,8 @@
 
 # Differentiable and accelerated wavelet transform on the sphere
 
-`S2WAV` is a JAX package for computing wavelet transforms on the sphere
-and rotation group. It leverages autodiff to provide differentiable
+`S2WAV` is a python package for computing wavelet transforms on the sphere
+and rotation group, both in JAX and PyTorch. It leverages autodiff to provide differentiable
 transforms, which are also deployable on modern hardware accelerators
 (e.g. GPUs and TPUs), and can be mapped across multiple accelerators.
 
@@ -20,10 +20,27 @@ wavelet transforms on the sphere and rotation group (for both real and
 complex signals), with support for adjoints where needed, and comes with
 a variety of different optimisations (e.g. precompute or not,
 multi-resolution algorithms) that one may select depending on available
-resources and desired angular resolution $L$. `S2WAV` is a sister package of [`S2FFT`](https://github.com/astro-informatics/s2fft), both of which are part of the `SAX` project, which aims to provide comprehensive support for differentiable transforms on the sphere and rotation group.
+resources and desired angular resolution $L$. `S2WAV` is a sister package of 
+[`S2FFT`](https://github.com/astro-informatics/s2fft), both of which are part of the `SAX` 
+project, which aims to provide comprehensive support for differentiable transforms on the 
+sphere and rotation group.
+
+> [!TIP]
+> As of version 1.0.0 `S2WAV` also provides partial frontend support for PyTorch. In future 
+> this will be expanded to full support. Also note that this release also provides JAX support 
+> for existing C spherical harmonic libraries, specifically `SSHT`. This works be wrapping 
+> python bindings with custom JAX frontends. Note that currently this C to JAX interoperability 
+> is limited to CPU.
 
 ## Wavelet Transform :zap: 
-`S2WAV` is an updated implementation of the scale-discretised wavelet transform on the sphere, which builds upon the papers of [Leistedt et al 2013](https://arxiv.org/abs/1211.1680) and [McEwen et al 2017](https://arxiv.org/abs/1509.06749). This wavelet transform is designed to have excellent localisation and uncorrelation properties, and has been successfully adopted for various applications e.g. scattering transforms on the sphere [McEwen et al 2022](https://arxiv.org/pdf/2102.02828.pdf). The wavelet dictionary is constructed by tiling the harmonic line with infinitely differentiable Cauchy-Schwartz functions, which can straightforwardly be performed in an efficient multiresolution manner, as in the Euclidean case. This is what the directional wavelet filters look like in pixel space.
+`S2WAV` is an updated implementation of the scale-discretised wavelet transform on the 
+sphere, which builds upon the papers of [Leistedt et al 2013](https://arxiv.org/abs/1211.1680) 
+and [McEwen et al 2017](https://arxiv.org/abs/1509.06749). This wavelet transform is designed to 
+have excellent localisation and uncorrelation properties, and has been successfully adopted for 
+various applications e.g. scattering transforms on the sphere [McEwen et al 2022](https://arxiv.org/pdf/2102.02828.pdf). 
+The wavelet dictionary is constructed by tiling the harmonic line with infinitely differentiable 
+Cauchy-Schwartz functions, which can straightforwardly be performed in an efficient multiresolution 
+manner, as in the Euclidean case. This is what the directional wavelet filters look like in pixel space.
 
 <p align="center">
   <img src="./docs/assets/figures/spherical_wavelets.png" width="700"/>
@@ -36,35 +53,75 @@ The Python dependencies for the `S2WAV` package are listed in the file
 into the active python environment by [pip](https://pypi.org) when running
 
 ``` bash
-pip install .        
+pip install s2wav     
 ```
+This will install the core functionality which includes JAX support (including PyTorch support).
 
-from the root directory of the repository. Unit tests can then be
-executed to ensure the installation was successful by running
+Alternatively, the `S2WAV` package may be installed directly from GitHub by cloning this 
+repository and then running
 
 ``` bash
-pytest tests/
+pip install .
 ```
 
-In the near future one will be able to install `S2WAV` directly from
-[PyPi](https://pypi.org) by `pip install s2wav` but this is not yet supported.
-Note that to run `JAX` on NVIDIA GPUs you will need to follow the
-[guide](https://github.com/google/jax#installation) outlined by Google.
+from the root directory.
+
+Unit tests can then be executed to ensure the installation was successful by first 
+installing the test requirements and then running pytest
+
+``` bash
+pip install -r requirements/requirements-tests.txt
+pytest tests/  
+```
+
+Documentation for the released version is available [here](https://astro-informatics.github.io/s2wav/).
+To build the documentation locally run 
+
+``` bash
+pip install -r requirements/requirements-docs.txt
+cd docs 
+make html
+open _build/html/index.html
+```
 
 ## Usage :rocket:
 
 To import and use `S2WAV` is as simple follows:
 
 ``` python
-import s2wav 
-
 # Compute wavelet coefficients
 f_wav, f_scal = s2wav.analysis(f, L, N)
 
 # Map back to signal on the sphere 
 f = s2wav.synthesis(f_wav, f_scal, L, N)
 ```
-however we strongly recommend that the multiresolution argument is set to true, as this will accelerate the transform by a factor of the total number of wavelet scales, which can be around an order of magnitude.
+> [!NOTE]  
+> However we strongly recommend that the multiresolution argument is set to true, as this 
+> will accelerate the transform by a factor of the total number of wavelet scales, which 
+> can be around an order of magnitude.
+
+## C JAX Frontends for SSHT :bulb:
+
+`S2WAV` also provides JAX support for SSHT, which is a highly optimised C library which 
+implements the underlying spherical harmonic transforms. This works by wrapping python 
+bindings with custom JAX frontends. Note that this C to JAX interoperability is currently 
+limited to CPU.
+
+For example, one may call these alternate backends for the spherical wavelet transform by:
+
+``` python
+# Compute wavelet coefficients using SSHT C library backend
+f_wav, f_scal = s2wav.analysis(f, L, N, use_c_backend=True)
+
+# Map back to signal on the sphere using SSHT C library backend
+f = s2wav.synthesis(f_wav, f_scal, L, N, use_c_backend=True)
+```
+These JAX frontends supports out of the box reverse mode automatic differentiation, 
+and under the hood is simply linking to the C packages you are familiar with. In this 
+way S2fft enhances existing packages with gradient functionality for modern scientific 
+computing or machine learning applications!
+
+For further details on usage see the associated [notebooks](https://astro-informatics.github.io/s2wav/tutorials/index.html).
 
 ## Contributors ✨
 We strongly encourage contributions from any interested developers; a
@@ -105,6 +162,21 @@ A BibTeX entry for `S2WAV` is:
 }
 ```
 
+we also request that you cite the following paper 
+
+``` 
+@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, submitted",
+   year        = "2023",
+   eprint      = "arXiv:2311.14670"        
+}
+```
+
+in which the core underlying algorithms for the spherical harmonic and Wigner transforms 
+are developed.
+
 ## License :memo:
 
 Copyright 2024 Matthew Price, Jessica Whtiney, Alicja Polanska, Jason

diff --git a/docs/api/filter_factory/filters.rst b/docs/api/filter_factory/filters.rst
@@ -1,7 +1,7 @@
 :html_theme.sidebar_secondary.remove:
 
 **************************
-Filter Generator
+Filter functions
 **************************
-.. automodule:: s2wav.filter_factory.filters
+.. automodule:: s2wav.filters
    :members: 
diff --git a/docs/api/filter_factory/index.rst b/docs/api/filter_factory/index.rst
@@ -10,26 +10,26 @@ Filter Factory
 
    * - Function Name
      - Description
-   * - :func:`~s2wav.filter_factory.filters.filters_axisym`
+   * - :func:`~s2wav.filters.filters_axisym`
      - Computes wavelet kernels :math:`\Psi^j_{\ell m}` and scaling kernel :math:`\Phi_{\ell m}` in harmonic space.
-   * - :func:`~s2wav.filter_factory.filters.filters_directional`
+   * - :func:`~s2wav.filters.filters_directional`
      - Generates the harmonic coefficients for the directional tiling wavelets in harmonic space.
-   * - :func:`~s2wav.filter_factory.filters.filters_axisym_vectorised`
-     - Vectorised implementation of :func:`~s2wav.filter_factory.filters.filters_directional`.
-   * - :func:`~s2wav.filter_factory.filters.filters_directional_vectorised`
-     - Vectorised implementation of :func:`~s2wav.filter_factory.filters.filters_directional`.
+   * - :func:`~s2wav.filters.filters_axisym_vectorised`
+     - Vectorised implementation of :func:`~s2wav.filters.filters_directional`.
+   * - :func:`~s2wav.filters.filters_directional_vectorised`
+     - Vectorised implementation of :func:`~s2wav.filters.filters_directional`.
 
 .. list-table:: Wavelet kernel functions.
    :widths: 25 25
    :header-rows: 1
 
    * - Function Name
      - Description
-   * - :func:`~s2wav.filter_factory.kernels.tiling_integrand`
+   * - :func:`~s2wav.filters.tiling_integrand`
      - Tiling integrand for scale-discretised wavelets.
-   * - :func:`~s2wav.filter_factory.kernels.part_scaling_fn`
+   * - :func:`~s2wav.filters.part_scaling_fn`
      - Computes integral used to calculate smoothly decreasing function :math:`k_{\lambda}`.
-   * - :func:`~s2wav.filter_factory.kernels.k_lam`
+   * - :func:`~s2wav.filters.k_lam`
      - Compute function :math:`k_{\lambda}` used as a wavelet generating function.
 
 .. list-table:: Wavelet tiling functions.
@@ -38,18 +38,16 @@ Filter Factory
 
    * - Function Name
      - Description
-   * - :func:`~s2wav.filter_factory.tiling.tiling_direction`
+   * - :func:`~s2wav.filters.tiling_direction`
      - Generates the harmonic coefficients for the directionality component of the tiling functions.
-   * - :func:`~s2wav.filter_factory.tiling.spin_normalization`
+   * - :func:`~s2wav.filters.spin_normalization`
      - Computes the normalization factor for spin-lowered wavelets, which is :math:`\sqrt{\frac{(l+s)!}{(l-s)!}}`.
-   * - :func:`~s2wav.filter_factory.tiling.spin_normalization_vectorised`
-     - Vectorised version of :func:`~s2wav.filter_factory.tiling.spin_normalization`.
+   * - :func:`~s2wav.filters.spin_normalization_vectorised`
+     - Vectorised version of :func:`~s2wav.filters.spin_normalization`.
 
 .. toctree::
    :hidden:
    :maxdepth: 2
-   :caption: Wavelet generators
+   :caption: Filter functions
 
-   filters
-   tiling
-   kernels
+   filters
diff --git a/docs/api/index.rst b/docs/api/index.rst
@@ -9,7 +9,7 @@ directory structure for the software.
 
 .. toctree::
    :hidden:
-   :maxdepth: 1
+   :maxdepth: 2
    :caption: Namespaces
 
    transforms/index

diff --git a/docs/api/filter_factory/kernels.rst → docs/api/transforms/base.rst b/docs/api/filter_factory/kernels.rst → docs/api/transforms/base.rst
@@ -1,7 +1,7 @@
 :html_theme.sidebar_secondary.remove:
 
 **************************
-Wavelet Kernels
+Numpy Transforms
 **************************
-.. automodule:: s2wav.filter_factory.kernels
+.. automodule:: s2wav.transforms.base
    :members: 
diff --git a/docs/api/filter_factory/tiling.rst → docs/api/transforms/construct.rst b/docs/api/filter_factory/tiling.rst → docs/api/transforms/construct.rst
@@ -1,7 +1,7 @@
 :html_theme.sidebar_secondary.remove:
 
 **************************
-Tiling Functions
+Matrice Precomputes
 **************************
-.. automodule:: s2wav.filter_factory.tiling
+.. automodule:: s2wav.transforms.construct
    :members: