s2fft
s2fft copied to clipboard
Add stable high-spin transforms (precompute, standard)
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 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.
@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.
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:
@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.