astro-informatics · willGraham01 · Mar 10, 2026 · Mar 10, 2026 · Mar 10, 2026 · Mar 10, 2026
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -38,7 +38,7 @@ jobs:
       run: |
         sudo apt install pandoc
         python -m pip install --upgrade pip
-        pip install .[docs]
+        pip install .[docs,examples]
 
     - name: Build documentation
       run: cd docs && make html

diff --git a/.gitignore b/.gitignore
@@ -36,3 +36,7 @@ install_manifest.txt
 compile_commands.json
 CTestTestfile.cmake
 _deps
+
+# Docs build
+docs/tutorials/
+docs/sg_execution_times.rst
diff --git a/README.md b/README.md
@@ -154,7 +154,7 @@ pytest
 
 ## Documentation 📖
 
-Documentation for the released version is available [here](https://astro-informatics.github.io/s2fft/).
+Documentation for the released version [is available here](https://astro-informatics.github.io/s2fft/).
 To install the documentation dependencies, clone the repository and install the package (in [editable mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html))
 with the extra documentation dependencies by running from the root of the repository
 
@@ -172,14 +172,13 @@ open _build/html/index.html
 
 ## Notebooks 📓
 
-A series of tutorial notebooks are included in the `notebooks` directory
-and rendered [in the documentation](https://astro-informatics.github.io/s2fft/tutorials/index.html).
+A series of tutorial notebooks have been rendered in, and are available for download from, [the documentation](https://astro-informatics.github.io/s2fft/tutorials/index.html).
 
 To install the dependencies required to run the notebooks locally, clone the repository and install the package (in [editable mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html))
-with the extra documentation and plotting dependencies by running from the root of the repository
+with the extra documentation, plotting, and examples dependencies by running from the root of the repository
 
 ```bash
-pip install -e ".[docs,plotting]"
+pip install -e ".[docs,examples,plotting]"
 ```
 
 To run the notebooks in Jupyter Lab, run from the root of the repository

diff --git a/docs/conf.py b/docs/conf.py
@@ -14,9 +14,12 @@
 import os
 import sys
 from importlib.metadata import version as get_version
+from pathlib import Path
 
 sys.path.insert(0, os.path.abspath(".."))
 
+DOCS_DIR = Path(__file__).parent
+
 
 # -- Project information -----------------------------------------------------
 
@@ -39,19 +42,18 @@
 # ones.
 extensions = [
     "sphinx_copybutton",
-    "nbsphinx_link",
     "sphinx.ext.autodoc",
     "sphinx.ext.napoleon",
     "sphinx.ext.mathjax",
     "sphinx.ext.githubpages",
     "sphinx_rtd_theme",
-    "nbsphinx",
     "IPython.sphinxext.ipython_console_highlighting",
     "sphinx_tabs.tabs",
     "sphinx_git",
     "sphinxcontrib.texfigure",
     "sphinx.ext.autosectionlabel",
     "sphinxemoji.sphinxemoji",
+    "sphinx_gallery.gen_gallery",
     "sphinx_mdinclude",
 ]
 
@@ -60,9 +62,18 @@
 napoleon_include_init_with_doc = True
 napoleon_numpy_docstring = False
 
+sphinx_gallery_conf = {
+    "examples_dirs": "../examples",
+    "gallery_dirs": "./tutorials",
+    "filename_pattern": "",
+    # For whatever reason, default_thumb_file is interpreted as
+    # relative to CWD in which build is run, unless an absolute path is provided.
+    "default_thumb_file": str(DOCS_DIR / "assets/sax_logo"),
+}
+
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]
-source_suffix = [".rst", ".ipynb"]
+source_suffix = [".rst"]
 
 # The master toctree document.
 master_doc = "index"

diff --git a/docs/tutorials/JAX_HEALPix/JAX_HEALPix_frontend.nblink b/docs/tutorials/JAX_HEALPix/JAX_HEALPix_frontend.nblink
diff --git a/docs/tutorials/JAX_SSHT/JAX_SSHT_frontend.nblink b/docs/tutorials/JAX_SSHT/JAX_SSHT_frontend.nblink
diff --git a/docs/tutorials/rotation/rotation.nblink b/docs/tutorials/rotation/rotation.nblink
diff --git a/docs/tutorials/spherical_harmonic/spherical_harmonic_transform.nblink b/docs/tutorials/spherical_harmonic/spherical_harmonic_transform.nblink
diff --git a/docs/tutorials/torch_frontend/torch_frontend.nblink b/docs/tutorials/torch_frontend/torch_frontend.nblink
diff --git a/docs/tutorials/wigner/wigner_transform.nblink b/docs/tutorials/wigner/wigner_transform.nblink
diff --git a/docs/tutorials/index.rst → examples/GALLERY_HEADER.rst b/docs/tutorials/index.rst → examples/GALLERY_HEADER.rst
@@ -1,17 +1,14 @@
-:html_theme.sidebar_secondary.remove:
-
-*****************************
 Tutorials
-*****************************
+=========
 
-This section contains a series of tutorial notebooks which go through some of the
+This section contains :ref:`a series of tutorial notebooks <tutorial-notebooks-label>` which go through some of the
 key features of the ``S2FFT`` package.
 
 At a high-level the ``S2FFT`` package is structured such that the two primary transforms,
 the Wigner and spherical harmonic transforms, can easily be accessed.
 
 Core usage |:rocket:|
------------------------------
+---------------------
 
 To import and use ``S2FFT`` is as simple follows: 
 
@@ -33,14 +30,11 @@ To import and use ``S2FFT`` is as simple follows:
 |   f = s2fft.inverse(flm, L, method="jax")             |   f = s2fft.wigner.inverse(flmn, L, N, method="jax")       |
 +-------------------------------------------------------+------------------------------------------------------------+
 
-.. toctree::
-   :hidden:
-   :maxdepth: 3
-   :caption: Tutorials
-
-   spherical_harmonic/spherical_harmonic_transform.nblink
-   wigner/wigner_transform.nblink
-   rotation/rotation.nblink
-   torch_frontend/torch_frontend.nblink
-   JAX_SSHT/JAX_SSHT_frontend.nblink
-   JAX_HEALPix/JAX_HEALPix_frontend.nblink
+.. _tutorial-notebooks-label:
+
+Tutorial notebooks
+------------------
+
+Below are a few short tutorials that cover how to use specific features of ``S2FFT``.
+
+We also have a notebook demonstrating how to use CUDA-accelerated HEALPix spherical harmonic transforms in ``S2FFT``, which `is accessible in notebook format here <https://github.com/astro-informatics/s2fft/blob/main/notebooks/JAX_CUDA_HEALPix.ipynb>`_, or alternatively can be `opened in Google Colab <https://colab.research.google.com/github/astro-informatics/s2fft/blob/main/notebooks/JAX_CUDA_HEALPix.ipynb>`_.
diff --git a/examples/JAX_HEALPix_frontend.py b/examples/JAX_HEALPix_frontend.py
@@ -0,0 +1,76 @@
+"""
+JAX HEALPix Frontend
+====================
+
+This short tutorial demonstrates how to use the custom ``JAX`` frontend support ``S2FFT`` provides for the `HEALPix <https://healpix.jpl.nasa.gov>`_ C++ library.
+"""
+
+# %%
+# ``S2FFT``'s support for the `HEALPix <https://healpix.jpl.nasa.gov>`_ C++ library resolves issues involving long JIT compile times for HEALPix when running on CPU.
+# As with the other introductions, let's import some packages and define an arbitrary bandlimited signal to work with.
+
+import jax
+import numpy as np
+
+import s2fft
+
+jax.config.update("jax_enable_x64", True)
+
+L = 128
+nside = 64
+method = "jax_healpy"
+sampling = "healpix"
+rng = np.random.default_rng(23457801234570)
+flm = s2fft.utils.signal_generator.generate_flm(rng, L)
+f = s2fft.inverse(flm, L, nside=nside, sampling=sampling, method=method)
+
+# %%
+# Calling forward HEALPix C++ function from JAX
+# ---------------------------------------------
+
+flm = s2fft.forward(f, L, nside=nside, sampling=sampling, method=method)
+
+# %%
+# Calling inverse HEALPix C++ function from JAX
+# ---------------------------------------------
+
+f_recov = s2fft.inverse(flm, L, nside=nside, sampling=sampling, method=method)
+
+# %%
+# Computing the roundtrip error
+# -----------------------------
+#
+# Let's check the associated error, which should be around ``1e-5`` for healpix, which is not an exact sampling of the sphere.
+# Note that increasing ``iters`` will reduce the numerical error here slightly, at the cost of linearly increased compute.
+
+print(f"Mean absolute error = {np.nanmean(np.abs(f_recov - f))}")
+
+# %%
+# Differentiating through HEALPix C++ functions
+# ---------------------------------------------
+#
+# So far all this is doing is providing an interface between ``JAX`` and ``HEALPix``, the real novelty comes when we differentiate through the C++ library.
+
+
+# Define an arbitrary JAX function
+def differentiable_test(flm) -> int:
+    f = s2fft.inverse(flm, L, nside=nside, sampling=sampling, method=method)
+    return jax.numpy.nanmean(jax.numpy.abs(f) ** 2)
+
+
+# Create the JAX reverse mode gradient function
+gradient_func = jax.grad(differentiable_test)
+
+# Compute the gradient automatically
+gradient = gradient_func(flm)
+
+# %%
+# Validating these gradients
+# --------------------------
+#
+# This is all well and good, but how do we know these gradients are correct?
+# Thankfully ``JAX`` provides a simple function to check this...
+
+from jax.test_util import check_grads
+
+check_grads(differentiable_test, (flm,), order=1, modes=("rev"))
diff --git a/examples/JAX_SSHT_frontend.py b/examples/JAX_SSHT_frontend.py
@@ -0,0 +1,72 @@
+"""
+JAX SSHT frontend
+=================
+
+This short tutorial demonstrates how to use the custom ``JAX`` frontend support ``S2FFT`` provides for the `SSHT <https://github.com/astro-informatics/ssht>`_ C library.
+"""
+
+# %%
+# As with the other introductions, let's import some packages and define an arbitrary bandlimited signal to work with.
+
+import jax
+import numpy as np
+
+import s2fft
+
+jax.config.update("jax_enable_x64", True)
+
+L = 128
+method = "jax_ssht"
+rng = np.random.default_rng(23457801234570)
+flm = s2fft.utils.signal_generator.generate_flm(rng, L)
+f = s2fft.inverse(flm, L, method=method)
+
+# %%
+# Calling forward SSHT C function from JAX
+# ----------------------------------------
+
+flm = s2fft.forward(f, L, method=method)
+
+# %%
+# Calling inverse SSHT C function from JAX
+# ----------------------------------------
+
+f_recov = s2fft.inverse(flm, L, method=method)
+
+# %%
+# Computing the roundtrip error
+# -----------------------------
+#
+# Let's check the associated error, which should be close to machine precision for the sampling scheme used.
+
+print(f"Mean absolute error = {np.nanmean(np.abs(f_recov - f))}")
+
+# %%
+# Differentiating through SSHT C functions
+# ----------------------------------------
+#
+# So far all this is doing is providing an interface between ``JAX`` and ``SSHT``, the real novelty comes when we differentiate through the C library.
+
+
+# Define an arbitrary JAX function
+def differentiable_test(flm) -> int:
+    f = s2fft.inverse(flm, L, method=method)
+    return jax.numpy.nanmean(jax.numpy.abs(f) ** 2)
+
+
+# Create the JAX reverse mode gradient function
+gradient_func = jax.grad(differentiable_test)
+
+# Compute the gradient automatically
+gradient = gradient_func(flm)
+
+# %%
+# Validating these gradients
+# --------------------------
+#
+# This is all well and good, but how do we know these gradients are correct?
+# Thankfully ``JAX`` provides a simple function to check this...
+
+from jax.test_util import check_grads
+
+check_grads(differentiable_test, (flm,), order=1, modes=("rev"))
diff --git a/notebooks/data/Gaia_EDR3_flux.npy → examples/data/Gaia_EDR3_flux.npy b/notebooks/data/Gaia_EDR3_flux.npy → examples/data/Gaia_EDR3_flux.npy
diff --git a/examples/spherical_harmonic_transform.py b/examples/spherical_harmonic_transform.py
@@ -0,0 +1,76 @@
+"""
+Spherical harmonic transform
+============================
+
+This tutorial demonstrates how to use ``S2FFT`` to compute spherical harmonic transforms.
+"""
+
+# %%
+# In this example we will adopt the sampling scheme of `McEwen & Wiaux (2012) <https://arxiv.org/abs/1110.6298>`_.
+# First let's load an input signal that is sampled on the sphere with this sampling scheme.
+
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import cartopy.crs as ccrs
+import numpy as np
+from matplotlib import pyplot as plt
+
+import s2fft
+
+sampling = "mw"
+f = np.load("data/Gaia_EDR3_flux.npy")
+L = f.shape[0]
+
+# %%
+# Let's look at the input signal:
+
+plt.figure(figsize=(10, 5))
+ax = plt.axes(projection=ccrs.Mollweide())
+im = ax.imshow(f, transform=ccrs.PlateCarree(), cmap="magma")
+plt.axis("off")
+plt.show()
+
+# %%
+# Computing the forward spherical harmonic transform
+# --------------------------------------------------
+#
+# Let's now run the ``JAX`` function to compute the spherical harmonic transform of this map.
+
+flm = s2fft.forward_jax(f, L)
+
+# %%
+# If you are planning on applying this transform many times (e.g. during training of a model) we recommend precomputing and storing some small arrays that are used every time.
+# This trades off additional memory usage for enhanced speed and should be fine at small and moderate bandlimits ``L``.
+#
+# To do this simply compute these and pass as a static argument.
+
+precomps = s2fft.generate_precomputes_jax(L, forward=True)
+flm_pre = s2fft.forward_jax(f, L, precomps=precomps)
+
+# %%
+# Computing the inverse spherical harmonic transform
+# --------------------------------------------------
+#
+# Let's run the ``JAX`` function to compute the inverse spherical harmonic transform to get back to the input map.
+
+f_recov = s2fft.inverse_jax(flm, L)
+
+# %%
+# Again, if you are planning on applying this transform many times we recommend precomputing and storing some small arrays that are used every time.
+# Recall, this trades off additional memory usage for enhanced speed and should be fine at small and moderate bandlimits ``L``.
+#
+# To do this simply compute these and pass as a static argument.
+
+precomps = s2fft.generate_precomputes_jax(L, forward=False)
+f_recov_pre = s2fft.inverse_jax(flm_pre, L, precomps=precomps)
+
+# %%
+# Computing the error
+# -------------------
+#
+# Let's check the associated error, which should be close to machine precision for the sampling scheme used.
+
+print(f"Mean absolute error = {np.nanmean(np.abs(f_recov - f))}")
+print(f"Mean absolute error using precomputes = {np.nanmean(np.abs(f_recov_pre - f))}")