{"id":13478039,"url":"https://github.com/astro-informatics/s2ball","last_synced_at":"2025-09-12T02:10:00.245Z","repository":{"id":220936209,"uuid":"587688640","full_name":"astro-informatics/s2ball","owner":"astro-informatics","description":"Differentiable and accelerated wavelet transform on the ball with JAX","archived":false,"fork":false,"pushed_at":"2024-02-05T09:03:20.000Z","size":4820,"stargazers_count":6,"open_issues_count":0,"forks_count":0,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-09-09T11:52:58.011Z","etag":null,"topics":["ball","differentiable-programming","harmonic-analysis","jax","wavelet-transform"],"latest_commit_sha":null,"homepage":"https://astro-informatics.github.io/s2ball/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/astro-informatics.png","metadata":{"files":{"readme":"README.rst","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null}},"created_at":"2023-01-11T10:51:24.000Z","updated_at":"2024-12-06T21:22:39.000Z","dependencies_parsed_at":"2024-02-05T09:50:38.515Z","dependency_job_id":"6c54e451-d311-4cb0-8b80-80f57ad97332","html_url":"https://github.com/astro-informatics/s2ball","commit_stats":null,"previous_names":["astro-informatics/s2ball"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/astro-informatics/s2ball","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/astro-informatics%2Fs2ball","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/astro-informatics%2Fs2ball/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/astro-informatics%2Fs2ball/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/astro-informatics%2Fs2ball/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/astro-informatics","download_url":"https://codeload.github.com/astro-informatics/s2ball/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/astro-informatics%2Fs2ball/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":274742855,"owners_count":25341132,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-12T02:00:09.324Z","response_time":60,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["ball","differentiable-programming","harmonic-analysis","jax","wavelet-transform"],"created_at":"2024-07-31T16:01:51.611Z","updated_at":"2025-09-12T02:10:00.221Z","avatar_url":"https://github.com/astro-informatics.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":".. image:: https://github.com/astro-informatics/s2ball/actions/workflows/tests.yml/badge.svg?branch=main\n    :target: https://github.com/astro-informatics/s2ball/actions/workflows/tests.yml\n.. image:: https://codecov.io/gh/astro-informatics/s2ball/graph/badge.svg?token=12JWTZPO96\n    :target: https://codecov.io/gh/astro-informatics/s2ball\n.. image:: https://img.shields.io/badge/License-MIT-yellow.svg\n    :target: https://opensource.org/licenses/MIT\n.. image:: https://badge.fury.io/py/s2ball.svg\n    :target: https://badge.fury.io/py/s2ball\n.. image:: http://img.shields.io/badge/arXiv-2402.01282-orange.svg?style=flat\n    :target: https://arxiv.org/abs/2402.01282\n.. image:: https://img.shields.io/badge/all_contributors-2-orange.svg?style=flat-square\n    :target: #contributors-\n.. image:: https://img.shields.io/badge/code%20style-black-000000.svg\n    :target: https://github.com/psf/black\n\n===================================================================\n|logo| Differentiable and accelerated wavelet transform on the ball\n===================================================================\n\n.. |logo| raw:: html\n\n   \u003cimg src=\"./docs/assets/sax_logo.png\" align=\"left\" height=\"100\" width=\"100\"\u003e\n\n`S2BALL` is a JAX package which provides support for automatically \ndifferentiable and greatly accelerated generalized Fourier transforms on the sphere, \nrotation group, and ball. Furthermore, we support extremely efficient and differentiable \nWigner-Laguerre wavelets which can be up to 22,000 times faster \nthan their ``C`` counterparts, discussed in related works by `Leistedt \u0026 McEwen 2012 \n\u003chttps://arxiv.org/pdf/1205.0792.pdf\u003e`_ and `Price \u0026 McEwen 2021 \u003chttps://arxiv.org/pdf/2105.05518.pdf\u003e`_.\n\nThe transforms `S2BALL` provides are optimally fast but come with a substantial memory \noverhead and cannot be used above a harmonic bandlimit of L ~ 256, at least with current GPU memory \nlimitations. That being said, many applications are more than comfortable at these resolutions, for \nwhich these JAX transforms are ideally suited, *e.g.* geophysical modelling, diffusion \ntensor imaging, multiscale molecular modelling. For those with machine learning in mind, \nit should be explicitly noted that these transforms are indeed equivariant to their respective groups.\n\nWavelet Filters on the Ball :ringed_planet:\n--------------------------------------------\n\nThe filters `S2BALL` provides were originally derived by `Leistedt \u0026 McEwen 2012 \n\u003chttps://arxiv.org/pdf/1205.0792.pdf\u003e`_ and are constructed by tesselating both harmonic space \nand the radial half-line with infinitely differentiable Cauchy-Schwartz functions. This tesselation \ngives rise to the follow frequency space localisation \n\n|filter_support|\n\n.. |filter_support| image:: ./docs/assets/figures/ball_filter_support.png\n    :width: 90%\n\nThe above is the frequency space localisation of the wavelet filters, however one can also view wavelet filters in pixel space. Visualising these filters is somewhat tricky as the ball is a 3-dimensional surface embedded in 4-dimensional space. We can, however, straightforwardly view a spherical slice of the ball for each radial node\n\n|filter_support_pixel|\n\n.. |filter_support_pixel| image:: ./docs/assets/figures/ball_filter_support_pixelspace.png\n    :width: 90%\n\nInstallation :computer:\n---------------------------\n\nThe Python dependencies for the `S2BALL` package are listed in the file\n`requirements/requirements-core.txt` and will be automatically installed\ninto the active python environment by `pip \u003chttps://pypi.org\u003e`_ when running\n\n.. code-block:: bash\n\n    pip install .        \n\n\nfrom the root directory of the repository. Unit tests can then be\nexecuted to ensure the installation was successful by running\n\n.. code-block:: bash\n\n    pytest tests/  \n\nAlternatively, the `S2BALL` package may be installed directly from PyPi by running \n\n.. code-block:: bash\n\n    pip install s2ball\n\nBasic Usage :rocket:\n---------------------------\n\nTo apply the generalised Fourier transforms supported by `S2BALL` one need only \nimport the package and apply the respective transform, which is as simple as doing the \nfollowing: \n\n.. code-block:: Python\n\n    from s2ball.transform import *\n    import numpy as np \n\n    # Load some data\n    f = np.load(\"path_to_your_data.npy\")\n\n+-------------------------------------------------------+------------------------------------------------------------+\n|and for data on the sphere with shape :math:`[L, 2L-1]`|or data on SO(3) with shape :math:`[2N-1, L, 2L-1]`         |\n|                                                       |                                                            |\n|.. code-block:: Python                                 |.. code-block:: Python                                      |\n|                                                       |                                                            |\n|   L = L                                               |   L = L; N = N                                             |\n|                                                       |                                                            |\n|   # Compute harmonic coefficients                     |   # Compute Wigner coefficients                            |\n|   flm = harmonic.forward(f, L)                        |   flmn = wigner.forward(f, L, N)                           |\n|                                                       |                                                            |\n|   # Sythensise signal f                               |   # Sythensise signal f                                    |\n|   f = harmonic.inverse(flm, L)                        |   f = wigner.inverse(flmn, L, N)                           |\n+-------------------------------------------------------+------------------------------------------------------------+\n\n+---------------------------------------------------+---------------------------------------------------------+\n|or data on the ball with shape :math:`[P, L, 2L-1]`|or with shape :math:`[P, 2N-1, L, 2L-1]`                 |\n|                                                   |                                                         |\n|.. code-block:: Python                             |.. code-block:: Python                                   |\n|                                                   |                                                         |\n|   L = L; P = P                                    |   L = L; N = N; P = P                                   |\n|                                                   |                                                         |\n|   # Compute spherical-Laguerre coefficients       |   # Compute Wigner coefficients                         |\n|   flmp = laguerre.forward(f, L, P)                |   flmnp = wigner_laguerre.forward(f, L, N, P)           |\n|                                                   |                                                         |\n|   # Sythensise signal f                           |   # Sythensise signal f                                 |\n|   f = laguerre.inverse(flmp, L, P)                |   f = wigner_laguerre.inverse(flmnp, L, N, P)           |\n+---------------------------------------------------+---------------------------------------------------------+\n\nHowever, for repeated application of these transforms it is optimal to instead precompile \nvarious kernels which can be placed on device to minimise i/o during *e.g.* training. This \noperational mode can be seen throughout our examples, found `here \n\u003chttps://github.com/astro-informatics/s2ball/tree/main/notebooks\u003e`_.\n\nComputational Benchmarking :hourglass_flowing_sand:\n--------------------------------------------------\n\nThe various generalized Fourier and wavelet transforms supported by `S2BALL` were \nbenchmarked against their ``C`` counterparts over a variety of parameter configurations. \nEach benchmark has been averaged over many runs, though here we provide only the mean. \nAll CPU based operations were executed on a single core from a AMD EPYC 7702 64-core \nprocessor, whereas all JAX operations were executed on a single NVIDIA A100 graphics \nprocessing unit. The Jupyter notebooks for each benchmark can be found `here \n\u003chttps://github.com/astro-informatics/s2ball/tree/main/notebooks\u003e`_.\n\nNote that benchmarking is restricted to scalar (spin 0 ) fields, though spin is supported \nthroughout `S2BALL`. Further note that for Wigner tests we set N=5, and in our \nLaguerre and wavelet benchmarking we set N=1, as FLAG/FLAGLET otherwise take \nexcessive compute. Finally, `S2BALL` transforms trivially support batching and \nso can, in many cases, gain several more orders of magnitude acceleration.\n    \n|harmonic| |wigner| \n\n|laguerre| |wavelet|\n\n.. |harmonic| image:: ./docs/assets/figures/harmonic.png\n    :width: 48%\n\n.. |wigner| image:: ./docs/assets/figures/wigner.png\n    :width: 48%\n\n.. |laguerre| image:: ./docs/assets/figures/laguerre.png\n    :width: 48%\n\n.. |wavelet| image:: ./docs/assets/figures/wavelet.png\n    :width: 48%\n\nContributors ✨\n-----------------------------------\n\nThanks goes to these wonderful people (`emoji\nkey \u003chttps://allcontributors.org/docs/en/emoji-key\u003e`_):\n\n.. raw:: html \n\n    \u003cembed\u003e\n        \u003ctable\u003e\n        \u003ctbody\u003e\n            \u003ctr\u003e\n            \u003ctd align=\"center\" valign=\"top\" width=\"16.66%\"\u003e\u003ca href=\"https://cosmomatt.github.io\"\u003e\u003cimg src=\"https://avatars.githubusercontent.com/u/32554533?v=4?s=100\" width=\"100px;\" alt=\"Matt Price\"/\u003e\u003cbr /\u003e\u003csub\u003e\u003cb\u003eMatt Price\u003c/b\u003e\u003c/sub\u003e\u003c/a\u003e\u003cbr /\u003e\u003ca href=\"https://github.com/astro-informatics/S2BALL/commits?author=CosmoMatt\" title=\"Code\"\u003e💻\u003c/a\u003e \u003ca href=\"https://github.com/astro-informatics/S2BALL/pulls?q=is%3Apr+reviewed-by%3ACosmoMatt\" title=\"Reviewed Pull Requests\"\u003e👀\u003c/a\u003e \u003ca href=\"#ideas-CosmoMatt\" title=\"Ideas, Planning, \u0026 Feedback\"\u003e🤔\u003c/a\u003e\u003c/td\u003e\n            \u003ctd align=\"center\" valign=\"top\" width=\"16.66%\"\u003e\u003ca href=\"http://www.jasonmcewen.org\"\u003e\u003cimg src=\"https://avatars.githubusercontent.com/u/3181701?v=4?s=100\" width=\"100px;\" alt=\"Jason McEwen \"/\u003e\u003cbr /\u003e\u003csub\u003e\u003cb\u003eJason McEwen \u003c/b\u003e\u003c/sub\u003e\u003c/a\u003e\u003cbr /\u003e\u003ca href=\"https://github.com/astro-informatics/S2BALL/commits?author=jasonmcewen\" title=\"Code\"\u003e💻\u003c/a\u003e \u003ca href=\"https://github.com/astro-informatics/S2BALL/pulls?q=is%3Apr+reviewed-by%3Ajasonmcewen\" title=\"Reviewed Pull Requests\"\u003e👀\u003c/a\u003e \u003ca href=\"#ideas-jasonmcewen\" title=\"Ideas, Planning, \u0026 Feedback\"\u003e🤔\u003c/a\u003e\u003c/td\u003e\n            \u003c/tr\u003e\n        \u003c/tbody\u003e\n        \u003c/table\u003e\n    \u003c/embed\u003e\n\nAttribution :books:\n-------------------\n\nShould this code be used in any way, we kindly request that the following article is\nreferenced. A BibTeX entry for this reference may look like:\n\n.. code-block:: \n\n    @article{price:s2ball, \n        author      = \"Matthew A. Price and Alicja Polanska and Jessica Whitney and Jason D. McEwen\",\n        title       = \"Differentiable and accelerated directional wavelet transform on the sphere and ball\",\n        year        = \"2024\",\n        eprint      = \"arXiv:2402.01282\"        \n    }\n\nThis work is provided as part of a collection of `JAX` harmonic analysis packages which include \n\n.. code-block:: \n\n    @article{price:s2fft, \n        author      = \"Matthew A. Price and Jason D. McEwen\",\n        title       = \"Differentiable and accelerated spherical harmonic and Wigner transforms\",\n        journal     = \"Journal of Computational Physics, submitted\",\n        year        = \"2023\",\n        eprint      = \"arXiv:2311.14670\"        \n    }\n    \nYou might also like to consider citing our related papers on which this code builds:\n\n.. code-block:: \n\n    @article{leistedt:flaglets,\n        author      = \"Boris Leistedt and Jason D. McEwen\",\n        title       = \"Exact wavelets on the ball\",\n        journal     = \"IEEE Trans. Sig. Proc.\",\n        year        = \"2012\",\n        volume      = \"60\",\n        number      = \"12\",\n        pages       = \"6257--6269\",        \n        eprint      = \"arXiv:1205.0792\",\n        doi         = \"110.1109/TSP.2012.2215030\"\n    }\n\n.. code-block:: \n\n    @article{mcewen:fssht,\n        author      = \"Jason D. McEwen and Yves Wiaux\",\n        title       = \"A novel sampling theorem on the sphere\",\n        journal     = \"IEEE Trans. Sig. Proc.\",\n        year        = \"2011\",\n        volume      = \"59\",\n        number      = \"12\",\n        pages       = \"5876--5887\",        \n        eprint      = \"arXiv:1110.6298\",\n        doi         = \"10.1109/TSP.2011.2166394\"\n    }\n\n.. code-block:: \n\n    @article{mcewen:so3,\n        author      = \"Jason D. McEwen and Martin B{\\\"u}ttner and Boris ~Leistedt and Hiranya V. Peiris and Yves Wiaux\",\n        title       = \"A novel sampling theorem on the rotation group\",\n        journal     = \"IEEE Sig. Proc. Let.\",\n        year        = \"2015\",\n        volume      = \"22\",\n        number      = \"12\",\n        pages       = \"2425--2429\",\n        eprint      = \"arXiv:1508.03101\",\n        doi         = \"10.1109/LSP.2015.2490676\"    \n    }\n\nLicense :memo:\n---------------\n\n`S2BALL` is released under the MIT license (see `LICENSE.txt \u003chttps://github.com/astro-informatics/s2ball/blob/main/LICENCE.txt\u003e`_).\n\n.. code-block::\n\n    We provide this code under an MIT open-source licence with the hope that it will be of use to a wider community.\n\n    Copyright 2024 Matthew Price, Jason McEwen and contributors.\n\n    `S2BALL` is free software made available under the MIT License. For details see the LICENSE file.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fastro-informatics%2Fs2ball","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fastro-informatics%2Fs2ball","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fastro-informatics%2Fs2ball/lists"}