Fix fmpe singularity on sde sampling

[manuelgloeckler]

This aims to fix an inefficiency in FMPE drift/diffusion definitions. This lead to unhealthy sampling trajectories (SDE divergence in the beginning). This should generally improve sampling with SDE on FMPE, but also likely fix #1656 .

[manuelgloeckler]

Running IID tests locally to check if the skip in FMPE can be removed.

[codecov]

Codecov Report

❌ Patch coverage is 66.66667% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.59%. Comparing base (bdf7d83) to head (214c182).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
...i/neural_nets/estimators/flowmatching_estimator.py	80.00%	1 Missing ⚠️
sbi/samplers/score/diffuser.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1661      +/-   ##
==========================================
- Coverage   86.70%   82.59%   -4.12%     
==========================================
  Files         134      134              
  Lines       11150    11151       +1     
==========================================
- Hits         9668     9210     -458     
- Misses       1482     1941     +459     

Flag	Coverage Δ
unittests	82.59% <66.66%> (-4.12%)	⬇️

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

Files with missing lines	Coverage Δ
sbi/inference/posteriors/vector_field_posterior.py	70.89% <ø> (-8.21%)	⬇️
...i/neural_nets/estimators/flowmatching_estimator.py	81.66% <80.00%> (-14.95%)	⬇️
sbi/samplers/score/diffuser.py	86.88% <0.00%> (-4.92%)	⬇️

... and 28 files with indirect coverage changes

[Copilot]

Pull Request Overview

This PR fixes numerical singularities in the Flow Matching Posterior Estimation (FMPE) that caused unstable SDE sampling trajectories. The fix changes the singularity bound from 1e-6 to 1e-2 to prevent drift/diffusion explosions.

• Increased singularity bounds in drift and diffusion functions from 1e-6 to 1e-2
• Fixed dimension concatenation bug in diffuser intermediate sample handling
• Added save_intermediate parameter to vector field posterior sampling interface

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
sbi/neural_nets/estimators/flowmatching_estimator.py	Updated singularity bounds in drift_fn and diffusion_fn to prevent numerical explosions
sbi/samplers/score/diffuser.py	Fixed concatenation dimension from 0 to 1 for intermediate samples
sbi/inference/posteriors/vector_field_posterior.py	Added save_intermediate parameter to diffusion sampling method

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[janfb]

Skipped tests here

sbi/tests/linearGaussian_vector_field_test.py

Lines 376 to 382 in bdf7d83

	if (
	vector_field_type == "fmpe"
	and prior_type == "uniform"
	and iid_method in ["gauss", "auto_gauss", "jac_gauss"]
	):
	# TODO: Predictor produces NaNs for these cases, see #1656
	pytest.skip("Known issue of IID methods with uniform priors, see #1656.")

need to be adapted.

please double check whether it all passes locally.

[manuelgloeckler]

Okay mean_t_fn also required the same treatment, which explains why it was prior dependent.

In FMPE at t=1 this became zero, which rendered some of the analytical equations used in IID setting ill-defined (as then there is nothing to denoise).

[manuelgloeckler]

Will run the test without the skip and remove it once passing.

[janfb]

For me, the three iid gauss variants are still failing with FMPE and uniform prior because all samples are NaN.

[manuelgloeckler]

For me, the three iid gauss variants are still failing with FMPE and uniform prior because all samples are NaN.

Where you on the most recent commit? (before the last one) `The mean_t_fn modification was necessary. Will check again, but for me all tests did pass.

[janfb]

For me, the three iid gauss variants are still failing with FMPE and uniform prior because all samples are NaN.

Where you on the most recent commit? (before the last one) `The mean_t_fn modification was necessary. Will check again, but for me all tests did pass.

Yes, all passing now. thanks! 🚀