Support for Healpix polarisation

[Magwos]

I have been trying to reproduce the healpix implementation of the E/B modes computation using s2fft tools in this script, using the map provided by healpy for their tests available here: https://github.com/healpy/healpy/blob/main/test/data/wmap_band_iqumap_r9_7yr_V_v4_udgraded32.fits

import healpy as hp
import numpy as np
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)
from s2fft.sampling.s2_samples import flm_2d_to_hp, flm_hp_to_2d


import s2fft

test_map = hp.read_map("wmap_band_iqumap_r9_7yr_V_v4_udgraded32.fits", field=None)

nside = int(np.sqrt(test_map.shape[-1]/12))
lmax = 2*nside

# Preparing s2fft 
method = "jax"

# Running healpy
alms_hp = hp.map2alm(test_map, lmax=lmax, pol=True, iter=0)[1:]

# Retreiving the maps from which we extract spin 2 and spin -2 components
map_S_plus = test_map[1] # Q
map_S_minus = test_map[2] # U

# Following the euqations given by https://healpix.jpl.nasa.gov/html/subroutinesnode12.htm#sub:alm2map_spin we have:
# In (6) and (7) we identify the fields S_2 and S_-2
map_S_2 = map_S_plus + 1j*map_S_minus
map_S_m2 = map_S_plus - 1j*map_S_minus

# Getting their alm's with their respective spins
flm_S_2 = s2fft.transforms.spherical.forward(
      f=map_S_2,
      L=lmax+1,
      spin=2,
      nside=nside,
      sampling="healpix",
      method=method,
      precomps=None,
      iter=0
    )

flm_S_m2 = s2fft.transforms.spherical.forward(
      f=map_S_m2,
      L=lmax+1,
      spin=-2,
      nside=nside,
      sampling="healpix",
      method=method,
      precomps=None,
      iter=0
    )

# Then retrieving alm_+ and alm_- from (2) and (3), which correspond respetively to alm_E and alm_B

flm_plus = -(flm_S_2 + flm_S_m2)/2
flm_minus = 1j*(flm_S_2 - flm_S_m2)/2
    
# Getting the alm's to compare with the ones from healpy
flm_E_hp = flm_2d_to_hp(flm_plus, lmax+1)
flm_B_hp = flm_2d_to_hp(flm_minus, lmax+1)


# Comparing the results
diff = np.abs(flm_E_hp - alms_hp[0]) + np.abs(flm_B_hp - alms_hp[1])
print("Maximum difference between the alms", diff.max())

# Checking where the difference is the largest
alm_hp_E = flm_hp_to_2d(alms_hp[0], lmax+1)
alm_hp_B = flm_hp_to_2d(alms_hp[1], lmax+1)


diff = np.abs(alm_hp_E - flm_plus)
print("ells and m where difference is greater than 10^-10 for E", np.where(diff > 1e-10))

diff = np.abs(alm_hp_B - flm_minus)
print("ells and m where difference is greater than 10^-10 for B", np.where(diff > 1e-10))

Following the procedure described in https://healpix.sourceforge.io/html/sub_map2alm_spin.htm or https://healpix.jpl.nasa.gov/html/subroutinesnode12.htm#sub:alm2map_spin, I manage to retrieve E and B modes although with an maximum error of 10^-4, which is way higher than what it should (for comparison the test in https://github.com/astro-informatics/s2fft/blob/main/tests/test_spherical_transform.py#L192 is done with tolerance 10^-14).

It should as well be noted that this high error arises only for a few flms, with all the other being below the tolerance of 10^-10. Given their distribution of the flm with high error, I suspect a bug in the alm reconstruction but I could not pinpoint it in the code.

I compared the healpy result with the one given by ducc0 and they agree up to a tolerance of 10^-10 (I can send you the corresponding test code if needed).

Is there something in my script which I am doing wrong, or is there a possible bug somewhere?