Tests for consistency of HEALPix FFT and IFFT implementations

[matt-graham]

Adds some tests for consistency between the NumPy and JAX implementations of HEALPix fast Fourier transform (FFT) and inverse FFT functions in s2fft.utils.healpix_ffts.

Assumes that

1. healpix_fft_numpy and healpix_fft_jax should give consistent outputs (up to floating point error) for equivalent input arguments
2. healpix_ifft_numpy and healpix_ifft_jax should give consistent outputs (up to floating point error) for equivalent input arguments,
3. healpix_fft_numpy and healpix_ifft_numpy as partial functions only on first arguments are inverses of each other,
4. healpix_fft_jax and healpix_ifft_jax as partial functions only on first arguments are inverses of each other.

I'm not sure that 3 and 4 are necessarily true or not.

At the moment only the tests for 1 all pass with tests for 2, 3 and 4 failing. This may well be due to the assumptions above being invalid.

[matt-graham]

@CosmoMatt has confirmed assumptions (3) and (4) above will not necessarily hold:

FFT step for HEALPix sampling is exactly where they lose accuracy, primarily because they throw away samples in the polar regions. So we wouldn't expect to get back to the same input, other than in the equitorial band that is.

so removed the corresponding tests.

[CosmoMatt]

@matt-graham thanks for setting this up. I've added a fix, basically the way the test was written you are comparing the consistency between two functions applied to an input vector. However, when the first (numpy baseline) function is applied it appears to update the input vector, so the second function is being applied to a different vector, hence producing different results.

I've added a "fix" to the test which explicitly copies the input vector to avoid this failure mode.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (11592c0) 91.63% compared to head (4fe1afc) 91.63%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #170   +/-   ##
=======================================
  Coverage   91.63%   91.63%           
=======================================
  Files          22       22           
  Lines        2510     2510           
=======================================
  Hits         2300     2300           
  Misses        210      210           

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

[matt-graham]

Should we add some tests here that check that the healpix_*fft_* don't mutate their arguments or do you think that would be better as a separate PR along with a fix for the current behaviour?

[CosmoMatt]

I think this should be fine as is @matt-graham. The arrays that are mutated are generated internally to the overall harmonic transform and it's unlikely that any user will run individual components of the overall transform independently (such as the FFT step)