s2fft
s2fft copied to clipboard
Gaussian random coefficients from power spectra
This PR adds a new function s2fft.utils.signal_generator.generate_flm_from_spectra() that generates a stack of Gaussian $f_{lm}$ from a stack of power spectra. Closes #201.
The function is implemented in terms of the existing generate_flm(), which gains a new size= parameter for generating more than one set of coefficients at a time.
The generated flm are then simply scaled by the matrix square root of the provided stack of power spectra. This is implemented here in terms of the SVD, not the Cholesky decomposition, so that it works more reliably with positive semi-definite spectra (e.g. E/B mode spectra for spin-2 fields where the first two entries vanish).
The most difficult part of the PR was the testing: here I am sampling a set of Gaussian random fields, then construct the Gaussian covariance and make sure that the chi2 of the realised fields is within 3 sigma of the expected value. Happy for any suggestions for a better test!
Finally, the whole thing is implemented using the Array API, so that it works with both numpy and jax (and both are tested). For jax, the first-order grads wrt to the spectra are checked as well. There is an issue with the reverse second-order grads that I don't fully understand, so I am not testing those here.
I am marking this as draft for the time being, please share your thoughts on the API and implementation.