Analytic sampling for conditional posterior instances trained with MDNs. #458

jnsbck · 2021-03-29T17:19:09Z

Hey,

I added code to analytically sample and condition DirectPosterior instances trained on MDNs. Wrapping a DirectPosterior instance that has been trained using an MDN in this way, replaces the .log_prob() and .sample() methods with analytical ones. This means it should be compatible with the rest of sbi.

Example

After training the posterior...

inference = SNPE(prior, density_estimator='mdn')
density_estimator = inference.append_simulations().train()
posterior = inference.build_posterior(density_estimator)
posterior.set_default_x(x=x_o)

...it can be wrapped, ...

analytic_posterior = MDNPosterior(posterior)

...and conditioned.

condition = torch.tensor([[theta_1, theta_2, 'nan']])
cond_posterior = analytic_posterior.conditionalise(condition)

Variables of interest are represented by 'nan's.

Then cond_posterior can be sampled and evaluated as before.

cond_posterior.sample((1,))
cond_posterior.log_prob(theta)

I am sure this could somehow be integrated into the DirectPosterior class directly as well, since the MDNPosterior inherits from it.
Let me know of any issues, I'd be happy to help integrating it into sbi. :)

…th MDNns.

michaeldeistler

Thanks a lot! I made a few preliminary comments. Nothing for you to do yet. Instead, we should re-write the MDN class in pyknos. This will avoid code repetition of sample and log_prob methods. I will take care of re-writing pyknos and will let you know. This might take a few days though.

sbi/utils/conditional_density.py

michaeldeistler · 2021-03-30T07:24:05Z

You can track progress here

…ings. Removed mentions.

…tion.

michaeldeistler

Okey, the PR is merged. We can now use .sample_mog() and .log_prob_mog() and replace basically the entire probability evaluations and sample functions by just calling these methods. Let me know in case I am missing something, and thanks again for taking over this endeavor!

sbi/utils/conditional_density.py

jnsbck · 2021-03-30T13:38:53Z

So much for "This might take a few days though." 😄
Sure. I'll keep you posted on the progress.

…terior.

jnsbck · 2021-03-31T22:24:25Z

sbi/utils/conditional_density.py

-            log_factor = torch.log(self.leakage_correction(x=self.default_x))
-            return torch.log(torch.sum(pdf, axis=1)) - log_factor
+        self.net.eval() # leakage correction requires eval mode
+        log_factor = torch.log(self.leakage_correction(x=self.default_x))


I think the leakage correction does not work for the conditioned MDNs, because it uses samples from self i.e. the full posterior rather then from the new, conditioned posterior.

sbi/utils/conditional_density.py

jnsbck

extract_and_transform_mog + __conditionalise can now be used to get new sets of mog logits, means etc. for an arbitrary condition which makes them compatible with sample_mog and log_prob_mog.

Next step: Introduce extract_and_transform_mog + __conditionalise into DirectPosterior.
Then insert them into sample_conditional together with sample_mog and sample_posterior_within_prior as well as into
log_prob_conditional together with log_prob_mog.

…niTb/sbi into analytic_mdn_conditioning

…sterior.

jnsbck

I have integrated everything into DirectPosterior as discussed and according to my tests it works.
Also... I just commited a bunch of files for review on accident. The only thing important for merging is the direct_posterior.py.

Looking forward to your remarks. Hope this round of review is the last :)

codecov-commenter · 2021-04-18T17:13:58Z

Codecov Report

Merging #458 (f989443) into main (9ed18ca) will decrease coverage by 0.99%.
The diff coverage is 12.00%.

@@            Coverage Diff             @@
##             main     #458      +/-   ##
==========================================
- Coverage   67.96%   66.96%   -1.00%     
==========================================
  Files          56       56              
  Lines        4117     4190      +73     
==========================================
+ Hits         2798     2806       +8     
- Misses       1319     1384      +65

Flag	Coverage Δ
unittests	`66.96% <12.00%> (-1.00%)`	⬇️

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

Impacted Files	Coverage Δ
sbi/utils/__init__.py	`100.00% <ø> (ø)`
sbi/utils/conditional_density.py	`61.98% <8.16%> (-36.63%)`	⬇️
sbi/inference/posteriors/direct_posterior.py	`69.40% <19.23%> (-11.51%)`	⬇️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9ed18ca...f989443. Read the comment docs.

jnsbck

Code should be ready for final review.

michaeldeistler

It's starting to look good, thanks! The main thing we still need are unit-tests. I suggest we write something akin to this. In your case, you would compare your samples from your conditional to samples from the ground truth via c2st

sbi/inference/posteriors/direct_posterior.py

jnsbck

So, I finally got round to implementing everything. I hope it still plays nicely with the new version of sbi that has come out in the meantime. I have also added something resembling a unit test, but I am not sure how to integrate it properly. Help would be greatly appreciated. I have just plugged it into an extra file for now. Looking forward to your comments.

Best,

Jonas

michaeldeistler

Thanks a lot! Super minor things only, we can merge it after this!

tests/testing_analytic_mdn_conditioning.py

sbi/utils/conditional_density.py

sbi/inference/posteriors/direct_posterior.py

jnsbck · 2021-07-06T16:19:16Z

I have just addressed your latest comments. Hope that it can be merged now. :)

michaeldeistler

This is great and will be super useful, thanks a lot for all your efforts!

michaeldeistler · 2021-07-09T07:16:37Z

Can be merged once merge conflicts are resolved.

jnsbck · 2021-07-18T18:40:09Z

I just resolved all conflicts, if it passes the tests, I guess you can merge it :)

Added analytic sampling of conditional posterior instances trained wi…

aaf7316

…th MDNns.

michaeldeistler self-assigned this Mar 29, 2021

michaeldeistler reviewed Mar 30, 2021

View reviewed changes

michaeldeistler mentioned this pull request Mar 30, 2021

Detach sample and log_prob from neural net sbi-dev/pyknos#17

Merged

jnsbck added 3 commits March 30, 2021 13:20

Adressing comments in sbi-dev#458: reformated code snippets in docstr…

0600280

…ings. Removed mentions.

Adressing comments in sbi-dev#458: Sampling nowvia Cholesky decomposi…

0978d7d

…tion.

Added Missing bracket.

630b7ca

michaeldeistler reviewed Mar 30, 2021

View reviewed changes

sbi/utils/conditional_density.py Outdated Show resolved Hide resolved

sbi/utils/conditional_density.py Outdated Show resolved Hide resolved

sbi/utils/conditional_density.py Outdated Show resolved Hide resolved

sbi/utils/conditional_density.py Outdated Show resolved Hide resolved

Intermediate step: sample_mog and log_prob now compatible with MDNPos…

d74fdae

…terior.

jnsbck commented Mar 31, 2021

View reviewed changes

sbi/utils/conditional_density.py Outdated Show resolved Hide resolved

jnsbck commented Mar 31, 2021

View reviewed changes

jnsbck added 4 commits April 6, 2021 00:33

Replaced with .

501bf68

Replaced with mulnormpdf with log_prob_mog.

2910bf0

Merge branch 'analytic_mdn_conditioning' of https://github.com/JBeckU…

aa4e1c5

…niTb/sbi into analytic_mdn_conditioning

Conditional sampling and evaluation for MDNs integrated into DirectPo…

06c3dc5

…sterior.

jnsbck commented Apr 18, 2021

View reviewed changes

jnsbck added 3 commits April 18, 2021 19:17

Removed accidentally added files from pull request

383c803

Small rewrites. Got rid of cholesky decomps.

a720a93

renamed vars.

a0e8800

jnsbck commented Apr 24, 2021

View reviewed changes

michaeldeistler requested changes Apr 28, 2021

View reviewed changes

michaeldeistler force-pushed the main branch from e70ff3a to 96d3a3b Compare May 4, 2021 07:12

famura mentioned this pull request May 5, 2021

Interface changes for SNPE-A #478

Merged

13 tasks

jnsbck added 3 commits May 13, 2021 22:04

Adressing critique: Fixed Docstrings. Restructure. Comments. Warnings.

660a9ce

looking for bug.

c7e0904

small fix.

ab20b6b

jnsbck added 2 commits June 13, 2021 15:00

BUGFIX import of DirectPosterior Type caused Problems.

794a74c

Added unittest. Works, but not sure how to integrate it properly.

293418f

jnsbck commented Jun 13, 2021

View reviewed changes

fixed docstring formatting. removed obsolete line.

449f1b7

michaeldeistler reviewed Jul 2, 2021

View reviewed changes

ran black, fixed formating, moved test to linearGaussian_snpe_test.

8d63816

michaeldeistler approved these changes Jul 9, 2021

View reviewed changes

Merge branch 'main' into analytic_mdn_conditioning

a39d186

jnsbck added 2 commits July 18, 2021 20:47

forgot arg for snpe_method in test_mdn_conditional_density.

43f0c51

fixed line that was accidently overwritten during resolve of conflicts.

f989443

jnsbck merged commit 00692fe into sbi-dev:main Jul 19, 2021

jnsbck mentioned this pull request Jan 25, 2022

first draft of FAQ to address #581 and #582 #595

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Analytic sampling for conditional posterior instances trained with MDNs. #458

Analytic sampling for conditional posterior instances trained with MDNs. #458

jnsbck commented Mar 29, 2021 •

edited

Loading

michaeldeistler left a comment

michaeldeistler commented Mar 30, 2021

michaeldeistler left a comment

jnsbck commented Mar 30, 2021

jnsbck Mar 31, 2021

jnsbck left a comment •

edited

Loading

jnsbck left a comment •

edited

Loading

codecov-commenter commented Apr 18, 2021 •

edited

Loading

jnsbck left a comment

michaeldeistler left a comment

jnsbck left a comment •

edited

Loading

michaeldeistler left a comment

jnsbck commented Jul 6, 2021

michaeldeistler left a comment

michaeldeistler commented Jul 9, 2021

jnsbck commented Jul 18, 2021

Analytic sampling for conditional posterior instances trained with MDNs. #458

Analytic sampling for conditional posterior instances trained with MDNs. #458

Conversation

jnsbck commented Mar 29, 2021 • edited Loading

Example

michaeldeistler left a comment

Choose a reason for hiding this comment

michaeldeistler commented Mar 30, 2021

michaeldeistler left a comment

Choose a reason for hiding this comment

jnsbck commented Mar 30, 2021

jnsbck Mar 31, 2021

Choose a reason for hiding this comment

jnsbck left a comment • edited Loading

Choose a reason for hiding this comment

jnsbck left a comment • edited Loading

Choose a reason for hiding this comment

codecov-commenter commented Apr 18, 2021 • edited Loading

Codecov Report

jnsbck left a comment

Choose a reason for hiding this comment

michaeldeistler left a comment

Choose a reason for hiding this comment

jnsbck left a comment • edited Loading

Choose a reason for hiding this comment

michaeldeistler left a comment

Choose a reason for hiding this comment

jnsbck commented Jul 6, 2021

michaeldeistler left a comment

Choose a reason for hiding this comment

michaeldeistler commented Jul 9, 2021

jnsbck commented Jul 18, 2021

jnsbck commented Mar 29, 2021 •

edited

Loading

jnsbck left a comment •

edited

Loading

jnsbck left a comment •

edited

Loading

codecov-commenter commented Apr 18, 2021 •

edited

Loading

jnsbck left a comment •

edited

Loading