.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/spherical_harmonic_transform.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_tutorials_spherical_harmonic_transform.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_tutorials_spherical_harmonic_transform.py:


Spherical harmonic transform
============================

This tutorial demonstrates how to use ``S2FFT`` to compute spherical harmonic transforms.

.. image:: https://colab.research.google.com/assets/colab-badge.svg
    :align: center
    :alt: Open in Google Colab
    :target: https://colab.research.google.com/github/astro-informatics/s2fft/tree/gh-pages/_colab_notebooks/spherical_harmonic_transform.ipynb

If you are working on this notebook in Google Colab; you will need to have Google Colab install ``cartopy``, ``s2fft`` and ``pyssht``, as well as fetch the data we'll be using.
You can do this by adding a cell to the top of the notebook with the following content:

.. code-block:: bash

    !pip install cartopy s2fft pyssht &> /dev/null
    !mkdir data/
    !wget https://github.com/astro-informatics/s2fft/raw/main/examples/data/Gaia_EDR3_flux.npy -P data/ &> /dev/null

and then running that cell.

.. GENERATED FROM PYTHON SOURCE LINES 25-27

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.

.. GENERATED FROM PYTHON SOURCE LINES 27-42

.. code-block:: Python


    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]








.. GENERATED FROM PYTHON SOURCE LINES 43-44

Let's look at the input signal:

.. GENERATED FROM PYTHON SOURCE LINES 44-51

.. code-block:: Python


    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()




.. image-sg:: /tutorials/images/sphx_glr_spherical_harmonic_transform_001.png
   :alt: spherical harmonic transform
   :srcset: /tutorials/images/sphx_glr_spherical_harmonic_transform_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 52-56

Computing the forward spherical harmonic transform
--------------------------------------------------

Let's now run the ``JAX`` function to compute the spherical harmonic transform of this map.

.. GENERATED FROM PYTHON SOURCE LINES 56-59

.. code-block:: Python


    flm = s2fft.forward_jax(f, L)








.. GENERATED FROM PYTHON SOURCE LINES 60-64

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.

.. GENERATED FROM PYTHON SOURCE LINES 64-68

.. code-block:: Python


    precomps = s2fft.generate_precomputes_jax(L, forward=True)
    flm_pre = s2fft.forward_jax(f, L, precomps=precomps)








.. GENERATED FROM PYTHON SOURCE LINES 69-73

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.

.. GENERATED FROM PYTHON SOURCE LINES 73-76

.. code-block:: Python


    f_recov = s2fft.inverse_jax(flm, L)








.. GENERATED FROM PYTHON SOURCE LINES 77-81

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.

.. GENERATED FROM PYTHON SOURCE LINES 81-85

.. code-block:: Python


    precomps = s2fft.generate_precomputes_jax(L, forward=False)
    f_recov_pre = s2fft.inverse_jax(flm_pre, L, precomps=precomps)








.. GENERATED FROM PYTHON SOURCE LINES 86-90

Computing the error
-------------------

Let's check the associated error, which should be close to machine precision for the sampling scheme used.

.. GENERATED FROM PYTHON SOURCE LINES 90-93

.. code-block:: Python


    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))}")




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Mean absolute error = 4.2080013578407947e-13
    Mean absolute error using precomputes = 4.2080013578407947e-13





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (4 minutes 10.629 seconds)


.. _sphx_glr_download_tutorials_spherical_harmonic_transform.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: spherical_harmonic_transform.ipynb <spherical_harmonic_transform.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: spherical_harmonic_transform.py <spherical_harmonic_transform.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: spherical_harmonic_transform.zip <spherical_harmonic_transform.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_