Repository: astro-informatics/s2fft (language: python)

Issue:
Derivatives for `healpix_forward` are incorrrect
I believe the derivatives defined for the `s2fft.transforms.c_backend_spherical.healpy_forward` function with `custom_vjp` are incorrect. As a minimal reproducing example:

```Python
import s2fft
from s2fft.transforms import c_backend_spherical as c_sph
import numpy as np
import jax
from jax.test_util import check_grads
jax.config.update("jax_enable_x64", True)

L = 32
nside = L // 2
method = "jax"
sampling = "healpix"
reality = True
rng = np.random.default_rng(23457801234570)
flm = s2fft.utils.signal_generator.generate_flm(rng, L, reality=reality)
f = s2fft.inverse(flm, L, nside=nside, sampling=sampling, method=method, reality=reality)
```

Running `check_grads` on `c_sph.healpy_inverse` to check gradients against numerical finite differencing completes without any error

```Python
check_grads(lambda x: c_sph.healpy_inverse(x, L, nside), (flm,), modes=("rev",), order=1)
```

Running the same on `c_sph.healpy_forward` however gives an `AssertionError`

```Python
check_grads(lambda x: c_sph.healpy_forward(x, L, nside, iter=0), (f,), modes=("rev",), order=1)
```

outputting

```Python
AssertionError: 
Not equal to tolerance rtol=1e-05, atol=1e-05
VJP cotangent projection
Mismatched elements: 1 / 1 (100%)
Max absolute difference: 1.84452973
Max relative difference: 1.69768919
 x: array(-2.931024)
 y: array(-1.086494)
```

Somewhat confusingly, running `check_grads` instead on specific scalar functions constructed using `c_sph.healpy_forward` does pass without error:

```Python
check_grads(lambda x: abs(c_sph.healpy_forward(x, L, nside, iter=0)).sum(), (f,), modes=("rev",), order=1)
```

Notice in all cases `iter` is fixed to zero so this is not due to using iterative refinement steps. While HEALPix sampling does not exhibit a sampling theorem and so round-tripping through the forward and backward ('inverse') transforms will introduce an error (which iterative refinement in the forward transform can reduce), this only affects the accuracy of the linear operator corresponding to the forward transform being the inverse of the linear operator corresponding to the backward transform. 

Distinct from this property, is that the linear operator represented by the forward transform is a (scaled and conjugated) transposition of the linear operator represented by the backward transform (and vice-versa), and this property holds exactly (modulo floating point error) for the HEALPix forward and backward transforms. Importantly it is this transposition property that is required for implementation of the derivative rules and so the incorrectness of the derivatives in the current implementation cannot be explained by the error in the inverse relationship.

Let the matrix represented by HEALPix forward spherical transform (`map2alm`) be $F \in \mathbb{C}^{m \times n}$ and the matrix represented by HEALPix inverse (backward) spherical transform (`alm2map`) be $B \in \mathbb{C}^{n \times m}$, with a harmonic bandlimit $\ell$ in both cases and HEALPix resolution parameter $r$ and $n = 12r^2$, $m = \ell(\ell + 1) / 2$. That is

$$\texttt{map2alm}(x) = F x \quad\text{and}\quad \texttt{alm2map}(y) = B y.$$

We can construct the matrices numerically by mapping the standard basis vectors through the `map2alm` and `alm2map` functions:

```Python

F = np.stack(
    list(
        map(
            lambda e: healpy.map2alm(e, lmax=L - 1),
            np.identity(12 * nside**2, dtype=float),
        )
    ),
    1,
)

# As alm2map argument is complex need to use both real and imaginary basis vectors

B = np.stack(
    list(
        map(
            lambda e: healpy.alm2map(e, nside=nside, lmax=L - 1),
            np.identity(flm_hp.shape[0], dtype=complex),
        )
    ),
    1,
) - 1j * np.stack(
    list(
        map(
            lambda e: healpy.alm2map(e, nside=nside, lmax=L - 1),
            1j * np.identity(L * (L + 1) // 2, dtype=complex),
        )
    ),
    1,
)
```

We then have the relationships

$$ B = F^H D \quad\text{and}\quad F = D^{-1} B^H$$

where $X^H= \textrm{conj}(X)^T$ (that is the conjugate / Hermitian transpose) and $D$ is a $n \times n$ diagonal matrix with the first $\ell$ entries along the diagonal equal to $3r^2 / \pi$ and the remaining $n - \ell = \ell(\ell - 1) / 2$ diagonal entries equal to $6r^2/\pi$.

Numerical verification:

```Python
D = np.diag([1] * L + [2] * (L * (L - 1) // 2)) *  (3 * nside**2) / np.pi
np.allclose(B, F.T.conj() @ D)
D_inv = np.diag(1 / D.diagonal())
np.allclose(D_inv @ B.T.conj(), F)
```
As `map2alm` and `alm2map` are both linear, their Jacobian vector product (JVP) functions are just the original maps:

$$
  \mathsf{jvp}(\texttt{map2alm})(x)(v) = \partial(x \mapsto F x) v = F v = \texttt{map2alm}(v)
$$
$$
  \mathsf{jvp}(\texttt{alm2map})(y)(v) = \partial(y \mapsto B y) v = B v = \texttt{alm2map}(v)
$$

For the vector Jacobian product (VJP) functions we have

$$
\begin{aligned}
\mathsf{vjp}(\texttt{map2alm})(x)(v) 
&= v^T \partial(x \mapsto F x) \\
&= v^TF \\ 
&= (F^T v)^T\\
&= \textrm{conj}((\textrm{conj}(F)^T\, \textrm{conj}(v)))^T\\
&= (F^H D D^{-1} \,\textrm{conj}(v))^H\\
&= (B D^{-1} \textrm{conj}(v))^H\\
& = \textrm{conj}(\texttt{alm2map}(D^{-1} \textrm{conj}(v)))
\end{aligned}
$$

$$
\begin{aligned}
\mathsf{vjp}(\texttt{alm2map})(y)(v) 
&= v^T\partial(y \mapsto B y) \\
&= v^TB \\ 
&= (B^T v)^T\\
&= \textrm{conj}((\textrm{conj}(B)^T\textrm{conj}(v)))^T\\
&= (D D^{-1} B^H\textrm{conj}(v))^H\\
&= (D F \textrm{conj}(v))^H\\
& = D  \textrm{conj}(\texttt{alm2map}(\textrm{conj}(v)))
\end{aligned}
$$

The current (correct) VJP definition for `healpy_inverse`  is 

https://github.com/astro-informatics/s2fft/blob/909e6f11c6b29ffcdaef6c5e055a3bc57addf297/s2fft/transforms/c_backend_spherical.py#L282-L297

It's a little difficult to relate this to the above derivation as this includes both the VJP for `healpy.alm2map` but also the VJP for `s2fft.sampling.reindex.flm_2d_to_hp_fast` with `healpy_inverse` defined as the composition of these:

https://github.com/astro-informatics/s2fft/blob/909e6f11c6b29ffcdaef6c5e055a3bc57addf297/s2fft/transforms/c_backend_spherical.py#L271-L273

Broadly though we can see the VJP rule corresponds to something which performs $\textrm{conj}(\texttt{alm2map}(\textrm{conj}(v)))$ plus an elementwise scaling operation corresponding to the application of $D$

The current (I think incorrect) VJP definition for `healpy_forward` is 

https://github.com/astro-informatics/s2fft/blob/909e6f11c6b29ffcdaef6c5e055a3bc57addf297/s2fft/transforms/c_backend_spherical.py#L341-L348

with `healpy_forward` here defined as a composition of `healpy.map2alm` and `s2fft.sampling.reindex.flm_hp_to_2d_fast`

https://github.com/astro-informatics/s2fft/blob/909e6f11c6b29ffcdaef6c5e055a3bc57addf297/s2fft/transforms/c_backend_spherical.py#L331-L332

Here we see there is no elementwise scaling before the application of `alm2map` corresponding to multiplication by $D^{-1}$ in $\textrm{conj}(\texttt{alm2map}(D^{-1} \textrm{conj}(v)))$
Cufft primitive
Creating a pure cuda healpix fft to avoid the long compile time when using pure JAX


Discussion / hints:

@matt-graham 
I have implemented what you asked

- use nanobind as pip package and not submodule
- remove _src imports when possible
- to be able to compile without cuda (but show error message when the user tries to call cuda functions)

I think that this PR can be merged (if it is all ok for you)
Things we need to take care of 

- [ ] removing dependency on pyssht if it is failing https://github.com/astro-informatics/s2fft/issues/206
- [ ] formatting https://github.com/astro-informatics/s2fft/issues/205
> Hi @ASKabalan. Thanks for the updates to the pull-request, we're hoping to try to get this merged in as soon as possible.
> 
> At the moment several (previously passing) tests in `tests/test_healpix_ffts.py` are failing because of a mismatch between the output of `healpix_ifft_jax` and `healpix_ifft_numpy`. The whitespace changes in `s2fft/utils/healpix_ffts.py` make it difficult to see where the actual substantive code changes in this file are which might have caused these tests to start failing. Could I ask that you revert the changes which are not connected to this PR to help us identify where this issue is arising?
> 
> We also need to have the tests which require a CUDA GPU to be skipped when running on non-GPU systems (such as the GitHub Action test runners). I can update the tests to do this if you are okay for me to commit to this PR branch directly.

I suggest doing a seperate formatting PR from the main branch
Format using a pre-commit formatter
Then I do the same and merge this way the comparaison is easier


I changed one thing in the JAX ifft function.
This https://github.com/astro-informatics/s2fft/blob/main/s2fft/utils/healpix_ffts.py#L478
to https://github.com/astro-informatics/s2fft/blob/cufft-primitive/s2fft/utils/healpix_ffts.py#L470

Using the actual IFFT instead of the FFT of the complex conjugate
We discussed this in the slack channel
> I suggest doing a seperate formatting PR from the main branch Format using a pre-commit formatter Then I do the same and merge this way the comparaison is easier

We can so a separate formatting PR and merge this in first, but this is likely to result in quite messy merge conflicts as if we are reformatting the whole codebase then we will be touching a lot of lines in the same files that are being updated in this PR. If you prefer to go this way we can do so, but I would say reverting the white space changes here would be a lot simpler.

> I changed one thing in the JAX ifft function. This https://github.com/astro-informatics/s2fft/blob/main/s2fft/utils/healpix_ffts.py#L478 to https://github.com/astro-informatics/s2fft/blob/cufft-primitive/s2fft/utils/healpix_ffts.py#L470
> 
> Using the actual IFFT instead of the FFT of the complex conjugate We discussed this in the slack channel

My understanding from the conversation in Slack was that the change shouldn't alter the numerical output, it was just meant to be removing a previous workaround for problems in 64-bit precision with the underlying JAX IFFT implementation? Either way as the test is failing, either the test or implementation need updating.

Looking at the change from

```Python
jnp.conj(jnp.fft.fft(jnp.fft.ifftshift(jnp.conj(fm_chunks), axes=-1), norm="backward")        
```

to

```Python
jnp.fft.ifft(jnp.fft.ifftshift(fm_chunks, axes=-1), norm="backward")
```

I think the problem maybe a difference in normalization factors as [this note](https://adamsiembida.com/how-to-compute-the-ifft-using-only-the-forward-fft/) suggests the equivalence is

$$
  \mathsf{IFFT}(X) = \frac{1}{N} \mathsf{conj}(\mathsf{FFT}(\mathsf{conj}(\mathsf{FFT}(X)))
$$

on the basis of a $\frac{1}{N}$ normalization factor being included in `jnp.fft.ifft` and no normalization in `jnp.fft.fft`, which from the [NumPy documentation](https://numpy.org/doc/stable/reference/routines.fft.html#normalization) I think should correspond to the default behaviour when `norm` arguments are not specified or when `norm = "backward"` (and I believe JAX should follow NumPy behaviour).

To retain the previous behaviour I think we need to account for the $\frac{1}{N}$ factor, either explicitly or by setting the `norm` argument to `ifft` accordingly. Looking at your changes, you appear to have `norm = "backward"` in `healpix_ifft_jax` while `healpix_ifft_numpy` uses `norm = "forward"` so that suggests you need to update to follow the NumPy implementation here, or if we decide on a different normalization convention we will need to update all of the IFFT function to use the same convention, though as this would be a breaking change I think we should avoid if possible.

EDIT: Tagging @CosmoMatt for comment here as he is probably best placed to advise on what the normalization conventions we should be following are!
@matt-graham You were right.
I fixed the issue all tests are good now
I also added a condition on the test to run only if GPU is available.
## [Codecov](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) Report
Attention: Patch coverage is `82.50000%` with `7 lines` in your changes missing coverage. Please review.
> Project coverage is 93.28%. Comparing base [(`5df8106`)](https://app.codecov.io/gh/astro-informatics/s2fft/commit/5df8106549eb2b19438e12bc29d424a40f3b7bab?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) to head [(`583aad4`)](https://app.codecov.io/gh/astro-informatics/s2fft/commit/583aad46dcf8de219af2a023341b1a7279b6ab91?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics).
> Report is 54 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) | Patch % | Lines |
|---|---|---|
| [s2fft/utils/jax\_primitive.py](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&filepath=s2fft%2Futils%2Fjax_primitive.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics#diff-czJmZnQvdXRpbHMvamF4X3ByaW1pdGl2ZS5weQ==) | 83.33% | [3 Missing :warning: ](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) |
| [s2fft/transforms/spherical.py](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&filepath=s2fft%2Ftransforms%2Fspherical.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics#diff-czJmZnQvdHJhbnNmb3Jtcy9zcGhlcmljYWwucHk=) | 66.66% | [2 Missing :warning: ](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) |
| [s2fft/utils/healpix\_ffts.py](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&filepath=s2fft%2Futils%2Fhealpix_ffts.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics#diff-czJmZnQvdXRpbHMvaGVhbHBpeF9mZnRzLnB5) | 87.50% | [2 Missing :warning: ](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #204      +/-   ##
==========================================
- Coverage   93.42%   93.28%   -0.15%     
==========================================
  Files          28       29       +1     
  Lines        3135     3170      +35     
==========================================
+ Hits         2929     2957      +28     
- Misses        206      213       +7     
```

| [Flag](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) | Coverage Δ | |
|---|---|---|
| [](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) | `?` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics).

> ## [Codecov](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) Report
> 
> Attention: Patch coverage is `59.84848%` with `53 lines` in your changes missing coverage. Please review.
> 
> > Project coverage is 91.76%. Comparing base [(`2c3d9e5`)](https://app.codecov.io/gh/astro-informatics/s2fft/commit/2c3d9e5af52940846b265d8912fc7ab0296f6891?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) to head [(`c552ce5`)](https://app.codecov.io/gh/astro-informatics/s2fft/commit/c552ce5718aea4a3377813c8214ba66bc7dd05d0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics).
> 
> [Files with missing lines](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics) 	Patch % 	Lines
> [s2fft/utils/healpix_ffts.py](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&filepath=s2fft%2Futils%2Fhealpix_ffts.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics#diff-czJmZnQvdXRpbHMvaGVhbHBpeF9mZnRzLnB5) 	38.35% 	[45 Missing ⚠️ ](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics)
> [s2fft/utils/jax_pritimive.py](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&filepath=s2fft%2Futils%2Fjax_pritimive.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics#diff-czJmZnQvdXRpbHMvamF4X3ByaXRpbWl2ZS5weQ==) 	86.44% 	[8 Missing ⚠️ ](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics)
> Additional details and impacted files
> 
> 		
> 		
> 
> [☔ View full report in Codecov by Sentry](https://app.codecov.io/gh/astro-informatics/s2fft/pull/204?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics). 📢 Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astro-informatics).

The tests are covered.
The fact that there is no GPU does not trigger codecov?
> The tests are covered. The fact that there is no GPU does not trigger codecov?

The codecov results are based on the coverage reported from the GitHub Actions workflow and so won't include the tests you added which are now only run when a GPU is available. I think we can just ignore the codecov checks here or [set the corresponding functions to be ignored in the coverage reports](https://coverage.readthedocs.io/en/latest/excluding.html). 



> > I changed one thing in the JAX ifft function. This https://github.com/astro-informatics/s2fft/blob/main/s2fft/utils/healpix_ffts.py#L478 to https://github.com/astro-informatics/s2fft/blob/cufft-primitive/s2fft/utils/healpix_ffts.py#L470
> > Using the actual IFFT instead of the FFT of the complex conjugate We discussed this in the slack channel
> 
> My understanding from the conversation in Slack was that the change shouldn't alter the numerical output, it was just meant to be removing a previous workaround for problems in 64-bit precision with the underlying JAX IFFT implementation? Either way as the test is failing, either the test or implementation need updating.
> 
> Looking at the change from
> 
> ```python
> jnp.conj(jnp.fft.fft(jnp.fft.ifftshift(jnp.conj(fm_chunks), axes=-1), norm="backward")        
> ```
> 
> to
> 
> ```python
> jnp.fft.ifft(jnp.fft.ifftshift(fm_chunks, axes=-1), norm="backward")
> ```
> 
> I think the problem maybe a difference in normalization factors as [this note](https://adamsiembida.com/how-to-compute-the-ifft-using-only-the-forward-fft/) suggests the equivalence is
> 
> IFFT ( X ) = 1 N conj ( FFT ( conj ( FFT ( X ) ) )
> 
> on the basis of a 1 N normalization factor being included in `jnp.fft.ifft` and no normalization in `jnp.fft.fft`, which from the [NumPy documentation](https://numpy.org/doc/stable/reference/routines.fft.html#normalization) I think should correspond to the default behaviour when `norm` arguments are not specified or when `norm = "backward"` (and I believe JAX should follow NumPy behaviour).
> 
> To retain the previous behaviour I think we need to account for the 1 N factor, either explicitly or by setting the `norm` argument to `ifft` accordingly. Looking at your changes, you appear to have `norm = "backward"` in `healpix_ifft_jax` while `healpix_ifft_numpy` uses `norm = "forward"` so that suggests you need to update to follow the NumPy implementation here, or if we decide on a different normalization convention we will need to update all of the IFFT function to use the same convention, though as this would be a breaking change I think we should avoid if possible.
> 
> EDIT: Tagging @CosmoMatt for comment here as he is probably best placed to advise on what the normalization conventions we should be following are!

Just had a chance to catch up on these discussions.  It seems like this has all been resolved now but for further background, @CosmoMatt found that `jnp.fft.ifft` wasn't fully supported in JAX(!) back when we first implemented this so as a workaround he implemented the IFFT via `jnp.fft.fft`.  It seems the JAX bug was resolved in more recent versions of JAX and so the code could be tidied up to use `jnp.fft.ifft` directly.  (@CosmoMatt was this applied anywhere else that we should also update now?)

The line itself shouldn't apply any normalisation (that is handled elsewhere), so using an `jnp.fft.fft` with `norm="backward"` achieves that.  Since we now switch to `jnp.fft.ifft` we should switch to `norm="forward"`, which ensures no normalisation is applied in that line of code.  

In general, I think we've tried to follow the convention where `jnp.fft.fft` and `jnp.fft.ifft` are performed without any normalisation and we instead apply the appropriate normalisation explicitly ourselves since often there are additional normalisation factors to be considered.  That corresponds to setting `norm="backward"` when calling `jnp.fft.fft` and setting `norm="forward"` when calling `jnp.fft.ifft` (although I couldn't say for sure this approach is applied absolutely everywhere in the code but that is least the generally approach we've typically been adopting).
Hi @ASKabalan. Thanks for your updates to the PR. There are still some previous review comments which I think are unresolved I've catalogued below and I've added some further comments from another review

- [x] Commented out code in `s2fft/utils/healpix_ffts.py` [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1638338763)
- [x] Fixing typo in name of `s2fft/utils/jax_pritimive.py` file [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1638277123)
- [x] Fixing indentation in docstrings in `s2fft/utils/jax_pritimive.py` file [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1638279292)
- [x] Fixing broken URL in comment in `s2fft/utils/jax_pritimive.py` [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1638280737)
- [x] Resolving use of external code by adjusting copyright statement in license or removing use of `BasePrimitive` abstract base class [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1638338763)
- [x] Resolving licensing for  external code in `lib/include/kernel_helpers.h` [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204#discussion_r1638357861)
- [x] Adjusting license in header for `lib/include/kernel_nanobind_helpers.h` to list modifications made [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204#discussion_r1638376027)
- [x] Fetching `argparse` library in `CMakesLists.txt` rather than including in repository [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204#discussion_r1638143845)
- [x] Remove `tools.setuptools.dynamic` section from `pyproject.toml` and remove redundant `requirements` files [(link to commment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1782453077)
- [x] Updating docstring and argument name for `use_cuda` argument [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204/files#r1782473001)
- [x] Revert renaming of `LICENCE.txt` to `LICENSE.txt` [(link to comment)](https://github.com/astro-informatics/s2fft/pull/204#discussion_r1771057518)

If you prefer I am happy to deal with addressing these.
@matt-graham  It is ok for me .. go for merge?
> @matt-graham It is ok for me .. go for merge?

The failing tests are as I forgot we're targetting Python 3.9 and above and I used some Python 3.10 syntax in the type hints. I will fix this now. Then hopefully if tests then pass from my side I'm good to merge, but also tagging @jasonmcewen and @CosmoMatt in case they want to do any final reviews.
@matt-graham Also I think that the test should be marked to not be included in the codecov workflow
Should i do it?

@all-contributors please add @ASKabalan for code, review, test
@matt-graham 

I've put up [a pull request](https://github.com/astro-informatics/s2fft/pull/232) to add @ASKabalan! :tada:

Fix patch:
diff --git a/s2fft/transforms/c_backend_spherical.py b/s2fft/transforms/c_backend_spherical.py
index 5b6df8f5..da6a9ed4 100644
--- a/s2fft/transforms/c_backend_spherical.py
+++ b/s2fft/transforms/c_backend_spherical.py
@@ -4,7 +4,8 @@
 
 # C backend functions for which to provide JAX frontend.
 import pyssht
-from jax import custom_vjp
+from jax import core, custom_vjp
+from jax.interpreters import ad
 
 from s2fft.sampling import reindex
 from s2fft.utils import quadrature_jax
@@ -241,83 +242,181 @@ def _ssht_forward_bwd(res, flm):
     return f, None, None, None, None, None
 
 
-@custom_vjp
-def healpy_inverse(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
-    r"""
-    Compute the inverse scalar real spherical harmonic transform (HEALPix JAX).
+# Link JAX gradients for C backend functions
+ssht_inverse.defvjp(_ssht_inverse_fwd, _ssht_inverse_bwd)
+ssht_forward.defvjp(_ssht_forward_fwd, _ssht_forward_bwd)
 
-    HEALPix is a C++ library which implements the scalar spherical harmonic transform
-    outlined in [1]. We make use of their healpy python bindings for which we provide
-    custom JAX frontends, hence providing support for automatic differentiation. Currently
-    these transforms can only be deployed on CPU, which is a limitation of the C++ library.
 
-    Args:
-        flm (jnp.ndarray): Spherical harmonic coefficients.
+def _complex_dtype(real_dtype):
+    """
+    Get complex datatype corresponding to a given real datatype.
 
-        L (int): Harmonic band-limit.
+    Derived from https://github.com/jax-ml/jax/blob/1471702adc28/jax/_src/lax/fft.py#L92
 
-        nside (int, optional): HEALPix Nside resolution parameter.  Only required
-            if sampling="healpix".  Defaults to None.
+    Original license:
 
-    Returns:
-        jnp.ndarray: Signal on the sphere.
+    Copyright 2019 The JAX Authors.
 
-    Note:
-        [1] Gorski, Krzysztof M., et al. "HEALPix: A framework for high-resolution
-        discretization and fast analysis of data distributed on the sphere." The
-        Astrophysical Journal 622.2 (2005): 759
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+        https://www.apache.org/licenses/LICENSE-2.0
 
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
     """
-    flm = reindex.flm_2d_to_hp_fast(flm, L)
-    f = jnp.array(healpy.alm2map(np.array(flm), lmax=L - 1, nside=nside))
-    return f
+    return (np.zeros((), real_dtype) + np.zeros((), np.complex64)).dtype
+
+
+def _real_dtype(complex_dtype):
+    """
+    Get real datatype corresponding to a given complex datatype.
+
+    Derived from https://github.com/jax-ml/jax/blob/1471702adc28/jax/_src/lax/fft.py#L93
+
+    Original license:
+
+    Copyright 2019 The JAX Authors.
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+        https://www.apache.org/licenses/LICENSE-2.0
 
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+    """
+    return np.finfo(complex_dtype).dtype
+
+
+def _healpy_map2alm_impl(f: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
+    return jnp.array(healpy.map2alm(np.array(f), lmax=L - 1, iter=0))
 
-def _healpy_inverse_fwd(flm: jnp.ndarray, L: int, nside: int):
-    """Private function which implements the forward pass for inverse jax_healpy."""
-    res = ([], L, nside)
-    return healpy_inverse(flm, L, nside), res
 
+def _healpy_map2alm_abstract_eval(
+    f: core.ShapedArray, L: int, nside: int
+) -> core.ShapedArray:
+    return core.ShapedArray(shape=(L * (L + 1) // 2,), dtype=_complex_dtype(f.dtype))
 
-def _healpy_inverse_bwd(res, f):
-    """Private function which implements the backward pass for inverse jax_healpy."""
-    _, L, nside = res
-    f_new = f * (12 * nside**2) / (4 * jnp.pi)
-    flm_out = jnp.array(
-        np.conj(healpy.map2alm(np.conj(np.array(f_new)), lmax=L - 1, iter=0))
+
+def _healpy_map2alm_transpose(dflm: jnp.ndarray, L: int, nside: int):
+    scale_factors = (
+        jnp.concatenate((jnp.ones(L), 2 * jnp.ones(L * (L - 1) // 2)))
+        * (3 * nside**2)
+        / jnp.pi
     )
-    # iter MUST be zero otherwise gradient propagation is incorrect (JDM).
-    flm_out = reindex.flm_hp_to_2d_fast(flm_out, L)
-    m_conj = (-1) ** (jnp.arange(1, L) % 2)
-    flm_out = flm_out.at[..., L:].add(
-        jnp.flip(m_conj * jnp.conj(flm_out[..., : L - 1]), axis=-1)
+    return (jnp.conj(healpy_alm2map(jnp.conj(dflm) / scale_factors, L, nside)),)
+
+
+_healpy_map2alm_p = core.Primitive("healpy_map2alm")
+_healpy_map2alm_p.def_impl(_healpy_map2alm_impl)
+_healpy_map2alm_p.def_abstract_eval(_healpy_map2alm_abstract_eval)
+ad.deflinear(_healpy_map2alm_p, _healpy_map2alm_transpose)
+
+
+def healpy_map2alm(f: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
+    """
+    JAX wrapper for healpy map2alm function (forward spherical harmonic transform).
+
+    This wrapper will return the spherical harmonic coefficients as a one dimensional
+    array using HEALPix (ring-ordered) indexing. To instead return a two-dimensional
+    array of harmonic coefficients use :py:func:`healpy_forward`.
+
+    Args:
+        f (jnp.ndarray): Signal on the sphere.
+
+        L (int): Harmonic band-limit. Equivalent to `lmax + 1` in healpy.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+    Returns:
+        jnp.ndarray: Harmonic coefficients of signal f.
+
+    """
+    return _healpy_map2alm_p.bind(f, L=L, nside=nside)
+
+
+def _healpy_alm2map_impl(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
+    return jnp.array(healpy.alm2map(np.array(flm), lmax=L - 1, nside=nside))
+
+
+def _healpy_alm2map_abstract_eval(
+    flm: core.ShapedArray, L: int, nside: int
+) -> core.ShapedArray:
+    return core.ShapedArray(shape=(12 * nside**2,), dtype=_real_dtype(flm.dtype))
+
+
+def _healpy_alm2map_transpose(df: jnp.ndarray, L: int, nside: int) -> tuple:
+    scale_factors = (
+        jnp.concatenate((jnp.ones(L), 2 * jnp.ones(L * (L - 1) // 2)))
+        * (3 * nside**2)
+        / jnp.pi
     )
-    flm_out = flm_out.at[..., : L - 1].set(0)
+    # Scale factor above includes the inverse quadrature weight given by
+    # (12 * nside**2) / (4 * jnp.pi) = (3 * nside**2) / jnp.pi
+    # and also a factor of 2 for m>0 to account for the negative m.
+    # See explanation in this issue comment:
+    # https://github.com/astro-informatics/s2fft/issues/243#issuecomment-2500951488
+    return (scale_factors * jnp.conj(healpy_map2alm(jnp.conj(df), L, nside)),)
 
-    return flm_out, None, None
+
+_healpy_alm2map_p = core.Primitive("healpy_alm2map")
+_healpy_alm2map_p.def_impl(_healpy_alm2map_impl)
+_healpy_alm2map_p.def_abstract_eval(_healpy_alm2map_abstract_eval)
+ad.deflinear(_healpy_alm2map_p, _healpy_alm2map_transpose)
+
+
+def healpy_alm2map(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
+    """
+    JAX wrapper for healpy alm2map function (inverse spherical harmonic transform).
+
+    This wrapper assumes the passed spherical harmonic coefficients are a one
+    dimensional array using HEALPix (ring-ordered) indexing. To instead pass a
+    two-dimensional array of harmonic coefficients use :py:func:`healpy_inverse`.
+
+    Args:
+        flm (jnp.ndarray): Spherical harmonic coefficients.
+
+        L (int): Harmonic band-limit. Equivalent to `lmax + 1` in healpy.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+    Returns:
+        jnp.ndarray: Signal on the sphere.
+
+    """
+    return _healpy_alm2map_p.bind(flm, L=L, nside=nside)
 
 
-@custom_vjp
 def healpy_forward(f: jnp.ndarray, L: int, nside: int, iter: int = 3) -> jnp.ndarray:
     r"""
     Compute the forward scalar spherical harmonic transform (healpy JAX).
 
     HEALPix is a C++ library which implements the scalar spherical harmonic transform
     outlined in [1]. We make use of their healpy python bindings for which we provide
-    custom JAX frontends, hence providing support for automatic differentiation. Currently
-    these transforms can only be deployed on CPU, which is a limitation of the C++ library.
+    custom JAX frontends, hence providing support for automatic differentiation.
+    Currently these transforms can only be deployed on CPU, which is a limitation of the
+    C++ library.
 
     Args:
         f (jnp.ndarray): Signal on the sphere.
 
         L (int): Harmonic band-limit.
 
-        nside (int, optional): HEALPix Nside resolution parameter.  Only required
-            if sampling="healpix".  Defaults to None.
+        nside (int): HEALPix Nside resolution parameter.
 
-        iter (int, optional): Number of subiterations for healpy. Note that iterations
-            increase the precision of the forward transform, but reduce the accuracy of
-            the gradient pass. Between 2 and 3 iterations is a good compromise.
+        iter (int, optional): Number of subiterations (iterative refinement steps) for
+            healpy. Note that iterations increase the precision of the forward transform
+            as an inverse of inverse transform, but with a linear increase in
+            computational cost. Between 2 and 3 iterations is a good compromise.
 
     Returns:
         jnp.ndarray: Harmonic coefficients of signal f.
@@ -328,28 +427,39 @@ def healpy_forward(f: jnp.ndarray, L: int, nside: int, iter: int = 3) -> jnp.nda
         Astrophysical Journal 622.2 (2005): 759
 
     """
-    flm = jnp.array(healpy.map2alm(np.array(f), lmax=L - 1, iter=iter))
+    flm = healpy_map2alm(f, L, nside)
+    for _ in range(iter):
+        f_recov = healpy_alm2map(flm, L, nside)
+        f_error = f - f_recov
+        flm += healpy_map2alm(f_error, L, nside)
     return reindex.flm_hp_to_2d_fast(flm, L)
 
 
-def _healpy_forward_fwd(f: jnp.ndarray, L: int, nside: int, iter: int = 3):
-    """Private function which implements the forward pass for forward jax_healpy."""
-    res = ([], L, nside, iter)
-    return healpy_forward(f, L, nside, iter), res
+def healpy_inverse(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
+    r"""
+    Compute the inverse scalar real spherical harmonic transform (HEALPix JAX).
 
+    HEALPix is a C++ library which implements the scalar spherical harmonic transform
+    outlined in [1]. We make use of their healpy python bindings for which we provide
+    custom JAX frontends, hence providing support for automatic differentiation.
+    Currently these transforms can only be deployed on CPU, which is a limitation of the
+    C++ library.
 
-def _healpy_forward_bwd(res, flm):
-    """Private function which implements the backward pass for forward jax_healpy."""
-    _, L, nside, _ = res
-    flm_new = reindex.flm_2d_to_hp_fast(flm, L)
-    f = jnp.array(
-        np.conj(healpy.alm2map(np.conj(np.array(flm_new)), lmax=L - 1, nside=nside))
-    )
-    return f * (4 * jnp.pi) / (12 * nside**2), None, None, None
+    Args:
+        flm (jnp.ndarray): Spherical harmonic coefficients.
 
+        L (int): Harmonic band-limit.
 
-# Link JAX gradients for C backend functions
-ssht_inverse.defvjp(_ssht_inverse_fwd, _ssht_inverse_bwd)
-ssht_forward.defvjp(_ssht_forward_fwd, _ssht_forward_bwd)
-healpy_inverse.defvjp(_healpy_inverse_fwd, _healpy_inverse_bwd)
-healpy_forward.defvjp(_healpy_forward_fwd, _healpy_forward_bwd)
+        nside (int): HEALPix Nside resolution parameter.
+
+    Returns:
+        jnp.ndarray: Signal on the sphere.
+
+    Note:
+        [1] Gorski, Krzysztof M., et al. "HEALPix: A framework for high-resolution
+        discretization and fast analysis of data distributed on the sphere." The
+        Astrophysical Journal 622.2 (2005): 759
+
+    """
+    flm = reindex.flm_2d_to_hp_fast(flm, L)
+    return healpy_alm2map(flm, L, nside)


Test patch:
diff --git a/tests/test_spherical_custom_grads.py b/tests/test_spherical_custom_grads.py
index 15f0dcd5..c3fb3938 100644
--- a/tests/test_spherical_custom_grads.py
+++ b/tests/test_spherical_custom_grads.py
@@ -307,22 +307,16 @@ def func(f):
 @pytest.mark.parametrize("nside", nside_to_test)
 @pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_healpix_c_backend_inverse_custom_gradients(flm_generator, nside: int):
-    sampling = "healpix"
     L = 2 * nside
     reality = True
     flm = flm_generator(L=L, reality=reality)
-    flm_target = flm_generator(L=L, reality=reality)
-    f_target = spherical.inverse_jax(
-        flm_target, L, nside=nside, sampling=sampling, reality=reality
-    )
 
     def func(flm):
-        f = spherical.inverse(
+        return spherical.inverse(
             flm, L, 0, nside, sampling="healpix", method="jax_healpy", reality=True
         )
-        return jnp.sum(jnp.abs(f - f_target) ** 2)
 
-    check_grads(func, (flm,), order=1, modes=("rev"))
+    check_grads(func, (flm,), order=2, modes=("fwd", "rev"))
 
 
 @pytest.mark.parametrize("nside", nside_to_test)
@@ -334,16 +328,12 @@ def test_healpix_c_backend_forward_custom_gradients(
     sampling = "healpix"
     L = 2 * nside
     reality = True
-    flm_target = flm_generator(L=L, reality=reality)
     flm = flm_generator(L=L, reality=reality)
     f = spherical.inverse_jax(flm, L, nside=nside, sampling=sampling, reality=reality)
 
     def func(f):
-        flm = spherical.forward(
+        return spherical.forward(
             f, L, nside=nside, sampling="healpix", method="jax_healpy", iter=iter
         )
-        return jnp.sum(jnp.abs(flm - flm_target) ** 2)
-
-    rtol = [1e-6, 1e-2, 5e-2, 1e-2][iter]
 
-    check_grads(func, (f,), order=1, modes=("rev"), rtol=rtol)
+    check_grads(func, (f,), order=2, modes=("fwd", "rev"))