feat: batched sampling and log prob methods. #1153

michaeldeistler · 2024-05-03T15:05:19Z

Edits by @janfb :
add methods to sample and evaluate given many observations, potentially vectorized.

It uses the capabilities of ConditionalDensityEstimators to broadcast sample and evaluate for multiple conditions.

not implementing it for MCMC for now, moved to #1176

manuelgloeckler · 2024-05-07T06:30:35Z

This currently won't work due to #1154.

… amortizedsample

…rs' into amortizedsample

…posteriors' into amortizedsample" This reverts commit 07084e2, reversing changes made to f16622d.

…from-different-posteriors' into amortizedsample

… reshapes in rejection

codecov · 2024-05-07T20:09:59Z

Codecov Report

Attention: Patch coverage is 75.94937% with 19 lines in your changes missing coverage. Please review.

Project coverage is 72.95%. Comparing base (3f722e3) to head (26444f7).
Report is 6 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1153      +/-   ##
==========================================
+ Coverage   72.91%   72.95%   +0.04%     
==========================================
  Files          93       93              
  Lines        7394     7459      +65     
==========================================
+ Hits         5391     5442      +51     
- Misses       2003     2017      +14

Flag	Coverage Δ
unittests	`72.95% <75.94%> (+0.04%)`	⬆️

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

Files	Coverage Δ
sbi/inference/posteriors/direct_posterior.py	`98.78% <100.00%> (+0.41%)`	⬆️
sbi/inference/posteriors/mcmc_posterior.py	`85.23% <100.00%> (+0.21%)`	⬆️
sbi/neural_nets/density_estimators/nflows_flow.py	`63.46% <ø> (ø)`
sbi/inference/posteriors/base_posterior.py	`86.04% <66.66%> (+1.70%)`	⬆️
sbi/inference/posteriors/importance_posterior.py	`55.38% <50.00%> (-0.18%)`	⬇️
sbi/inference/posteriors/rejection_posterior.py	`80.95% <50.00%> (-1.55%)`	⬇️
sbi/inference/posteriors/vi_posterior.py	`80.86% <50.00%> (-0.39%)`	⬇️
sbi/samplers/rejection/rejection.py	`88.00% <92.85%> (+0.94%)`	⬆️
sbi/inference/snpe/snpe_a.py	`62.60% <0.00%> (-0.95%)`	⬇️
sbi/inference/posteriors/ensemble_posterior.py	`50.00% <11.11%> (-3.54%)`	⬇️

... and 5 files with indirect coverage changes

This reverts commit 17c5343.

manuelgloeckler · 2024-05-08T08:21:29Z

Okey SNPE stuff works now.

MCMC stuff still needs to be done and tested.

janfb · 2024-06-13T14:46:01Z

ok, merging main into amortizedsample fixed the conflicts and removed all the tracked changes from the other PRs. 👍

janfb

great effort! 👏
I added a couple of comments for clarification.

sbi/inference/posteriors/direct_posterior.py

sbi/inference/posteriors/importance_posterior.py

sbi/inference/posteriors/mcmc_posterior.py

sbi/samplers/rejection/rejection.py

tests/density_estimator_test.py

tests/posterior_nn_test.py

manuelgloeckler · 2024-06-14T09:37:49Z

Oh wow thanks for fixing this mess :D. I am not sure what went wrong on the rebase (but probably merging main in should be preferred if there were many commits into main; I will keep that in mind).

Co-authored-by: Jan <janfb@users.noreply.github.com>

janfb

Looks good to me now! 🎉

sbi/inference/posteriors/mcmc_posterior.py

tests/density_estimator_test.py

janfb · 2024-06-18T10:57:44Z

tests/posterior_nn_test.py

+        pytest.param(SNLE_A, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(SNRE_A, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(SNRE_B, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(SNRE_C, marks=pytest.mark.xfail(raises=NotImplementedError)),


Happy with this test - but not sure if it belongs under posterior_nn_test -> instead should it be under mcmc_test?

I think this test is fine here (just checks that all posterior produce the same shapes).

I think in MCMC_test we should properly test batched MCMC sampling works as expected (as the current tests do not catch that the MCMC implementation did produce wrong results) i.e. should be part of #1176

janfb · 2024-06-18T12:00:10Z

there seems to be a failing tests during CD:

pytest tests/linearGaussian_snpe_test.py::test_c2st_multi_round_snpe_on_linearGaussian[snpe_a]

manuelgloeckler · 2024-06-18T12:41:06Z

@janfb I have already responded to the reshape (but you can only see it in code review ???, and I cannot answer the new comment you made).

In short, This reshape was not introduced in this PR. It was already there; I just moved it a line upwards and added a note (currently, MCMC will only work with a single event_dim anyway).

gmoss13

Great work @manuelgloeckler, thanks!

gmoss13 · 2024-06-18T13:01:16Z

tests/posterior_nn_test.py

+        pytest.param(SNLE_A, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(SNRE_A, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(SNRE_B, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(SNRE_C, marks=pytest.mark.xfail(raises=NotImplementedError)),


Happy with this test - but not sure if it belongs under posterior_nn_test -> instead should it be under mcmc_test?

gmoss13 · 2024-06-18T13:04:52Z

tests/ensemble_test.py

+    else:
+        samples = posterior.sample_batched((10,), ones(x_o_batch_dim, num_dim))
+
+    assert samples.shape == (10, x_o_batch_dim, num_dim), "Sample shape wrong"


gmoss13 · 2024-06-18T13:06:59Z

tests/density_estimator_test.py

-    samples = density_estimator.sample((1000,), condition=condition)
+    samples = density_estimator.sample(sample_shape, condition=condition)
+    samples = samples.reshape(-1, batch_dim, *input_event_shape)  # Flat for comp.
+
    samples_separate1 = density_estimator.sample(


nice! I would suggest in a future issue/PR to test this behavior for all batched sample/batched log_prob tests, as in some of them we currently only test the shape.

gmoss13 · 2024-06-18T13:08:05Z

sbi/inference/posteriors/ensemble_posterior.py

+        )
+        samples = []
+        for posterior_index, sample_size in torch.vstack(
+            posterior_indizes.unique(return_counts=True)


indices typo

sbi/inference/posteriors/mcmc_posterior.py

manuelgloeckler · 2024-06-18T14:00:01Z

Okay, the SNPE_A test fails because the SNPE_A posterior switches the "batching" behavior in different rounds

After first round > Sample from the proposal are of shape (500, 1, 2), which is correct as x_o has batch_dim 1
After second round > Sample from the proposal are of shape (1,500,2) which is incorrec.

This was because in the first round it sample from the DensityEstimator ( which followed the new convention). And after the first round it sampled using a custom implementation ( which followed the old convention).

manuelgloeckler · 2024-06-18T14:10:34Z

Main is currently failing at tests due to:

max_err = np.max(error)
>       assert max_err < 0.0027
E       assert 0.0033298< 0.0027

This seems like a random fluctation (due to some change).

Should I increase the tolerance for a bit to fix this? [Edit: I increase tolerance to 0.005, just for the record]

janfb · 2024-06-18T14:37:56Z

Main is currently failing at tests due to:
max_err = np.max(error)
>       assert max_err < 0.0027
E       assert 0.0033298< 0.0027
This seems like a random fluctation (due to some change).

Should I increase the tolerance for a bit to fix this? [Edit: I increase tolerance to 0.005, just for the record]

No, please keep the tolerance as it was. This will be fixed in #1177

This reverts commit 2aac705.

manuelgloeckler · 2024-06-18T14:40:53Z

Alright, I will merge after my local GPU/Slow tests are finished.

janfb · 2024-06-18T14:48:09Z

Alright, I will merge after my local GPU/Slow tests are finished.

Cool!

manuelgloeckler and others added 5 commits April 29, 2024 09:04

Base estimator class

17c5343

intermediate commit

705e9df

make autoreload work

07b53cd

amortized_sample works for MCMCPosterior

dd02e22

fixes current bug!

663185b

manuelgloeckler added the blocked Something is in the way of fixing this. Refer to it in the issue label May 7, 2024

Added tests

df8899a

manuelgloeckler removed the blocked Something is in the way of fixing this. Refer to it in the issue label May 7, 2024

manuelgloeckler and others added 14 commits May 7, 2024 17:25

batched_rejection_sampling

aa82aab

intermediate commit

00cdade

make autoreload work

cb8e4d8

amortized_sample works for MCMCPosterior

d64557f

Merge branch 'amortizedsample' of https://github.com/sbi-dev/sbi into…

f16622d

… amortizedsample

Merge branch '990-add-sample_batched-and-log_prob_batched-to-posterio…

07084e2

…rs' into amortizedsample

Revert "Merge branch '990-add-sample_batched-and-log_prob_batched-to-…

e54a2fb

…posteriors' into amortizedsample" This reverts commit 07084e2, reversing changes made to f16622d.

Merge branch '1154-density-estimator-batched-sample-mixes-up-samples-…

52d0e7e

…from-different-posteriors' into amortizedsample

sample works, try log_prob_batched

cd808d5

log_prob_batched works

f542224

abstract method implement for other methods

48a1a28

temp fix mcmcposterior

5a37330

meh for general use i.e. in the restriction prior we have to add some…

2b23e42

… reshapes in rejection

... test class

6362051

manuelgloeckler added 5 commits May 8, 2024 07:26

Revert "Base estimator class"

294609d

This reverts commit 17c5343.

removing previous change

99abbb1

removing some artifacts

ef9e99c

revert wierd change

5eb1007

docs and tests

82127ab

janfb requested changes Jun 13, 2024

View reviewed changes

janfb changed the title ~~feat: sample method for batched conditions~~ feat: batched sampling and log prob methods. Jun 13, 2024

janfb mentioned this pull request Jun 14, 2024

categorial mass estimator and mnle fails for 2D density estimator test #1172

Closed

manuelgloeckler and others added 5 commits June 14, 2024 11:52

Apply suggestions from code review

93ca374

Co-authored-by: Jan <janfb@users.noreply.github.com>

Adding comment about squeeze

86f3531

Update sbi/inference/posteriors/direct_posterior.py

2a5f357

Co-authored-by: Jan <janfb@users.noreply.github.com>

fixing formating

79273a2

reverting MCM posterior changes

7b23d60

janfb mentioned this pull request Jun 18, 2024

Allow sampling the posterior given different x (batched) #944

Closed

manuelgloeckler added 2 commits June 18, 2024 10:03

xfail mcmc tests

d4f9e46

Exclude MCMC from ensamble batched_sample test

6798c97

manuelgloeckler requested review from gmoss13 and janfb June 18, 2024 08:55

janfb approved these changes Jun 18, 2024

View reviewed changes

gmoss13 approved these changes Jun 18, 2024

View reviewed changes

SNPE_A Bug fix

b1724a5

typo fix

a6f4845

preamtive main fix

2aac705

Revert "preamtive main fix"

26444f7

This reverts commit 2aac705.

manuelgloeckler merged commit 4951439 into main Jun 18, 2024
6 checks passed

janfb deleted the amortizedsample branch June 20, 2024 10:31

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: batched sampling and log prob methods. #1153

feat: batched sampling and log prob methods. #1153

michaeldeistler commented May 3, 2024 •

edited by janfb

Loading

manuelgloeckler commented May 7, 2024

codecov bot commented May 7, 2024 •

edited

Loading

manuelgloeckler commented May 8, 2024

janfb commented Jun 13, 2024

janfb left a comment

manuelgloeckler commented Jun 14, 2024

janfb left a comment

janfb Jun 18, 2024

gmoss13 Jun 18, 2024

manuelgloeckler Jun 18, 2024

janfb commented Jun 18, 2024

manuelgloeckler commented Jun 18, 2024

gmoss13 left a comment

gmoss13 Jun 18, 2024

gmoss13 Jun 18, 2024

gmoss13 Jun 18, 2024

gmoss13 Jun 18, 2024

manuelgloeckler commented Jun 18, 2024

manuelgloeckler commented Jun 18, 2024 •

edited

Loading

janfb commented Jun 18, 2024 •

edited

Loading

manuelgloeckler commented Jun 18, 2024

janfb commented Jun 18, 2024

feat: batched sampling and log prob methods. #1153

feat: batched sampling and log prob methods. #1153

Conversation

michaeldeistler commented May 3, 2024 • edited by janfb Loading

manuelgloeckler commented May 7, 2024

codecov bot commented May 7, 2024 • edited Loading

Codecov Report

manuelgloeckler commented May 8, 2024

janfb commented Jun 13, 2024

janfb left a comment

Choose a reason for hiding this comment

manuelgloeckler commented Jun 14, 2024

janfb left a comment

Choose a reason for hiding this comment

janfb Jun 18, 2024

Choose a reason for hiding this comment

gmoss13 Jun 18, 2024

Choose a reason for hiding this comment

manuelgloeckler Jun 18, 2024

Choose a reason for hiding this comment

janfb commented Jun 18, 2024

manuelgloeckler commented Jun 18, 2024

gmoss13 left a comment

Choose a reason for hiding this comment

gmoss13 Jun 18, 2024

Choose a reason for hiding this comment

gmoss13 Jun 18, 2024

Choose a reason for hiding this comment

gmoss13 Jun 18, 2024

Choose a reason for hiding this comment

gmoss13 Jun 18, 2024

Choose a reason for hiding this comment

manuelgloeckler commented Jun 18, 2024

manuelgloeckler commented Jun 18, 2024 • edited Loading

janfb commented Jun 18, 2024 • edited Loading

manuelgloeckler commented Jun 18, 2024

janfb commented Jun 18, 2024

michaeldeistler commented May 3, 2024 •

edited by janfb

Loading

codecov bot commented May 7, 2024 •

edited

Loading

manuelgloeckler commented Jun 18, 2024 •

edited

Loading

janfb commented Jun 18, 2024 •

edited

Loading