Add stable high-spin transforms (precompute, standard)

[CosmoMatt]

This PR wraps the existing Risbo Wigner d-function recursions to generate the necessary precompute matrices for the forward/inverse spherical harmonic/Wigner transforms.

Method:

To do this for each $\ell \in [0, L)$ the Wigner d plane for $m,n$ is computed for $\beta = \frac{\pi}{2}$. Then the coefficients $d^{\ell}_{mn}(\beta)$ for all $\lbrace m,n,\beta \rbrace$ are efficiently computed through an IRFFT (see equation 8 of this paper, or discussion 3.1 of this paper).

A single IRFFT is necessary for each $\lbrace \ell, m, n \rbrace$ each of which is of length $\mathcal{O}(L)$ with complexity $\mathcal{O}(L \log L)$ and therefore the complexity here is $\mathcal{O}(NL^3 \log L)$. Alternatively we can manually perform the recursion for each $\beta$ with complexity $\mathcal{O}(L^4)$. Note that even though N may be $\leq L$ for the Risbo recursion we still need all $N$ unfortunately. Therefore when $N <= L / \log L$ the FFT approach is better, otherwise it's better to compute the recursion manually.

I've added a switch to automatically switch between these modes depending on the parameter configuration. In any case, the same $\mathcal{O}(NL^3 \sim L^4)$ memory overhead still applies.

Further, more efficient, alternatives will be added in a subsequent PR.

[codecov]

Codecov Report

Attention: Patch coverage is 95.89744% with 8 lines in your changes missing coverage. Please review.

Project coverage is 96.08%. Comparing base (1254b92) to head (b3e0333).
Report is 11 commits behind head on main.

Files with missing lines	Patch %	Lines
s2fft/precompute_transforms/construct.py	94.52%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #230      +/-   ##
==========================================
+ Coverage   93.27%   96.08%   +2.80%     
==========================================
  Files          29       29              
  Lines        3167     3319     +152     
==========================================
+ Hits         2954     3189     +235     
+ Misses        213      130      -83     

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

[CosmoMatt]

@jasonmcewen I've now reworked the existing functions which precompute the necessary Wigner d-function elements to switch between recursion depending on the users configuration.

Ultimately, a user should now be able to run the precompute harmonic and Wigner transforms at arbitrary spins and recover machine precision (except for HEALPix sampling of course, but that is another issue entirely...)

@ElisR hopefully this should solve your immediate problems, at least up to L = N = 128. I'll pick up the more memory efficient routines in a separate PR.

[jasonmcewen]

Is all looking great @CosmoMatt!

A few comments but they mostly relate to very minor clarity related things. Note that for some of these I've only highlighted the first instance, so please be sure to address all related instances. Thanks!

One question. Do we only support Risbo for the Wigner case then? I don't recall all the details of what we discussed the other day. Was there no setting where McEwen-Price was better? Perhaps not actually.

And test coverage seems to be stopping the merge. Is looks like the new code has good coverage so not exactly sure what's going on there?

[ElisR]

Hi @CosmoMatt, this is looking very promising, thanks!

I ran my minimal example on this branch (after changing which API function I called) and the results look excellent, at least on CPU.

Old figure from #209:

New figure generated with dec8cb1e:

[CosmoMatt]

@ElisR Ok great that it seems to have solved your immediate problem! If there's any further functionality that you'd like to see (or think would be useful) in the package, let us know and we can add it to our board.

@jasonmcewen I'll pick up the suggested changes when I next get a chance, then we can merge before implementing the various memory efficient versions of the Risbo transforms.

[jasonmcewen]

@CosmoMatt I caught a couple of very minor typos. Can you correct these and then merge? I've approved the PR already so no need for me to take another look. Please just merge when you've addressed these very minor outstanding points.