Integrate new density estimator interface into SBI

[gmoss13]

Addresses #957. Should implement fully for all methods before merging.

NOTE:

• Also removed support for torch < 1.8 distributions support (as torch >= 1.8 is required)
• Started preparations to support batch_shape for x_o.
• MultipleIndependent prior support check on dim=1 -> dim=-1

[manuelgloeckler]

Great work @gmoss13 ! I will join this branch and try to update the NPE classes.

Edit: Are the changes up to date, then I would commit my change for NPE?

[gmoss13]

Added some basic functionality for incorporating Zuko flows with embedding nets, works with SNLE. A couple of questions:

1. For NFlowsFlow, the embedding net is under density_estimator.net._embedding_net, whereas for ZukoFlow it is under density_estimator._embedding_net. This is a bit confusing. The easiest solution would be to add an ._embedding_net property to all DensityEstimators, which for NFlowsFlow would just return self.net._embedding_net. What do you think @manuelgloeckler?

2. Do we want to add separate builders for ZukoFlow estimators, or add it as an option to the current builders. E.g. build_maf(..., backend="zuko")? This looks nicer but is a bit annoying as the nflows and zuko take different kwargs.

[michaeldeistler]

Regarding 2), I would prefer build_maf_zuko() over build_maf(..., backend="zuko"). Let's make this very explicit.

[codecov]

Codecov Report

Attention: Patch coverage is 77.60417% with 43 lines in your changes are missing coverage. Please review.

Project coverage is 76.37%. Comparing base (3c1bb5a) to head (b4b00ba).

Files	Patch %	Lines
sbi/inference/snpe/snpe_a.py	30.76%	18 Missing ⚠️
sbi/inference/snpe/snpe_c.py	18.18%	9 Missing ⚠️
sbi/inference/posteriors/direct_posterior.py	71.42%	4 Missing ⚠️
sbi/utils/get_nn_models.py	20.00%	4 Missing ⚠️
sbi/neural_nets/flow.py	93.93%	2 Missing ⚠️
sbi/analysis/conditional_density.py	83.33%	1 Missing ⚠️
.../inference/potentials/posterior_based_potential.py	85.71%	1 Missing ⚠️
sbi/neural_nets/density_estimators/base.py	75.00%	1 Missing ⚠️
sbi/neural_nets/density_estimators/nflows_flow.py	75.00%	1 Missing ⚠️
sbi/neural_nets/density_estimators/zuko_flow.py	97.77%	1 Missing ⚠️
... and 1 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #979      +/-   ##
==========================================
+ Coverage   76.30%   76.37%   +0.07%     
==========================================
  Files          83       84       +1     
  Lines        6398     6507     +109     
==========================================
+ Hits         4882     4970      +88     
- Misses       1516     1537      +21     

Flag	Coverage Δ
unittests	76.37% <77.60%> (+0.07%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[manuelgloeckler]

In the integration process, I also noticed that we have to expose some more properties (currently, I temporarily fixed it only for NFlows, to first get all tests to run through). But we indeed should in no line have to call self._neural_net.net (which I currently have to do in snpe_a/c).

1. I think we need to find a common interface for quite a few properties @gmoss13 :

• Yes, we have to impose a few arguments e.g., embedding_net. I would add this as an optional property which by default returns None.
• We also often also use transforms, especially the last transform is often necessary (i.e. to map to a certain support or the z_score_transform. This can also be an optional argument that, by default returns an empty list.
• And further, we sometimes use base_dist i.e. the base measure for the flow.

2. For SNPE_A and SNPE_C an MDN estimator has a special meaning and might require additional properties. Currently the "correction" is also implemented as a DensityEstimator Wrapper.
3. We have to decide how/if/when we allow a batch_shape on x_o. There is now a difference between an x_o of shape (1,d) and one of (d,) as e.g. the density estimator would then sample (1, sample_shape, d).

[gmoss13]

We also often also use transforms, especially the last transform is often necessary (i.e. to map to a certain support or the z_score_transform. This can also be an optional argument that, by default returns an empty list.

Can we get around having a transforms property? Not all density estimators will be flows. Otherwise I agree that we should add common properties as you suggested.

For SNPE_A and SNPE_C an MDN estimator has a special meaning and might require additional properties. Currently the "correction" is also implemented as a DensityEstimator Wrapper.

Since this is MDN-specific, I wonder if we can leave it as it is?

[manuelgloeckler]

Can we get around having a transforms property? Not all density estimators will be flows. Otherwise I agree that we should add common properties as you suggested.

Agree, but e.g. MDNs will always be wrapped as a flow that either has an identity_transform or z_score_transform. The most general would probably be to add an abstract child class i.e. FlowBasedDensityEstimator, which has these additional attributes.

Since this is MDN-specific, I wonder if we can leave it as it is?

Probably.

[gmoss13]

Also added now RatioEstimator wrapper - currently not used. I've assumed that there can only be one x passed to the classifier (potentially the same x repeated). I don't know what the expected behaviour should be if one passes several values of x to the classifier. I also wonder if we want to generalise the classifier type? We define our classifiers to take as input a list of [theta,x]. So an externally defined classifier would need to also use this interface to work.

[janfb]

Great effort! 👏

I added lots of questions and comments. happy to discuss online tomorrow.

[janfb]

Looks great!
Added a couple of minor comments, and one important one about the loss vs log_prob.

[manuelgloeckler]

I will run all slow tests locally over launch to check if they also pass.

[manuelgloeckler]

Okay, all 204 slow tests did pass.

[janfb]

Requires a rebase on main before merging.

Given the long commit history with merge commits etc, I suggest to squash the commits into one big commit. Alternatively, you could try to reorganize and summarize the changes into a smaller number of commits during rebasing.

[gmoss13]

Requires a rebase on main before merging.

Given the long commit history with merge commits etc, I suggest to squash the commits into one big commit. Alternatively, you could try to reorganize and summarize the changes into a smaller number of commits during rebasing.

I would go for the one big commit. Let's wait to see if @michaeldeistler has any comments before merging.

[michaeldeistler]

This is really amazing work @manuelgloeckler and @gmoss13, thanks so much! I have made a few comments below, but most of them are to create follow-up issues and add a few comments. Once addressed this is good to go!