Significant round trip error when using the method = "jax" and sampling = "healpix"

[sijil-jose]

Hi , I was checking the round trip errors in a similar manner to that of the tutorial notebooks with method =" jax" and sampling= 'healpix.' I noticed that the mean absolute error was around 0.02. I have noticed a similar range while using method =" numpy" with sampling = "healpix". Is this expected? I would like to use the method = ''jax" to use vmap and jit. Is there a work around for this?

Thanks
Sijil

L = 128
nside = 64
method = "jax"
sampling = "healpix"
rng = np.random.default_rng(23457801234570)
flm = s2fft.utils.signal_generator.generate_flm(rng, L,)

f = s2fft.inverse(flm, L, nside=nside, sampling = sampling, method=method)

flm = s2fft.forward(f, L, nside=nside, sampling=sampling, method=method)
f_recov = s2fft.inverse(flm, L, nside=nside, sampling=sampling,  method=method)

print(f"Mean absolute error = {np.nanmean(np.abs(f_recov - f))}")

Mean absolute error = 0.019590334447510353

[matt-graham]

Hi @sijil-jose. The HEALPix sampling scheme does not have an associated sampling theorem, and so the forward transform is not an exact inverse of the backward transform, which is why you are seeing these large roundtrip errors. This is inherent to the HEALPix sampling scheme rather than somehting specific to the implementation in s2fft. At a higher computational cost you can achieve a lower roundtrip error by using iterative refinement in the forward transform - the number of iterate refinements steps to use is exposed in the s2fft.forward function via the iter argument. For instance in your example changing to iter=3

import s2fft
import numpy as np
import jax

jax.config.update("jax_enable_x64", True)

L = 128
nside = 64
method = "jax"
sampling = "healpix"
rng = np.random.default_rng(23457801234570)
flm = s2fft.utils.signal_generator.generate_flm(rng, L,)

f = s2fft.inverse(flm, L, nside=nside, sampling = sampling, method=method)

flm = s2fft.forward(f, L, nside=nside, sampling=sampling, method=method, iter=3)
f_recov = s2fft.inverse(flm, L, nside=nside, sampling=sampling,  method=method)

print(f"Mean absolute error = {np.nanmean(np.abs(f_recov - f))}")

outputs

Mean absolute error = 1.4234506825162816e-05