Iterative refinement support for JAX and NumPy forward (spherical) transform implementations

[matt-graham]

Should resolve #139

Adds support for using iter keyword argument to s2fft.forward when using method = "numpy" and method = "jax", using an iterative refinement approach equivalent to that implemented in healpy to improve accuracy of forward transform (in terms of acting as the inverse of the inverse transform).

Currently this is only added for spherical transform, and is available for all sampling types, even though in practice probably only useful for HEALPix (though iter=1 does seem to give a small gain in accuracy for other sampling types).

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.21%. Comparing base (dc9b2bc) to head (23fb7af).
Report is 56 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #241      +/-   ##
==========================================
+ Coverage   96.07%   96.21%   +0.13%     
==========================================
  Files          31       32       +1     
  Lines        3565     3591      +26     
==========================================
+ Hits         3425     3455      +30     
+ Misses        140      136       -4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[CosmoMatt]

@matt-graham looks great, one thing that would also be good to add is the same iterations but for the precompute versions of the transforms.

[jasonmcewen]

Good point @CosmoMatt (re precompute). We could either do that in the current PR if quick or else do it in a next one.

[matt-graham]

@CosmoMatt - just to check by precompute versions, do you mean the forward and inverse functions in s2fft/precompute_transforms/spherical.py?

Also, with regards to the precomps argument to the forward and inverse functions in s2fft/transforms/spherical.py - at the moment the current implementation passes whatever value for precomps is passed to the top level forward call to both the underlying forward and inverse function calls. With precomps = None this works fine, but thinking about it I think this will give incorrect results if precomps != None as presumable the precomputed values for forward and inverse transforms differ? Also as we are making multiple calls to the forward and inverse transforms when iter > 0 presumably it is always worth pre-computing in this situation?

[matt-graham]

• Add tests checking base transform with iter > 0 against healpy + check reconstruction for large n_iter close to zero

[CosmoMatt]

Great to have this functionality and to extend it throughout all the different versions of the functions!