Neural ODEs for FLOWPM and ODE TF correction

[dlanzieri]

In this notebook I try to find the correction terms to compensate for the PM approximations (like PGD in the n-body), following what has been done in this notebook.
The idea was using the same correction parameters found with the jax neural network, but using the TF neural network to inject them.
To do that, I followed this example. The problem is that the Dormand-Prince solver for the neural ODE takes forever.
@EiffL any idea of what could have gone wrong?

[EiffL]

This looks great, my best bet for what might be going wrong is that the model doesnt actually produces the correctt output, and so that completetly messes up the ODE.

Could you try to check whether your TF neural network outputs exactly the same thing as the jax version.

[EiffL]

I had a closer look and did a comparison with the jax code and it seems to work nicely! I'm very happy with this solution @dlanzieri :-)
https://colab.research.google.com/drive/1WsbCGPF5qWgYkczMQXvUexmB0J1l1u0n?usp=sharing

:-)
I don't know what might be the issue :-/

[EiffL]

One small trick maybe... I'm getting a factor 10 speed up in just running this little neural network if I do:

net = JaxNSFF(params, fun.apply)

@tf.function(jit_compile=True)
def apply_model(k, a):
  return net(k,a)

and then use this apply_model instead of the neural network directly

[dlanzieri]

Good news @EiffL ! I figured out what was going wrong. As you suggested, something wrong was completely messing up the ODE. Still, it wasn't the TF neural network outputs to be wrong, but rather the results of the forces computation.
The results seem very similar to the ones found in the Jax-notebook! :-)

[EiffL]

Fantastic Denise :-) now I'm really curious to see what our HOS and lensing statistics look like with this new simulation code

[dlanzieri]

Not sure this is the right Github issue, but I want to show you a very preliminary result:

For this plot seems that the angular power spectrum from the ODE implementation agrees with the TNGone better than the PGD implementation up to ell~1500, after this point there is a lack of power.
Following the results with different smoothing size:

The vertical lines define ell = 300 and ell=300, the Cls are computed by averaging over 200 maps.
@EiffL comment on that?

[EiffL]

hummmm I dont like the shape of this powerspectrum too much... It's worse than without correction?

Could you make a plot of the 3D PS at each snapshot of the simulation?

[dlanzieri]

Hello @EiffL. Here the plot with the 3D pk at each snapshot :

And, following the Cls, (I indicated with FASTPM our old implementation and with DLL ODE the new one).

So the Cls from our new implementation are definitely better with the correction , but yes, weird shape.

[EiffL]

Can I see your code somewhere? The most obvious reason why the power on the cls would go up on the right is because of shot noise, but it's surprising that the smoothing of the lens planes we used for the old flowpm is no longer appropriate here

[dlanzieri]

Sure! https://github.com/LSSTDESC/DifferentiableHOS/blob/main/scripts/compute_maps_ode_noia_nocorrr.py

[dlanzieri]

So, here how the HOS looks like with this new implementation.


To have an idea of the same results we got before with the PGD, I add here the plots:


So, I guess that we can use the new implementation since it seems to be not worse than PGD.