gh-794: create spherical harmonic transform wrapper functions

CONTRIBUTING.md

@@ -19,7 +19,7 @@ The developer installation of _GLASS_ comes with several optional dependencies -
 These options can be used with `uv` in the following way
 
 ```bash
-uv sync --all-extras --group docs --group test
+uv sync --all-extras --all-groups
 ```
 
 ## Tooling

docs/modules/harmonics.rst

@@ -0,0 +1,26 @@
+.. module:: glass.harmonics
+
+:mod:`glass.harmonics` --- Spherical harmonics utilities
+========================================================
+
+.. currentmodule:: glass.harmonics
+
+This module contains utilities for working with spherical harmonics which are
+used by GLASS, but are otherwise unrelated to GLASS functionality.
+
+This module should be imported manually if used outside of GLASS::
+
+    import glass.harmonics
+
+
+General
+-------
+
+.. autofunction:: multalm
+
+
+Spherical harmonic transforms
+-----------------------------
+
+.. autofunction:: transform
+.. autofunction:: inverse_transform

docs/modules/index.rst

@@ -7,3 +7,4 @@ Modules
 
    algorithm
    grf
+   harmonics

glass/algorithm.py

@@ -111,8 +111,8 @@ def cov_clip(
     The relative tolerance *rtol* is defined as for
     :func:`~array_api.linalg.matrix_rank`.
 
-    Parameter
-    ---------
+    Parameters
+    ----------
     cov
         A symmetric matrix (or a stack of matrices).
     rtol

glass/fields.py

@@ -407,7 +407,9 @@ def _generate_grf(
 
         # transform alm to maps
         # can be performed in place on the temporary alm array
-        yield hp.alm2map(alm, nside, pixwin=False, pol=False, inplace=True)
+        yield glass.harmonics.inverse_transform(
+            alm, nside=nside, polarised_input=False, inplace=True
+        )
 
 
 @deprecated("use glass.generate() instead")

glass/harmonics.py

@@ -2,8 +2,11 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+import healpy as hp
+
 import array_api_compat
 
 if TYPE_CHECKING:
@@ -17,24 +20,22 @@ def multalm(
     """
     Multiply alm by bl.
 
-    The alm should be in GLASS order:
+    The alm should be in GLASS order::
 
-    [
-        00,
-        10, 11,
-        20, 21, 22,
-        30, 31, 32, 33,
-        ...
-    ]
+        [
+            00,
+            10, 11,
+            20, 21, 22,
+            30, 31, 32, 33,
+            ...
+        ]
 
     Parameters
     ----------
     alm
         The alm to multiply.
     bl
         The bl to multiply.
-    inplace
-        Whether to perform the operation in place.
 
     Returns
     -------
@@ -43,3 +44,94 @@ def multalm(
     """
     xp = array_api_compat.array_namespace(alm, bl, use_compat=False)
     return alm * xp.repeat(bl, xp.arange(bl.size) + 1)
+
+
+def transform(
+    maps: FloatArray,
+    *,
+    lmax: int | None = None,
+    polarised_input: bool = True,
+    use_pixel_weights: bool = False,
+) -> ComplexArray:
+    """
+    Computes the spherical harmonic transform.
+
+    Parameters
+    ----------
+    maps
+        The input map(s).
+    lmax
+        The maximum multipole to use.
+    polarised_input
+        Whether the input maps represent polarised data.
+    use_pixel_weights
+        Whether to use pixel weights in the transform.
+
+    Returns
+    -------
+        The spherical harmonic coefficients.
+
+    """
+    return hp.map2alm(
+        maps,
+        lmax=lmax,
+        pol=polarised_input,
+        use_pixel_weights=use_pixel_weights,
+    )
+
+
+def inverse_transform(  # noqa: PLR0913
+    alm: ComplexArray | Sequence[ComplexArray],
+    *,
+    nside: int,
+    inplace: bool = False,
+    lmax: int | None = None,
+    polarised_input: bool = True,
+    spin: int | None = None,
+) -> FloatArray | list[FloatArray]:
+    """
+    Computes the inverse spherical harmonic transform.
+
+    Parameters
+    ----------
+    alm
+        The spherical harmonic coefficients. If computing a spin-s map (spin is
+        not None), this should be a sequence of two coefficient arrays:
+        [alm, blm]. Otherwise, it is the single alm array.
+    nside
+        The nside of the output map if using HEALPix.
+    inplace
+        Whether to perform the operation in place (only applicable when
+        spin=None).
+    lmax
+        The maximum multipole to use.
+    polarised_input
+        Whether the input alm represents polarised data (only applicable when
+        spin=None).
+    spin
+        The spin s of the field being transformed.
+
+    Returns
+    -------
+        The output map(s).
+
+    """
+    if spin is None:
+        # alm is the single alm array or a sequence/array of 3 for polarised data
+        return hp.alm2map(
+            alm,
+            nside,
+            inplace=inplace,
+            lmax=lmax,
+            pol=polarised_input,
+        )
+
+    # check if alm is of the correct form for spin transforms
+    if not (hasattr(alm, "shape") and alm.shape[0] == 2) and not (
+        isinstance(alm, Sequence) and len(alm) == 2
+    ):
+        # alm must be a sequence [alm, blm]
+        msg = "for spin transforms, alm must be of the form [alm, blm]"
+        raise ValueError(msg)
+    # does not accept "inplace" or "polarised_input" arguments
+    return hp.alm2map_spin(alm, nside, spin, lmax)

glass/lensing.py

@@ -39,6 +39,7 @@
 import array_api_compat
 
 import glass._array_api_utils as _utils
+import glass.harmonics
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -279,7 +280,9 @@ def from_convergence(  # noqa: PLR0913
         lmax = 3 * nside - 1
 
     # compute alm
-    alm = hp.map2alm(kappa, lmax=lmax, pol=False, use_pixel_weights=True)
+    alm = glass.harmonics.transform(
+        kappa, lmax=lmax, polarised_input=False, use_pixel_weights=True
+    )
 
     # mode number; all conversions are factors of this
     ell = np.arange(lmax + 1)
@@ -293,7 +296,7 @@ def from_convergence(  # noqa: PLR0913
 
     # if potential is requested, compute map and add to output
     if potential:
-        psi = hp.alm2map(alm, nside, lmax=lmax)
+        psi = glass.harmonics.inverse_transform(alm, nside=nside, lmax=lmax)
         results += (psi,)
 
     # if no spin-weighted maps are requested, stop here
@@ -310,7 +313,9 @@ def from_convergence(  # noqa: PLR0913
 
     # if deflection is requested, compute spin-1 maps and add to output
     if deflection:
-        alpha = hp.alm2map_spin([alm, blm], nside, 1, lmax)
+        alpha = glass.harmonics.inverse_transform(
+            [alm, blm], nside=nside, spin=1, lmax=lmax
+        )
         alpha = alpha[0] + 1j * alpha[1]
         results += (alpha,)
 
@@ -328,7 +333,9 @@ def from_convergence(  # noqa: PLR0913
     hp.almxfl(alm, fl, inplace=True)
 
     # transform to shear maps
-    gamma = hp.alm2map_spin([alm, blm], nside, 2, lmax)
+    gamma = glass.harmonics.inverse_transform(
+        [alm, blm], nside=nside, spin=2, lmax=lmax
+    )
     gamma = gamma[0] + 1j * gamma[1]
     results += (gamma,)
 
@@ -370,7 +377,9 @@ def shear_from_convergence(
         lmax = 3 * nside - 1
 
     # compute alm
-    alm = hp.map2alm(kappa, lmax=lmax, pol=False, use_pixel_weights=True)
+    alm = glass.harmonics.transform(
+        kappa, lmax=lmax, polarised_input=False, use_pixel_weights=True
+    )
 
     # zero B-modes
     blm = np.zeros_like(alm)
@@ -390,7 +399,7 @@ def shear_from_convergence(
     hp.almxfl(alm, fl, inplace=True)
 
     # transform to shear maps
-    return hp.alm2map_spin([alm, blm], nside, 2, lmax)
+    return glass.harmonics.inverse_transform([alm, blm], nside=nside, spin=2, lmax=lmax)
 
 
 class MultiPlaneConvergence:

tests/core/test_harmonics.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING
 
+import healpy as hp
 import pytest
 
+import array_api_extra as xpx
+
 import glass.harmonics
 
 if TYPE_CHECKING:
@@ -54,3 +58,83 @@ def test_multalm(compare: type[Compare], xp: ModuleType) -> None:
 
     result = glass.harmonics.multalm(alm, bl)
     compare.assert_allclose(result, alm)
+
+
+def test_transform(compare: type[Compare], xp: ModuleType) -> None:
+    if xp.__name__ in {"array_api_strict", "jax.numpy"}:
+        pytest.skip("transform depends on healpy which is not Array API compatible")
+
+    # prepare inputs
+    constant = 1.23
+    lmax = 0
+    nside = 8
+
+    # create an unpolarised map where every pixel is the same
+    simple_map = xp.full(hp.nside2npix(nside), constant)
+
+    a00 = math.sqrt(4 * xp.pi) * constant
+    # the alm array for lmax=0 has size 1, so it only contains a_00
+    alm_expected = xp.asarray([a00], dtype=xp.complex128)
+
+    alm_result = glass.harmonics.transform(
+        simple_map,
+        lmax=lmax,
+        polarised_input=False,
+    )
+
+    compare.assert_allclose(alm_result, alm_expected, rtol=1e-15)
+
+
+def test_inverse_transform(compare: type[Compare], xp: ModuleType) -> None:
+    if xp.__name__ in {"array_api_strict", "jax.numpy"}:
+        pytest.skip("transform depends on healpy which is not Array API compatible")
+
+    # prepare inputs
+    lmax = 1
+    nside = 8
+    spin = 1
+
+    # create a simple monopole component ensuring a predictable output map value
+    alm_size = hp.Alm.getsize(lmax)
+    alm_monopole = xp.zeros(alm_size, dtype=xp.complex128)
+    xpx.at(alm_monopole, 0).set(math.sqrt(4 * math.pi))
+
+    # standard unpolarised transform without spin
+
+    map_result = glass.harmonics.inverse_transform(
+        alm_monopole,
+        lmax=lmax,
+        nside=nside,
+        polarised_input=False,
+    )
+
+    assert hasattr(map_result, "ndim")
+    assert map_result.ndim == 1
+    compare.assert_array_equal(map_result, 1.0)
+
+    # valid spin transform
+
+    alm_input_list = [alm_monopole, alm_monopole]
+
+    map_result_spin = glass.harmonics.inverse_transform(
+        alm_input_list,
+        lmax=lmax,
+        nside=nside,
+        spin=spin,
+    )
+
+    assert isinstance(map_result_spin, list)
+    assert len(map_result_spin) == 2
+    assert all(hasattr(m, "ndim") and m.ndim == 1 for m in map_result_spin)
+
+    # invalid spin input
+
+    with pytest.raises(
+        ValueError, match="for spin transforms, alm must be of the form"
+    ):
+        glass.harmonics.inverse_transform(
+            alm_monopole,
+            lmax=lmax,
+            nside=nside,
+            spin=spin,
+        )