additional features for NPSE

[gmoss13]

What does this implement/fix? Explain your changes

This introduces some additional features for score estimation named in #1226, namely:

• allow enable_transform = True for score-based potentials
• implement MAP calculation for score-based posteriors
• Implements rejection sampling for score-based posteriors to ensure prior coverage
• Allow batched sampling for score-based posteriors
• Allow IID observations for score-based posteriors (@manuelgloeckler has started working on this - let's discuss what's missing and merge our branches)
• implements custom converged() method for NPSE

Does this close any currently open issues?

#1226

Any relevant code examples, logs, error output, etc?

Any other comments?

• Currently, calling score_based_posterior.map() is still quite slow. We get the gradient of the log probs with respect to theta by using the score estimator, but still computing the log-probs explicitly in gradient_ascent, which is more expensive. To get around this, we save a low-accuracy ode_flow to calculate the log-probs more quickly. Ideally, we might want to write a custom gradient_ascent function for calculating the MAP for score estimators to avoid doing this altogether.
• I increased the tolerance of the test in linearGaussian_npse_test.py::test_npse_map - as far as I can tell, the reason this failed with the lower tolerance is not because of MAP calculation, but because score-based posteriors are currently slightly less accurate (at least for our test tasks).

[gmoss13]

specified torch<2.6.0 to avoid type checking errors as mentioned in #1380

[gmoss13]

I've requested review now. While batched sampling for score-based posteriors is now possible and tested for, IID sampling is still not possible, but talking to @manuelgloeckler about this, maybe this can be done in a new PR. Other than that, I've also noticed while testing that sampling from the posterior with ode can be much less accurate than via diffusion. So the test linear_Gaussian_npse_test::test_c2st_npse_on_linearGaussian can sometimes fail with sample_with="ode", but this is independent of any of the changes made in this PR.

[codecov]

Codecov Report

Attention: Patch coverage is 57.73196% with 41 lines in your changes missing coverage. Please review.

Project coverage is 78.18%. Comparing base (18f92b1) to head (205d59e).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
sbi/inference/posteriors/score_posterior.py	43.18%	25 Missing ⚠️
sbi/inference/potentials/score_based_potential.py	57.89%	16 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1370       +/-   ##
===========================================
- Coverage   89.31%   78.18%   -11.13%     
===========================================
  Files         119      119               
  Lines        8779     8844       +65     
===========================================
- Hits         7841     6915      -926     
- Misses        938     1929      +991     

Flag	Coverage Δ
unittests	78.18% <57.73%> (-11.13%)	⬇️

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

Files with missing lines	Coverage Δ
sbi/inference/posteriors/direct_posterior.py	97.67% <ø> (ø)
sbi/inference/trainers/npse/npse.py	96.50% <100.00%> (ø)
sbi/samplers/rejection/rejection.py	87.75% <100.00%> (-0.25%)	⬇️
sbi/samplers/score/diffuser.py	85.18% <100.00%> (+0.27%)	⬆️
sbi/utils/restriction_estimator.py	76.31% <ø> (-8.65%)	⬇️
sbi/inference/potentials/score_based_potential.py	80.80% <57.89%> (-16.17%)	⬇️
sbi/inference/posteriors/score_posterior.py	74.00% <43.18%> (-23.02%)	⬇️

... and 31 files with indirect coverage changes

[manuelgloeckler]

I started integrating the IID stuff into the current version of this branch and created a new PR for it (#1381). So, lets first get this merged the IID PR still requires some work from my side.

[janfb]

Great! Thanks a lot for addressing all these issues with the current NPSE. 🚀

I left a couple of suggestions and questions for my understanding. Happy to discuss in person if needed.

[gmoss13]

Thanks for the comments @janfb! I've tried to answer some of your questions, and will make the appropriate changes soon!

[janfb]

Thanks for addressing my questions and comments!

I have some final minor things and a renaming suggestion, plus an idea for better type checking.

Also, I would suggest that we add more info about how to use the score-based methods in the tutorials as well, e.g., the convergence details about the validation times, about sampling via sde vs ode, how to use MAP, how to do i.i.d. inference. (could also be a separate PR). I think this will be important for users to actually be able to use all these new features.

[gmoss13]

Thanks @janfb! I agree that adding a tutorial on the new features would be important - as we plan to rework the tutorials in the hackathon, I think this would belong in a later PR 😄

[janfb]

looks good now. great work @gmoss13 ! 🎉 🚀

[janfb]

Thanks @janfb! I agree that adding a tutorial on the new features would be important - as we plan to rework the tutorials in the hackathon, I think this would belong in a later PR 😄

moved to issue #1392