gh-27: computation in harmonic space

docs/examples.rst

@@ -34,6 +34,7 @@ To get started, these examples focus on simulating one thing at a time.
    examples/1-basic/matter.ipynb
    examples/1-basic/density.ipynb
    examples/1-basic/lensing.ipynb
+   examples/1-basic/lensing-harmonic.ipynb
    examples/1-basic/photoz.ipynb
 
 

glass/__init__.py

@@ -37,6 +37,8 @@
     "getcl",
     "glass_to_healpix_spectra",
     "grf",
+    "harmonic_shear",
+    "harmonic_space",
     "healpix_to_glass_spectra",
     "iternorm",
     "linear_bias",
@@ -92,6 +94,7 @@
     generate_lognormal,
     getcl,
     glass_to_healpix_spectra,
+    harmonic_space,
     healpix_to_glass_spectra,
     iternorm,
     lognormal_fields,
@@ -112,6 +115,7 @@
     MultiPlaneConvergence,
     deflect,
     from_convergence,
+    harmonic_shear,
     multi_plane_matrix,
     multi_plane_weights,
     shear_from_convergence,

glass/fields.py

@@ -1109,3 +1109,29 @@ def regularized_spectra(
 
     # return regularised spectra from cov matrix array
     return [cov[:, i, j] for i, j in spectra_indices(cov.shape[-1])]
+
+
+def harmonic_space(
+    maps: Iterable[NDArray[Any]],
+    lmax: int | None = None,
+) -> Iterator[NDArray[Any]]:
+    """
+    Transform simulated fields to harmonic space.
+
+    This function takes an iterator over maps, such as the one returned by
+    :func:`generate`, and turns it into an iterator over their spherical
+    harmonic coefficients :math:`a_{lm}`.
+
+    Parameters
+    ----------
+    maps
+        Maps to be transformed to harmonic space.
+
+    Returns
+    -------
+    alms
+        Spherical harmonic decomposition of the maps.
+
+    """
+    for m in maps:
+        yield hp.map2alm(m, lmax=lmax, use_pixel_weights=True)

glass/lensing.py

@@ -31,6 +31,7 @@
 
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING, Literal, overload
 
 import healpy as hp
@@ -43,6 +44,7 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
     from types import ModuleType
+    from typing import Any
 
     from numpy.typing import NDArray
 
@@ -687,3 +689,88 @@ def deflect(
     d = xp.atan2(sa * sg, st * ca - ct * sa * cg)
 
     return lon - uxpx.degrees(d), uxpx.degrees(tp)
+
+
+def harmonic_shear(
+    kappa: NDArray[Any],
+    lon: NDArray[Any],
+    lat: NDArray[Any],
+) -> NDArray[Any]:
+    r"""
+    Compute shear at given positions in harmonic space.
+
+    Takes a convergence field in harmonic space (i.e. spherical harmonic
+    coefficients :math:`\kappa_{lm}`) and returns the *shear* values in the
+    given positions.
+
+    Parameters
+    ----------
+    kappa
+        Spherical harmonic coefficients of the convergence field.
+    lon, lat
+        Longitude and latitude (in degrees) of positions to evaluate.
+
+    Returns
+    -------
+    gamma
+        Complex shear in the given positions.
+
+    Notes
+    -----
+    The spherical harmonic coefficients of convergence and shear for
+    :math:`l \ge 2` are related as [Tessore23]_
+
+    .. math::
+
+        \gamma_{lm}
+        = -\sqrt{\frac{(l+2) \, (l-1)}{l \, (l+1)}} \, \kappa_{lm} \;,
+
+    with :math:`\gamma_{lm} = ` for :math:`l < 2`.
+
+    """
+    try:
+        import ducc0  # noqa: PLC0415
+    except ModuleNotFoundError as exc:
+        raise RuntimeError("function requires the ducc0 package") from exc
+
+    # should be moved to helper module
+    from glass.fields import _multalm  # noqa: PLC0415
+
+    # turn inputs into spherical coordinates
+    theta = (90.0 - lat.reshape(-1)) / 180 * np.pi
+    phi = (lon.reshape(-1) / 180 * np.pi) % (2 * np.pi)
+    if theta.size != phi.size:
+        raise ValueError("inconsistent input shapes")
+
+    # skip empty list of points
+    if theta.size == 0:
+        return np.zeros(0) + 1j * np.zeros(0)
+
+    # compute lmax from size of kappa alms
+    lmax = (math.isqrt(8 * kappa.shape[-1] + 1) - 3) // 2
+
+    # array for kappa to gamma conversion
+    ell = np.arange(2, lmax + 1)
+    fl = np.concat(
+        [
+            np.zeros(2),
+            -np.sqrt(((ell + 2) * (ell - 1)) / (ell * (ell + 1))),
+        ]
+    )
+
+    # compute inputs to transform
+    alm = np.stack([_multalm(kappa, fl), np.zeros_like(kappa)], axis=0)
+    loc = np.stack([theta, phi], axis=1)
+
+    # compute shear values in positions
+    gamma = ducc0.sht.synthesis_general(
+        alm=alm,
+        spin=2,
+        lmax=lmax,
+        loc=loc,
+        epsilon=2e-13,
+        nthreads=0,
+    )
+
+    # return shear as complex array
+    return np.dot(np.asarray([1, 1j]), gamma)  # type: ignore[no-any-return]