Density Estimator Class

[gmoss13]

fix #932

Goals:

• General density estimator class to wrap around different models (Flows, MDNs, ...).
• Training Objectives for SNPE, SNRE, SNLE taking density estimators as args
• Reduce code repetition of .train() methods

[gmoss13]

@michaeldeistler @janfb this is bare bones at the moment, could use some feedback. Specifically, how much do we want DensityEstimator to handle? Should it already implement different potentials and samplers?

[michaeldeistler]

Overall, I like the structure. I believe that a more natural file location would be sbi/neural_nets/density_estimators/....

Beyond that I made many detailed comments below. Most important would be:

1. I do not think that the FlowDensityEstimator is working right now because it does not follow the nflows API.
2. Please write of the FlowDensityEstimator, even if they are only API tests.

[michaeldeistler]

As for your question: I think the scope of the DensityEstimator is just fine. It should not handle potentials and samplers.

[gmoss13]

Thanks for the comments! I applied most of the suggestions, but have some questions regarding a couple of your comments, where I could use some extra clarification.

[michaeldeistler]

New PR:

• utils.posterior_nn() should return wrapped density estimator --> add tests
• rewrite .train() to use .loss and MCMCPosterior.log_prob() should use DensityEstimator.log_prob() and DirectPosterior.sample() should use DensityEstimator.sample(), etc...

[gmoss13]

Thanks again for the comments! New commit should resolve these.

New PR:

• utils.posterior_nn() should return wrapped density estimator --> add tests
• rewrite .train() to use .loss and MCMCPosterior.log_prob() should use DensityEstimator.log_prob() and DirectPosterior.sample() should use DensityEstimator.sample(), etc...

Once this PR is is merged, we can work on these!

[codecov]

Codecov Report

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

Project coverage is 76.08%. Comparing base (3aeb775) to head (0853270).

Files	Patch %	Lines
sbi/neural_nets/density_estimators/base.py	72.72%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #952      +/-   ##
==========================================
+ Coverage   76.02%   76.08%   +0.06%     
==========================================
  Files          80       83       +3     
  Lines        6332     6361      +29     
==========================================
+ Hits         4814     4840      +26     
- Misses       1518     1521       +3     

Flag	Coverage Δ
unittests	76.08% <89.65%> (+0.06%)	⬆️

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.

[janfb]

great work! 👏

Added another round of comments. Happy to discuss later when we meet.

[manuelgloeckler]

Great, thanks for finishing it. Let's get it merged.

Outdated.

[janfb]

looks good, thanks a lot!

I think the unticked goals in the PR description can be moved to future PRs / separate issues.

[francois-rozet]

Hello, just a comment that it might be worth adding a sample_and_log_prob or rsample_and_log_prob method to the DensityEstimator interface as some flows cannot (or very slowly) get the density of arbitrary values, but can get the density of their own samples. This is the case of inverse autoregressive flows (IAFs, basically reversed MAFs), which are very fast to sample from but whose density is slow to evaluate. This could be beneficial for algorithms that rely on samples and their density (almost all SNPE methods).

By the way, Zuko flows already have an rsample_and_log_prob method 😉