Gaussian random coefficients from power spectra

[ntessore]

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.

[matt-graham]

Thanks for this @ntessore I've added some initial review comments below.

There is an issue with the reverse second-order grads that I don't fully understand, so I am not testing those here.

Out of interest what was is the error / issue you are having here?

[matt-graham]

Suggested change

	# always use reality=True, this could be real fields or E/B modes
	# always use reality=True, this could be real fields or E/B modes

'E/B modes' should ideally have some clarification here - from a bit of searching it looks like this may be a common term in cosmology settings, but as someone without a cosmology background I don't know what it means and I'd guess I'm representative of this in other potential users and developers from outside cosmology!

[ntessore]

Makes sense! I added a more generic comment now.

In any case, my idea was to only tackle flm with the reality condition here. Sampling flm without that condition is trickier due to the multiple spectra that are required for each field. Even then, the easiest way is to use this "E/B" decomposition into two fields with the reality symmetry, and later assemble the complex fields from there. Ideally, there would be a generic helper function for this conversion (complex <-> E/B), at which point the functionality could be added to generate_flm_from_spectra() as well.

[jasonmcewen]

Thanks @ntessore , I agree that sounds like a good approach, i.e. considering E/B spectra, since that is the typical convention in cosmology. Adding some more documentation on E/B fields as @matt-graham suggests is also a good idea.

[ntessore]

Out of interest what was is the error / issue you are having here?

The auto and numerical derivatives don't match. I suspect this has something to do with the random sampling being called multiple times.

The differentiability in light of the random sampling is a bit of a headache. It almost makes more sense to split this functionality up into two parts:

1. Sample random coefficients flm = generate_flm(...).
2. Scale with the power spectra flm = apply_power_spectra(spectra, flm).

This would make it straightforward to keep the random coefficients fixed and compute derivatives only wrt. the spectra.

[matt-graham]

The auto and numerical derivatives don't match. I suspect this has something to do with the random sampling being called multiple times.

Okay that make sense, I was a bit surprised that closing over the random number generator here by creating an instance with a fixed seed in the outer function and the differentiating through this worked at all, but in retrospect I guess it makes sense as JAX will just see the generated values when it traces the computation and they will then be baked in to the trace as constants.

It does seem possible for this sort of pattern to work with second-order derivatives - for example

import numpy as np
from jax.test_util import check_grads

def func(x):
    rng = np.random.default_rng(42)
    return x**2 * rng.standard_normal(x.shape)

check_grads(func, (np.array([1., 2.]),), 2)

passes without any errors, but I guess this doesn't necessarily carry through to more complex cases.

The differentiability in light of the random sampling is a bit of a headache. It almost makes more sense to split this functionality up into two parts:

1. Sample random coefficients `flm = generate_flm(...)`.

2. Scale with the power spectra `flm = apply_power_spectra(spectra, flm)`.

This would make it straightforward to keep the random coefficients fixed and compute derivatives only wrt. the spectra.

Agree that this sort of design would be better if supporting differentiating through the function is important and it would still also allow exposing a similar interface as currently by having a wrapper function which just calls these two functions.

It also makes it simpler to use the functionality in other contexts - for example a probabilistic programming / Bayesian inference setting. There if you for example wish to sample the random variates using MCMC you would need to pass in values directly into the function. I used a similar approach in a demo I worked on recently of using s2fft to reconstruct a signal on the sphere from noisy samples to allow interfacing with an MCMC package.

[jasonmcewen]

This is looking great @ntessore! Thanks for adding.

It would also be useful to add a function specialised to the scalar case only, i.e. just generating a scalar field from a single power spectrum (or similarly many fields at once for batch operation). That is a very common use case and then doesn't require any SVD.