Add function to obtain the MAP

[EiffL]

This issue is to track the development of a script to obtain the Maximum A Posterior (MAP) solution using neural score matching.

It's essentially a matter of implementing a gradient descent procedure, using an existing JAX optimizer for instance.

The deliverable for this issue are:

• : a script MAP_solution.py similar to scripts/sample_maps.py.

[b-remy]

At the moment, the procedure to obtain the MAP is documented here:
https://github.com/CosmoStat/jax-lensing/blob/training-script/notebooks/MAP.ipynb

[EiffL]

Here are some thoughts on annealing for MAP:

• In general, try to look at the evolution of the Power Spectrum with the number of iteration
• Look at the evolution of the Power Spectrum as a function of iterations, after a sudden drop in temperature.
• At the minimum use ADAM to have an automatic scaling of the step size, maybe L-BFGS (https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/optimizer/lbfgs_minimize)

These points above can be done in a notebook, and then implemented in the new reconstruction script

In the Gaussian case:

• We should be able to compare to the Wiener Filter @NiallJeffrey is implemeting
• We can compare to the mean of Gaussian samples

[b-remy]

Here are some results for the MAP with full prior computed with gradient descent, where the temperature is decreased at every step, changing the step size as:

temp_list = jnp.arange(temp_init, temp_final, step=(temp_final-temp_init)/n_steps)

for i in range(n_steps):
  step_size = initial_step_size * (temp_list[i] / temp_final)**2
  gradient_step(...)

and power spectrum compared to the ground truth convergence and the MAP with a gaussian prior

[EiffL]

This is still and open issue, we know how to do this with SDE interpretation though

[EiffL]

So.... we now know how to do it properly in principle... using the ODE to obtain the map log likelihood. While this is very cool in principle, we are thinking of flagging this as out of scope for this paper.

Presumably though..... the MAP should be better point estimate than the mean......