Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Contrastive Neural Ratio Estimation #787

Merged
merged 21 commits into from
Dec 19, 2022
Merged

Conversation

bkmi
Copy link
Contributor

@bkmi bkmi commented Nov 22, 2022

Dear SBI,

I would like to add our algorithm NRE-C to the SNRE section. I created the following checklist which I will update as the draft nears completion.

Let me know if there's anything else I can do!
Ben

@janfb
Copy link
Contributor

janfb commented Nov 22, 2022

Dear @bkmi
that's great, thank you for contributing NRE-C!
At the moment, the tests are failing because of black. So please run black, isort locally and push again.

@janfb
Copy link
Contributor

janfb commented Nov 22, 2022

black and isort are passing now 👍
However, pyright is complaining about some type erros. You can see them locally by running pyright, or on github by expanding the test logs:

Run pyright sbi
[7](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:8)
Loading configuration file at /home/runner/work/sbi/sbi/pyrightconfig.json
[8](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:9)
Assuming Python platform Linux
[9](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:10)
Searching for source files
[10](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:11)
Found 85 source files
[11](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:12)
pyright 1.1.280
[12](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:13)
/home/runner/work/sbi/sbi/sbi/inference/snre/snre_base.py
[13](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:14)
  /home/runner/work/sbi/sbi/sbi/inference/snre/snre_base.py:242:56 - error: Argument expression after ** must be a mapping with a "str" key type (reportGeneralTypeIssues)
[14](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:15)
  /home/runner/work/sbi/sbi/sbi/inference/snre/snre_base.py:272:60 - error: Argument expression after ** must be a mapping with a "str" key type (reportGeneralTypeIssues)
[15](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:16)
/home/runner/work/sbi/sbi/sbi/inference/snre/snre_c.py
[16](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:17)
  /home/runner/work/sbi/sbi/sbi/inference/snre/snre_c.py:165:31 - error: "float" is not iterable
[17](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:18)
    "__iter__" method not defined (reportGeneralTypeIssues)
[18](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:19)
  /home/runner/work/sbi/sbi/sbi/inference/snre/snre_c.py:180:16 - error: Expression of type "tuple[float, float]" cannot be assigned to return type "float"
[19](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:20)
    "tuple[float, float]" is incompatible with "float" (reportGeneralTypeIssues)
[20](https://github.com/mackelab/sbi/actions/runs/3524072067/jobs/5909002539#step:7:21)
4 errors, 0 warnings, 0 informations

@bkmi
Copy link
Contributor Author

bkmi commented Nov 22, 2022

The tests are extremely slow. Any hints here? Is there a way to only run the ratio-relevant ones?

I could run only the tests that I changed, but I did slightly alter the RatioEstimator in snre_base, so maybe I cannot get away with it.

@janfb
Copy link
Contributor

janfb commented Nov 22, 2022

yes, you can run specific tests like this (in the sbi folder):

pytest tests/linearGaussian_snre_test.py::test_function

or run all tests in that file (exlcuding marked gpu and slow tests):

pytest -m "not slow and not gpu" tests/linearGaussian_snre_test.py

@bkmi
Copy link
Contributor Author

bkmi commented Nov 23, 2022

I found a bug in RatioEstimator while running the tests #788.

If you like my proposal I will include the fix in this PR.

@bkmi
Copy link
Contributor Author

bkmi commented Nov 23, 2022

I found a bug in RatioEstimator while running the tests #788.

If you like my proposal I will include the fix in this PR.

We decided to make this a different PR.

@codecov-commenter
Copy link

codecov-commenter commented Nov 23, 2022

Codecov Report

Merging #787 (4a535e5) into main (6cdf678) will increase coverage by 0.17%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #787      +/-   ##
==========================================
+ Coverage   74.62%   74.80%   +0.17%     
==========================================
  Files          79       80       +1     
  Lines        6144     6187      +43     
==========================================
+ Hits         4585     4628      +43     
  Misses       1559     1559              
Flag Coverage Δ
unittests 74.80% <100.00%> (+0.17%) ⬆️

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

Impacted Files Coverage Δ
sbi/inference/posteriors/direct_posterior.py 98.21% <ø> (ø)
...inference/potentials/likelihood_based_potential.py 100.00% <ø> (ø)
.../inference/potentials/posterior_based_potential.py 100.00% <ø> (ø)
sbi/inference/potentials/ratio_based_potential.py 100.00% <ø> (ø)
sbi/inference/snre/snre_b.py 100.00% <ø> (ø)
sbi/utils/sbiutils.py 79.76% <ø> (ø)
sbi/inference/__init__.py 100.00% <100.00%> (ø)
sbi/inference/snre/__init__.py 100.00% <100.00%> (ø)
sbi/inference/snre/snre_base.py 94.87% <100.00%> (ø)
sbi/inference/snre/snre_c.py 100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

@bkmi bkmi marked this pull request as ready for review November 24, 2022 23:27
@bkmi
Copy link
Contributor Author

bkmi commented Nov 24, 2022

I think this pull request is ready to be reviewed!

Copy link
Contributor

@janfb janfb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks again for adding this to SBI!
looks great already, I added some minor comments. One general question I have concerns the sequential version of NRE_C. If it is not exact this should be clearly stated in the docstrings and we should add a warning or even an error when it is used in sequential mode, informing about the practical consequences.
But maybe I am missing something here, I am happy to discuss this in more detail.

EDIT: Slow and gpu tests are passing!

sbi/inference/potentials/ratio_based_potential.py Outdated Show resolved Hide resolved
sbi/inference/snre/__init__.py Show resolved Hide resolved
sbi/inference/snre/snre_b.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_base.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_base.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_base.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_c.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_c.py Show resolved Hide resolved
sbi/inference/snre/snre_c.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_c.py Show resolved Hide resolved
@bkmi
Copy link
Contributor Author

bkmi commented Nov 25, 2022

I will comment more carefully soon, but just to be clear... It is not that NRE-C is anymore inaccurate than any other sequential method, rather the guarantee that, at optimum, $r(x \mid \theta) = \frac{p(x \mid \theta)}{p(x)}$ does not hold. Instead, there would be a multiplicative term which depends on $x$.

In NRE-A this optimum holds (only in the first round). In NRE-B, this never holds. NRE-C is just like NRE-A in this regard. NRE-B is the default method but does not carry the warning.

We could add the warning to all NRE methods, but I don't think that is a good answer. After all, this does not affect posteriors but rather ways to test them. Sequential methods generally make diagnostics on the posterior more difficult (or impossible). If we added a warning here, it is my opinion that a general statement that sequential methods are not easily testable would be more appropriate and honest.

Perhaps the best answer would be for me to be very specific in the comment about what is meant regarding the sequential behavior of NRE-C, namely that an unknown multiplicative term in $x$ appears when doing sequential training?

@janfb
Copy link
Contributor

janfb commented Nov 25, 2022

I see, thanks for clarifying!
I agree, it would be good to point this out in the docstring of SNRE_C (and also update the ones for the other NRE methods), e.g., sth like that the likelihood-to-evidence ratio approximation is exact only in the first round, but that this does not affect the accuracy of the posterior samples, with a pointer to the paper for more details.

@bkmi
Copy link
Contributor Author

bkmi commented Dec 7, 2022

Thanks for your patience, NeurIPS delayed me on this.

(There are still TODOs)

@michaeldeistler
Copy link
Contributor

Jumping in here regarding warnings for unnormalized posteriors: the way this is currently handled is that the MCMCPosterior (regardless of whether it is created by NRE, NLE, or sequential methods) warns if one calls .log_prob() here. To avoid the warning, the user should use .potential() for MCMCPosteriors.

I suggest to describe in the docstring of NRE_C that its .potential() will be incouraged to be normalized, but to not issue further warnings.

Copy link
Contributor

@michaeldeistler michaeldeistler left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot Ben! It looks great overall, I made just a few comments.

README.md Outdated Show resolved Hide resolved
sbi/inference/snre/snre_c.py Outdated Show resolved Hide resolved
sbi/inference/snre/snre_c.py Show resolved Hide resolved
tutorials/16_implemented_methods.ipynb Show resolved Hide resolved
@bkmi
Copy link
Contributor Author

bkmi commented Dec 10, 2022

Jumping in here regarding warnings for unnormalized posteriors: the way this is currently handled is that the MCMCPosterior (regardless of whether it is created by NRE, NLE, or sequential methods) warns if one calls .log_prob() here. To avoid the warning, the user should use .potential() for MCMCPosteriors.

I suggest to describe in the docstring of NRE_C that its .potential() will be incouraged to be normalized, but to not issue further warnings.

I'm going to go with this approach.

Copy link
Contributor

@janfb janfb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great now, thanks a lot for adding SNRE_C and for the detailed discussion!

We merged the other PR with BNRE already, which might have caused merge conflicts with this branch, sorry!
Once those are resolved, this is good to into main as well!

sbi/inference/snre/snre_c.py Show resolved Hide resolved
…idate the discussion of the effects on .potential to that function / description of nre-c.
@bkmi
Copy link
Contributor Author

bkmi commented Dec 19, 2022

I clarified the effects and left it to .potential to do the heavy lifting of warning the user w.r.t. unnormalized ratios of densities.

I also resolved the merge conflicts.

Assuming the tests pass I'd say it's good to go~

Copy link
Contributor

@janfb janfb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for fixing the conflicts and the pyright problems with the python versions.
Looks great now!

@michaeldeistler michaeldeistler merged commit 0ae6126 into sbi-dev:main Dec 19, 2022
@bkmi bkmi deleted the snre_c branch December 19, 2022 16:44
@bkmi
Copy link
Contributor Author

bkmi commented Dec 19, 2022

Thanks everyone! Really appreciate the inclusion of NRE-C in sbi and I found the process to be enjoyable and collaborative!

@michaeldeistler
Copy link
Contributor

Thanks a lot for contributing, Ben! And congrats on the work itself!

@janfb
Copy link
Contributor

janfb commented Dec 20, 2022

Yeah, thanks a lot @bkmi! It was a pleasure, and kudos for your work!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants