add stable forward/inverse memory efficient Wigner transforms

[CosmoMatt]

This PR extends the Risbo precompute Wigner transform to the setting where we instead ONLY compute the Wigner d coefficients for $\theta = \beta = \pi/2$ which we denote $\Delta^l_{mn} = d^l_{mn}(\pi/2)$. The coefficients for colatitudes other than $\pi/2$ are accessed implicitly at run-time through the Fourier relation outlined by this McEwen et al 2016 with further details in this McEwen et al 2011.

Ultimately this algorithm costs an extra FFT at run-time but reduces the memory overhead from $\mathcal{O}(NL^3)\sim\mathcal{O}(L^4)$ to $\mathcal{O}(L^3)$. A further reduction in memory by a factor of 8 can be achieved by exploiting the various symmetries of $\Delta^l_{mn}$, however this has yet to be implemented.

As the Risbo recursion is extremely stable we can also consider reduced precision computation (f32, f16, f8) which should recover order ($10^{-8}, 10^{-4}, 10^{-2}$) respectively with corresponding memory reduction and acceleration by a factor of ($2, 4, 8$) respectively.

Tip

When building flax layers for $SO(3)$ convolutions using this implementation the dominant memory overhead is entirely independent from the number of channels, so the memory requirement for the precompute transform is fixed.

Note

With both the memory reductions from aforementioned symmetry relations and from the switched algorithm, the > precompute algorithm should be runnable at $L \sim 2048$ and potentially slightly beyond. However, this has yet to be tested explicitly.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.22%. Comparing base (b3e0333) to head (8e33569).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #238      +/-   ##
==========================================
+ Coverage   96.08%   96.22%   +0.14%     
==========================================
  Files          29       30       +1     
  Lines        3319     3448     +129     
==========================================
+ Hits         3189     3318     +129     
  Misses        130      130              

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