N-body simulation Jacobian

[dlanzieri]

@EiffL I tried to compute the Jacobian for the N-body simulation here: https://github.com/LSSTDESC/DifferentiableHOS/blob/7c836494d75cd06174b8729291322c21177003df/scripts/jacobian_nbody.py. When I make the code run I get :

jacobian_nbody.py:81 compute_jacobian  *
    k, power_spectrum = pkl(final_field,kmin=0.01,dk=0.01,shape=final_field.shape,boxsize=np.array([FLAGS.box_size, FLAGS.box_size, 10 * FLAGS.box_size]))
/Users/dl264294/Desktop/github/DifferentiableHOS/DifferentiableHOS/pk.py:60 pk  *
    dig, Nsum, xsum, W, k, kedges = _initialize_pk(shape,boxsize,kmin,dk)
/Users/dl264294/Desktop/github/DifferentiableHOS/DifferentiableHOS/pk.py:34 _initialize_pk  *
    dig = tf.Variable(dig,dtype=tf.int32)
/Users/dl264294/.local/lib/python3.8/site-packages/tensorflow/python/ops/variables.py:262 __call__  **
    return cls._variable_v2_call(*args, **kwargs)
/Users/dl264294/.local/lib/python3.8/site-packages/tensorflow/python/ops/variables.py:244 _variable_v2_call
    return previous_getter(
/Users/dl264294/.local/lib/python3.8/site-packages/tensorflow/python/ops/variables.py:67 getter
    return captured_getter(captured_previous, **kwargs)
/Users/dl264294/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py:701 invalid_creator_scope
    raise ValueError(

ValueError: tf.function-decorated function tried to create variables on non-first call.`

Can you help me? I can't understand where the problem is, basically I've emulated the same steps in the compute_jacobian.py code https://github.com/LSSTDESC/DifferentiableHOS/blob/7c836494d75cd06174b8729291322c21177003df/scripts/compute_jacobian.py, changing the function computing the power spectrum and removing the ray tracing part. Also, I've noted https://github.com/DifferentiableUniverseInitiative/DHOD/blob/3ac05c8f57e391138a896c27edec43fe29507e97/diffhod/mock_observables/pk.py#L9, this " not differentiable!" is it not a problem for us?

[EiffL]

Ah yes ok... Well the error message is actually telling you what the problem is :-) The power spectrum code is using some tf.variable for some reason.... That's not a good idea. So you can remove the tf.variables

[dlanzieri]

@EiffL I added the latest version of the script computing the Jacobian of 3D power spectrum I'm working to. https://github.com/LSSTDESC/DifferentiableHOS/blob/08a59462ec65b1e014fad48525c8d4df50e7dd22/scripts/jacobian_nbody.py If you could take a look at it, it would help me, because in the following notebook https://github.com/LSSTDESC/DifferentiableHOS/blob/08a59462ec65b1e014fad48525c8d4df50e7dd22/notebooks/Nbody_jacobian.ipynb I found a significant mismatch between the Power Spectrum computed analytically and that one from simulation

[EiffL]

Ok so I must admit I'm a little lost in all the test we did ^^ and if I'm lost after 2 weeks, most likely in 6months from now there is no chance that we can recover the conclusions on our tests :-)

So, for the particular test of the jacobian of the N-body simulation, could you do the following so that we can close this issue:

• [ ] : Document here the conclusions of your study, with the plot showing the nice agreement between numerical and autodiff jacobians, a link to the code/notebook you used to get this plot, and some explanations of what you did.

• [ ] : Add a test in FlowPM for the jacobian of the N-body power spectrum. When we tests things like this, a nice way to close an investigation is to write a new test and add it to the battery of tests, so that we never have to wonder again whether FlowPM derivatives are ok ;-) Could you make a small Pull Request on the FlowPM project to add a test that compares finite difference jacobian to the autodiff one? It can be with a very low resolution simulation, that;s ok, have a look for instance at the tests in https://github.com/DifferentiableUniverseInitiative/flowpm/blob/d9f947eec7ce615f472552d9d316a6b13b9d8ff0/tests/test_tfpm.py#L130