astro-informatics / s2wav

Differentiable and accelerated wavelet transform on the sphere with JAX
https://astro-informatics.github.io/s2wav/
MIT License
12 stars 0 forks source link

Add vectorised filters and specific tiling functions #42

Closed CosmoMatt closed 1 year ago

CosmoMatt commented 1 year ago

This PR just vectorises the functions which generate directional and axisymmetric filters. The rate limiting step for the directional functions is the computation of the kernels k_lam however this relates to the quadrature integration we wish to replace so for now this implementation should be sufficient.

Basic %timeit tests show minimal improvement in speed which isn't too surprising as the implementation evaluates many unnecessary values (for el values outside of the j wavelet support) to be easily vectorised. Moreover, the filter generating functions are a very minor contribution to the overall transform wall-time so probably not wasting effort here (for now).

CosmoMatt commented 1 year ago

@alicjapolanska @JessWhitney also important to note that I've moved from storing harmonic coefficients as L^2 1D coefficients to L, 2L-1 2D coefficients. This double the memory overhead but is consistent with the indexing in S2FFT which we should try to match.

This will, however, make certain loops within the top-level transform (e.g. synthesis) a little more awkward until S2FFT is public (hopefully soon).