32 fix snle

[maho3]

Fixes SBIRunner to allow for SNLE and SNRE models. Solves #32

[maho3]

Results of the new SNLE integration shows robustness for the toy example:

[maho3]

Added emcee samplers for sbi modules, and integrated it in with a modular interface in the validation step. Now, users can choose which sampler (emcee or sbi/Pyro) they want to use for each metric in the validation phase. This could be useful if you want high resolution sampling for a single example, but low resolution sampling for the entire test set (to check posterior coverage).

Here's an example posterior constraint from using emcee on an SNLE model.

This is built to work fluidly with the pydelfi version, but hasn't yet been tested. It should address #47 if everything works well.

[maho3]

Pydelfi version works with emcee! Here's an example:

[maho3]

This is another fairly extensive PR, so here's a summary of major changes:

• sbi's SNLE and SNRE models now work within our framework. To accomplish this, there have been major changes to SBIRunner. Specifically, in this area we build each method's neural network model and in this area we define different training procedures with which to train different methods. These had to be customized because of the different ways these models dealt with networks, proposals, and embedding nets. These might be able to be coded more elegantly than I have them now.
• There's a new module at ili.utils.samplers which provides emcee- and pyro-backed MCMC samplers for the likelihood and ratio estimation models. The emcee sampler interface works for both the pydelfi and sbi backends, whereas the pyro sampler only works for sbi. The interface allows you to customize the number of samples, chains, burn in steps, and thinning intervals. Different chains should sample in parallel, for maximum efficiency.
• The pydelfi models now have a .potential() function which is equivalent to their existing .log_posterior_stacked() function. This is to parallel with sbi's .potential() function, and is useful for uniting them in the emcee sampler interface.
• I've introduced _SampleBasedMetric as a base class for other validation metrics in ili.validation.metrics. This simply provides a sampling interface with the emcee and pyro samplers under ._build_sampler to avoid duplicating code. All other metrics now inherit from _SampleBasedMetric, but _BaseMetric still exists in case people want to build metrics without samples (i.e. using log probabilities).
• The examples/toy_sbi.py script now has argparse to allow users to try different methods (SNPE/SNLE/SNRE) without needing a new script. It was getting kind of clunky having different scripts for each method.
• Now, each validation metric that you specify in the config files now allows you to specify a custom sampling method (e.g. number of chains). This is relevant if you want many samples for PlotSinglePosterior, but few samples per input for RankStatistics.

And here are some minor changes:

• To make more sense with SNLE/SNRE models, the attribute SBIRunner.neural_posteriors has now been renamed to SBIRunner.nets in all of its references. This attribute stores an array of pytorch neural networks which can be used for amortized posterior estimation with SNPE, likelihood estimation with SNLE, or ratio estimation with SNRE. So, the label 'neural_posteriors' no longer necessarily applies. This change has also been reflected in DelfiRunner as well as all the example configurations in examples/configs/infer.
• To avoid confusion, 'val_loss' in the SBIRunner weighting scheme has been changed to 'val_logprob,' as that's what it technically is. the val_loss should be the negative of val_logprob.
• Some base classes without real functionality have been renamed with a leading underscore _ so that the autodoc ignores them.

Questions:

• If you use SNPE to train a model, it will always use direct sampling to generate samples. We currently don't support MCMC for SNPE, though sbi does. Should we include this functionality?

A non-exhaustive list of things to test:

• Make sure that the toy SNPE/SNLE/SNRE configurations work and give good posterior coverage.
• Make sure you can play with the sampler configurations and make them faster/slower.
• Try different SNLE and SNRE models (e.g. SNLE_A, SNLE_B, SNPE_A, SNRE_B, etc.) to make sure they are consistent (if you give them the right nets). I haven't exhaustively tested all of these yet.

[CompiledAtBirth]

(sbi backend only) On a new Python 3.10.12 environment with ili-summarizer also installed, SNRE, SNPE and SNLE work without issue in the toy_sbi example. Changing the _samplemethod attribute to emcee instead of "slice_np_vectorize" for the utils.samplers instances also runs without apparent issue. Next I'm going to look at the actual posterior coverage, try to change the configurations and check the type A/B/C instances (toy_sbi.py only for now).

[CompiledAtBirth]

Temporary workarounds with emcee sampler instances invoked in metrics.py:

• in TARP and RankStatistics with sbi backend, emcee already returns a numpy array, so .detach().numpy() returns errors
• in TARP, in the" call" method with emcee , posterior_samples had the wrong shape

[CompiledAtBirth]

[sbi backend] Hello! I do not understand the second major bullet point, namely that emcee- and pyro-backed MCMC samplers are available for SNLE or SNRE instances.

• _ili/inference/runnersbi/line 228: the _buildposterior method is called without argument for any trained instance (SNPE/SNRE/SNLE) in the SBIRunner class.
• _sbi/inference/snle/snlebase/line 287: said method has the "sample_with" argument allowing for MCMC, rejection sampling or Variational inference (idem in snpe_base)

1. Did I miss where we modify the sampling backends for SNRE/SNLE models? Currenlty ili.utils.samplers is not imported in ili/inference/*.py
2. If not, do we want to run with the default sbi behavior when dealing with trained likelihood(-ratio)s, or should we include the possibility to use MCMC(emcee or pyro) , VI or rejection as in sbi?

[maho3]

Hey Nicolas, thanks for looking into this.

Yes so here are the answers to some questions:

• The modules within ili/inference/.py are only designed to train models and then save them to file. They do not do any sampling inherently, all of that is done in the validation step (ili/validation/.py). You can see where we setup the samplers here (https://github.com/maho3/ltu-ili/blob/55a567d4762b54364a1fa581231c70b464037715/ili/validation/metrics.py#L65C9-L65C9).
• sbi's functionality is a little weird, due to the fact that its codebase has changed over time. Yes, the .build_posterior() method gives you a NeuralPosterior object that has a default sampler preset through the sbi package. However, the main thing we want it for is that it sets up a .potential function (for example) which is proportional to the log posterior probability and which we use for sampling later on in the validation step. This way, we can customize the sampler we want in each stage of the validation step, and simply refer that sampler to the potential function, instead of enforcing a single sampler configuration for all of the neural posterior samples.
• We currently don't support the Rejection or VI sampling procedures, but its something to include in a new ticket for the future. From my testing, they didn't perform significantly more efficiently than the MCMC sampling, and weren't necessarily more robust.

[CompiledAtBirth]

Hello, moving now to testing the toy examples with the pydelfi backend on a fresh Python==3.6.0 environment. Have you installed the version of this PR as such ? There is no compatible version of "mpsort" when installing ili-summarizer, looking into it.

[DeaglanBartlett]

When do you get this issue? If I run

git clone -b 32-fix-snle git@github.com:maho3/ltu-ili.git
git clone https://github.com/florpi/ili-summarizer.git
conda create -n ili-pydelfi python=3.6 -y
conda activate ili-pydelfi
pip install -e "ltu-ili[pydelfi]"
pip install -e ili-summarizer

then everything seems to install fine. Does this happen at installation stage or when you run a particular script?

[CompiledAtBirth]

It does work with that way actually, the fault was me being stubborn on using a venv ==3.6 instead of conda and asking for trouble (installation stage) Thank you

[CompiledAtBirth]

Question

• Currently TARP(_SampleBasedMetric) does not work with pydelfi/python3.6 backend. Do we want to deal with that later?

Answer about MCMC for SNPE (sbi) If sbi already allows MCMC sampling for trained posteriors from SNPE instances, we should include this possibility as well.

Some notes before finishing tests for review

• trying to install ltu-ili[sbi] in a brand new conda environment python ==3.11.4 environment fails; python==3.10.12 has no issue.
• it seems we only use the Dataset class from ili_summarizer. Even though, running toypydelfi.py , we get warnings such as "ili-summarizer/summarizer/init.py:33: UserWarning: Running without the following backends due to ModuleNotFoundError: nbodykit, jax, Corrfunc, kymatio"_, it does not seem we need these backends for ltu-ili classes. As per my previous question answered by @DeaglanBartlett , I actually tried installing ili-summarizer[backends] instead of just ili-summarizer , hence the incompatibility error at installation with pydelfi backend.
• With sbi backend, installing ili-summarizer[backends] works with Python 3.10.12. Only the nbodykit backends fails when importing summarizer, which is fixed by running

  python -m pip install git+https://github.com/cosmodesi/pycorr

  Given the previous point, it seems however we do not needs summarizer backends to work as of now.

[maho3]

Hey Nicolas,

Thanks a lot for the detailed review. Some notes:

• The sbi's DirectPosterior objects that come out of SNPE don't natively allow for mcmc sampling (see here). There's a way you can hack it together, by defining a custom PyroSampler to target the log_prob function, but it will be a bit of a headache. That being said, the native DirectSampling is as explorative as MCMC and way way faster, so I think its okay to just default to that behavior if you have a Direct Posterior
• Thanks for catching that dependency issue! I think its safe to say few people will be trying to use pydelfi and the ili-summarizer backends at the same time, so it is okay that we don't support it. That being said, I've added a line in INSTALL.md specifically warning about this conflicting dependency

[maho3]

Hey Nicolas,

• Yes, I noticed the same thing for the mlp version of SNLE, that it was quite rubbish even for default sbi. The resnet is quite reliable though, and I would encourage people to use it.
• Yes, I haven't tested MNLE extensively because on the backend it seems quite different from SNLE. Since SNLE is the main one that we want to compare with e.g. pydelfi I think its okay if it only works for that. Also, MNLE is primarily used for mixed data (e.g. continuous and discrete data simultaneously) which we don't often have for our use cases.
• The idea of adding customizable saving signatures is really good. I have added that as issue #73.

Additonal notes:

• Also, I couldn't get the nsf or made architectures to work in the cuda-integration runs, I think because their backends (nflows) aren't set up well for pytorch-cuda. This is not something we can fix, necessarily, but an important thing to note for future use.