search

astrodeepnet / sbi_experiments

Simulation Based Inference experiments
MIT License
3 stars 3 forks source link

Implement a simple WL toy likelihood model (e.g. fitting gaussian mixture to actual WL likelihood) #19

Open EiffL opened 2 years ago

EiffL commented 2 years ago

This issue is complementary to #18 which contains a simulator with latent variables which make it an implicit model. Here we will consider a case where the marginal likelihood is known, this is much easier and this is what we use in normal cosmological inference. I think this might be a safer bet for getting some nice results on a short timescale.

I'm thinking of writing a simple weak lensing likelihood model, trying to do inference over typical 6 cosmological parameters, so without nuisance params. This is already not trivial to fit this likelihood in 6 dimensions.

I started with this notebook: https://colab.research.google.com/drive/1Jui18xm9lol00q7dXm_dDYAG2MNqF6Yx?usp=sharing But I'll try to write a proper simulation tool.

On a 2D problem, it works shockingly well to try to learn the distribution simply by the score: image This is a posterior distribution on only 2 params, fitted using only 6 calls to the likelihood model. The fitted distribution is the shaded region in the background, the truth are the magma contours, the prior are the green coutours, data points are in white.

I also tried to push it to a DES Y1 problem with 21 parameters, but without success so far.

EiffL commented 2 years ago

Made a better notebook, to learn a 5d posterior here: https://github.com/astrodeepnet/sbi_experiments/blob/u/EiffL/lensing_simuilation/notebooks/lensing_sims/AnalyticLensingSimulator_SIREN_5d.ipynb

Works pretty well: image The blue contours are learned by fitting the posterior with 200 model evaluations, the purple contours are reference ones obtained by HMC, they take O(3e4) model evaluations!

EiffL commented 2 years ago

For now I'm using a SIREN to make this plot, and it's annoying because I can't directly sample from it, so I need to first fit the SIREN using the score to learn the log_prob, and then I run an HMC on it to get the samples.

It would be much better to be able to use a Flow here, but for that we need a flow that works in more that 2d.