C/C++ custom JAX support#

s2fft.transforms.c_backend_spherical.healpy_alm2map(flm: Array, L: int, nside: int) → Array#

JAX wrapper for healpy alm2map function (inverse spherical harmonic transform).

This wrapper assumes the passed spherical harmonic coefficients are a one dimensional array using HEALPix (ring-ordered) indexing. To instead pass a two-dimensional array of harmonic coefficients use healpy_inverse().

Parameters:

  • flm (jnp.ndarray) – Spherical harmonic coefficients.

  • L (int) – Harmonic band-limit. Equivalent to lmax + 1 in healpy.

  • nside (int) – HEALPix Nside resolution parameter.

Returns:

Signal on the sphere.

Return type:

jnp.ndarray

s2fft.transforms.c_backend_spherical.healpy_forward(f: Array, L: int, nside: int, iter: int = 3) → Array#

Compute the forward scalar spherical harmonic transform (healpy JAX).

HEALPix is a C++ library which implements the scalar spherical harmonic transform outlined in [1]. We make use of their healpy python bindings for which we provide custom JAX frontends, hence providing support for automatic differentiation. Currently these transforms can only be deployed on CPU, which is a limitation of the C++ library.

Parameters:

  • f (jnp.ndarray) – Signal on the sphere.

  • L (int) – Harmonic band-limit.

  • nside (int) – HEALPix Nside resolution parameter.

  • iter (int, optional) – Number of subiterations (iterative refinement steps) for healpy. Note that iterations increase the precision of the forward transform as an inverse of inverse transform, but with a linear increase in computational cost. Between 2 and 3 iterations is a good compromise.

Returns:

Harmonic coefficients of signal f.

Return type:

jnp.ndarray

Note

[1] Gorski, Krzysztof M., et al. “HEALPix: A framework for high-resolution discretization and fast analysis of data distributed on the sphere.” The Astrophysical Journal 622.2 (2005): 759

s2fft.transforms.c_backend_spherical.healpy_inverse(flm: Array, L: int, nside: int) → Array#

Compute the inverse scalar real spherical harmonic transform (HEALPix JAX).

HEALPix is a C++ library which implements the scalar spherical harmonic transform outlined in [1]. We make use of their healpy python bindings for which we provide custom JAX frontends, hence providing support for automatic differentiation. Currently these transforms can only be deployed on CPU, which is a limitation of the C++ library.

Parameters:

  • flm (jnp.ndarray) – Spherical harmonic coefficients.

  • L (int) – Harmonic band-limit.

  • nside (int) – HEALPix Nside resolution parameter.

Returns:

Signal on the sphere.

Return type:

jnp.ndarray

Note

[1] Gorski, Krzysztof M., et al. “HEALPix: A framework for high-resolution discretization and fast analysis of data distributed on the sphere.” The Astrophysical Journal 622.2 (2005): 759

s2fft.transforms.c_backend_spherical.healpy_map2alm(f: Array, L: int, nside: int) → Array#

JAX wrapper for healpy map2alm function (forward spherical harmonic transform).

This wrapper will return the spherical harmonic coefficients as a one dimensional array using HEALPix (ring-ordered) indexing. To instead return a two-dimensional array of harmonic coefficients use healpy_forward().

Parameters:

  • f (jnp.ndarray) – Signal on the sphere.

  • L (int) – Harmonic band-limit. Equivalent to lmax + 1 in healpy.

  • nside (int) – HEALPix Nside resolution parameter.

Returns:

Harmonic coefficients of signal f.

Return type:

jnp.ndarray