s2fft icon indicating copy to clipboard operation
s2fft copied to clipboard

Add stable high-spin transforms (precompute, standard)

Open CosmoMatt opened this issue 1 year ago • 3 comments

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.

CosmoMatt avatar Sep 30 '24 12:09 CosmoMatt

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 :warning:
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     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

codecov[bot] avatar Sep 30 '24 12:09 codecov[bot]

@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.

CosmoMatt avatar Oct 01 '24 13:10 CosmoMatt

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:

broken_s2fft_inverse_wigner

New figure generated with dec8cb1e:

fixed_s2fft_inverse_wigner

ElisR avatar Oct 03 '24 09:10 ElisR

@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.

CosmoMatt avatar Oct 06 '24 13:10 CosmoMatt