Speaker
Magdy Morshed
(INFN Ferrara)
Description
The Healpy package is widely used in cosmology, both for the handling of the HEALPix pixelization and the associated spherical harmonic transforms. The current Healpy package implementation is not JAX friendly, as the spherical harmonics implementation is based in C++, which prevents us to possibly accelerate the associated routines using JIT or to run them on GPUs. In this talk, we will present JAX-Healpy, a package aimed at proposing Healpy routines implemented using JAX. In particular, we will review the current status of the codes, with the spherical harmonics implementation wrapping the JAX package S2FFT.
Primary authors
Magdy Morshed
(INFN Ferrara)
Pierre Chanial
(APC)
