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?

[Magwos]

I have done the same exercise with alm2map, again from the equations described in the healpy documentation here: https://healpix.jpl.nasa.gov/html/subroutinesnode12.htm#sub:alm2map_spin

Here is the script to reproduce my implementation

import healpy as hp
import numpy as np
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from functools import partial

jax.config.update("jax_enable_x64", True)
from s2fft.sampling.reindex import flm_2d_to_hp_fast, flm_hp_to_2d_fast

from s2fft.transforms.spherical import forward, inverse

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)

# Retreiving the alms corresponding to E and B modes
L = lmax + 1
alms_input = jax.vmap(partial(flm_hp_to_2d_fast, L=L), in_axes=(0,))(alms_hp) # Transforming the ordering into the one accepted by s2fft 


# Following the euqations given by https://healpix.jpl.nasa.gov/html/subroutinesnode12.htm#sub:alm2map_spin, 
# we identify the spin 2 and spin -2 alms as in (2) and (3) 
alm_S_2 = -(alms_input[1,...]+1j*alms_input[2,...])
alm_S_m2 = -(alms_input[1,...]-1j*alms_input[2,...])

# Getting the corresponding maps with their respective spins
map_plus2 = inverse(
      flm=alm_S_2,
      L=lmax+1,
      spin=2,
      nside=nside,
      sampling="healpix",
      method=method,
      precomps=None,
    )

map_minus2 = inverse(
      flm=alm_S_m2,
      L=lmax+1,
      spin=-2,
      nside=nside,
      sampling="healpix",
      method=method,
      precomps=None,
    )

# Then retrieving map Q and U from (6) and (7), which correspond respetively to alm_E and alm_B
map_Q_s2fft = (map_plus2 + map_minus2) /2
map_U_s2fft = (map_plus2 - map_minus2) / (2*1j)
    
# Getting the output map of alm2map from healpy to get a comparison (we can as well compare with the input map)
map_hp = hp.alm2map(alms_hp, lmax=lmax, nside=nside, pol=True)[1:] # We only save the Q and U maps

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

# We acn then plot the difference between the two maps
plt.figure(figsize=(15,10))
for i, map_s2fft in enumerate([map_Q_s2fft, map_U_s2fft]):
  hp.mollview(np.abs(map_s2fft - map_hp[i]), sub=(2,1,i+1), title='QU'[i])
plt.show()

# Now, if we remove the pixels where the difference is above 10^-5

tol = 1e-6

diff = np.abs(map_Q_s2fft - map_hp[0])
print(f"Maximum difference for pixels with tolerance lower than {tol}, for the Q map: {diff[diff<tol].max()}")

diff = np.abs(map_U_s2fft - map_hp[1])
print(f"Maximum difference for pixels with tolerance lower than {tol}, for the U map is {diff[diff<tol].max()}")

In this case, we can clearly identify 9 rings on each map which are badly reconstructed and associated to an error above 10^-6.

Outside those 9 rings, the reconstructed maps from s2fft are completely consistent with the healpy ones, with an absolute tolerance of 10^-15.

It is very similar to the example described in the previous post: except for a few rings (above was for few ells), the alms/maps are very well reconstructed. We only have this important error for a few ells/rings.

[jasonmcewen]

Thanks for the messages @Magwos. We're pretty busy with a number of s2* related things at the moment.

I haven't taken a detailed read of your posts above but just took a quick skim.

The healpix accuracy will be very low if you don't have iterations turned on. We recently added support to s2fft to iterate to improve healpix accuracy (in this PR). I don't recall what the default is but it should be described in the documentation. You also might need to pick up a more recent version of s2fft if you have an older version installed before this suport was added.

Please let us know whether that resolves the issue? If not, do also let us know and we'll try to take a more detailed look soon.

[Magwos]

Thanks for your message @jasonmcewen, no worries!

The healpix accuracy will be very low if you don't have iterations turned on. We recently added support to s2fft to iterate to improve healpix accuracy (in #241). I don't recall what the default is but it should be described in the documentation. You also might need to pick up a more recent version of s2fft if you have an older version installed before this suport was added.

I just pulled the very latest s2fft version and retested with a greater number of iterations for the map2alm operations, and this does not help getting under the $10^{-10}$ tolerance (the tolerance is then at maximum $10^{-5}$, which is too high).

I retested as well the alm2map script I sent here (leading to the maps plotted above) with the latest s2fft github version, and it leads to the same issue. In particular, in this script I compare the result of inverse operations and alm2map, which do not need any iterative refinements.