Implemented initial conditions with local non-Gaussianity

[tsfloss]

I modified the IC generator to allow for local non-Gaussianity, to this end I added a new cosmology parameter f_nl_loc and a new function to give the linear transfer function (linear_transfer) in the same spirit as the linear power spectrum (linear_power). The ICs are generated on a 2x larger grid, to make sure that squaring the field does not lead to aliasing effects, afterwards generating it downsamples to the target grid.

[adrianbayer]

Hey Yin, we're planning to do some field-level inference with fnl using this code, let us know what you think and if you're happy for us to merge!

[eelregit]

Hi Thomas and Adrian,
Thanks, this looks really good! Maybe we can further save memory and computation by Orszag's rule to pad by 1.5x. We have some utility for that in some other branch or stash, which I can dig out and add to this PR if you don't mind?
I see that you comment out the checkpoint decorator. If that's because of the JAX bug, you can call the function using all positional arguments to avoid that. There's also a workaround which I can add later.

[tsfloss]

Hey Yin, sorry for the late reply!

Indeed padding by 1.5x would already be good enough, but the padding steps do not seem to slow down the IC generation significantly. I leave that choice up to you, happy to help if needed!

Indeed I had to disable the checkpoint decorator because of some error, but I'm not sure what is the right thing to do there!

Best,
Thomas