PGD scheme

[dlanzieri]

Add functions for the PGD implementation

[EiffL]

Thanks @dlanzieri for opening this PR :-) let me add a little bit of context. This PR aims to address issue #89 and provide the tools for PGD sharpening of flowpm simulations.

We still have a number of open issue:

• [ ] : compatibility with PGD parameters from madlens. Apparently we have different conventions somewhere because when applying the parameters obtained by biwei it doesn't give the expected result.

• [ ] : differentiability with respect to PGD parameters. We will probably want to optimize or even sample over PGD params but I think right now the PGD kernel is in simple python.

@modichirag do you have time for a review and try to check these two boxes? To merge this PR we want to at least have a set of PGD params that works as expected

[modichirag]

I tried fitting for the parameters in the last commit fc92237 The only thing we basically needed to do was make the kernel function tf expression, and so I moved it into tfpm instead of kernels (since kernels in kernels.py were simple np functions, but that can be moved back)

We need to figure out the loss function for fitting the parameters. https://arxiv.org/pdf/1804.00671.pdf Appendix B in Biwei's paper suggests MSE of power spectrum, weighted by reference power spectrum. One place where we might differ is that our reference power spectrum comes from jax while for the paper it comes from higher resolution simulation and hence does not have cosmic variance. This matters on the largest scales, but it is possible that the response to PGD parameters is 0 due to suppressing large scales with the PGD kernel and so we might not need to worry about it. We can also weigh it by 1 - k/knyquist to suppress smaller scales than we need i.e. k>2

This is what PGD seems to do for L=200, N=128, B=1, T=40 steps. I am showing transfer function with jax for different iterations in optimization for a=1 snapshot

This is what it seems to do for L=200, N=128, B=1, T=10 steps.

I also asked Yu Feng. While MADlens uses PGD for post processing, c-fastpm uses it as a part of the simulation. We can try to do that if we want. To keep things simple, we can start with just fitting independent/different parameters at all redshifts before tying ourselves to some interpolating function. In his paper, Biwei gives parametric form in Eq. 4.1-4.5, but I am not sure how useful it is given that it is only for resolutions > 1.

[modichirag]

The last commit b52c8b9 adds online PGD in tfpm_pgd.py. Online pgd means that after every drift step, we do a PGD update. Currently I have set it up so that it takes in the reference power spectra at all intermediate points and fits the PGD parameters independently at every redshift to match this reference spectra. The results are sensible, so that is good.

This is how it looks for 400 Mpc/h box, 10 steps, B=1, nc=128. (My 200 Mpc/h runs did not match the spectra on the large scales even for normal N-Body, which is worrying...? This was what Denise was seeing for 50Mpc/h box)

It needs to be compared wether it's better to do online fitting or just fit each snapshot separately in post processing, as the previous commit was doing. I am hoping Denise can look into that?

Also, as Francois pointed out, we are comparing CIC-smoothed spectra with reference spectra. Maybe we can smooth the reference spectra somehow? Otherwise maybe we can run some test simulations on fastpm with high resolution and just save their snapshots on the low resolution of interest where we want to fit the parameters and use them to fit PGD parameters. This should not take a lot of effort and is probably easier thing in the short run.

[dlanzieri]

Thank you @modichirag for this review! I'll sure take a look at it. Just one thing, maybe I'm missing something, if the last plot shows the ratio between our matter power spectrum and the reference power spectrum, what's the difference between the dashed and the solide lines?

[EiffL]

Ok I made some modifs to respond to my own comments above, and saved them there: #92