1446 add api for guidance by manuelgloeckler · Pull Request #1482 · sbi-dev/sbi

manuelgloeckler · 2025-03-18T15:47:47Z

This adds an API for post-hoc modifications of the trained score estimator, allowing modifications of the likelihood or prior, the support, or other additional constraints. This addresses issue #1446 .

To does:

Adds general API inline with current iid_method
Refactor the parameterizations from stringly type to strongly typed (also for iid!)
Adds some useful explanatory guidance approaches:
- Classifier free guidance
- Universal guidance - interval truncations.

…v/sbi into 1446-add-api-for-guidance

- add pytest-split plugin to dev dependencies - use split in ci workflow based on test durations in .test_durations

…v/sbi into 1446-add-api-for-guidance

codecov · 2025-03-20T08:53:49Z

Codecov Report

❌ Patch coverage is 67.21311% with 160 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.83%. Comparing base (937efc2) to head (672514b).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sbi/utils/vector_field_utils.py	65.78%	103 Missing ⚠️
sbi/inference/potentials/vector_field_adaptor.py	68.18%	56 Missing ⚠️
sbi/inference/potentials/vector_field_potential.py	90.90%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1482      +/-   ##
==========================================
- Coverage   88.54%   83.83%   -4.72%     
==========================================
  Files         137      136       -1     
  Lines       11515    11797     +282     
==========================================
- Hits        10196     9890     -306     
- Misses       1319     1907     +588

Flag	Coverage Δ
fast	`83.83% <67.21%> (?)`
full	`?`

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	`69.12% <ø> (-8.06%)`	⬇️
...i/neural_nets/estimators/flowmatching_estimator.py	`81.66% <ø> (-15.00%)`	⬇️
sbi/inference/potentials/vector_field_potential.py	`75.20% <90.90%> (-15.87%)`	⬇️
sbi/inference/potentials/vector_field_adaptor.py	`78.19% <68.18%> (ø)`
sbi/utils/vector_field_utils.py	`66.12% <65.78%> (-17.21%)`	⬇️

... and 25 files with indirect coverage changes

…ance

Update the branch

…her i.e. score -> vector field

janfb

Great work @manuelgloeckler 👏 Great to have PriorGuide in here as well 🚀

added mostly minor formatting and edge case handling comments.

More high-level, I am concerned about the long signature sample(...) has by now. We should think about introducing config classes for the different methods (sde config, iid, guidance), or method chaining (see comment below). But this would be a larger refactoring - let's discuss.

On the test side, maybe add another test on the combination of guidance and iid settings, given that this implemented as an option?

Documentation: great to have a tutorial as part of advanced tutorial 20 already. As a follow-up, I suggest a refactoring of tutorials 19 and 20, and 1-2 how-to-guide for the VF and guidance methods.

janfb · 2026-02-25T05:41:58Z

docs/advanced_tutorials/20_score_based_methods_new_features.ipynb

Great to have this tutorial. Overall, the VF and NPSE variants tutorials 19 and 20 are very long now and with overlap. I am already working on a refactoring and will take the guidance part into account as well.

janfb · 2026-02-25T05:47:57Z

sbi/inference/posteriors/vector_field_posterior.py

+            guidance_method: Method to guide the diffusion process. If None, no guidance
+                is used. currently we support `affine_classifier_free`, which allows to
+                scale and shift the "likelihood" or "prior" score contribution. This can
+                be used to perform "super" conditioning i.e. shring the variance of the


typo: shrink

janfb · 2026-02-25T05:48:08Z

sbi/inference/posteriors/vector_field_posterior.py

+                scale and shift the "likelihood" or "prior" score contribution. This can
+                be used to perform "super" conditioning i.e. shring the variance of the
+                likelihood. `Universal` can be used to guide the diffusion process with
+                a general guidance function. `Interval` is an isntance of that where


typo: instance

janfb · 2026-02-25T05:50:44Z

sbi/inference/posteriors/vector_field_posterior.py

+        guidance_method: Optional[str] = None,
+        guidance_params: Optional[Dict] = None,


the sample method now has >20 kwargs. we could think about introducing config classes instead, e.g., sth like

from sbi.inference.posteriors import GuidanceConfig, PriorGuideCfg guidance = GuidanceConfig( method="prior_guide", params=PriorGuideCfg(train_prior=..., test_prior=..., K=5) ) samples = posterior.sample((1000,), x=x_o, guidance=guidance)

but this would be a larger, possibly follow-up refactoring. what do you think?

or even sth like

samples = posterior.with_guidance("prior_guide", ...).sample((1000,), x=x_o)

Mhh, I think config classes would be fine (I internally use config classes anyway). But at least at the time of implementation most user API was still configured via dicts so I made this the external API.

I do actually like your suggestions with with_guidance and with_iid to make it more explicit. But yes I would make this a follow-up refactoring.

janfb · 2026-02-25T05:52:45Z

sbi/inference/posteriors/vector_field_posterior.py

            iid_params: Additional parameters passed to the iid method. See the specific
                `IIDScoreFunction` child class for details.
+            guidance_method: Method to guide the diffusion process. If None, no guidance
+                is used. currently we support `affine_classifier_free`, which allows to


janfb · 2026-02-25T06:16:25Z

sbi/inference/potentials/vector_field_adaptor.py

        return denoising_posterior_precision


+def compute_score(p: Distribution, inputs: Tensor):


add docstring and return type please.

janfb · 2026-02-25T06:22:25Z

sbi/utils/vector_field_utils.py

+    max_log_ratio: float = 50.0,
+    device: Union[str, torch.device] = "cpu",
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """Implementation for fitting a generalized GMM to the prior ratio


Drop "Implementation for fitting", replace by "Fit" to have it in one line

janfb · 2026-02-25T06:34:20Z

sbi/utils/vector_field_utils.py

+class _HashableById:
+    __slots__ = ("obj", "_id")
+
+    def __init__(self, obj: Distribution):
+        """Wraps a non-hashable Distribution to make it cache-key compatible."""
+        self.obj = obj
+        self._id = id(obj)
+
+    def __hash__(self) -> int:
+        return self._id
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, _HashableById) and self._id == other._id


overall great to have this here! the main purpose is to not refit the GMM on every call from the potential, right?

just wondering, this caches by object ID not distribution parameters, so if someone passes the same dist with same params but a new instance, this won't have an effect. Accordingly, when someone changes the distribution params in place this will not be noticed in the cache.
But I think both are edge cases so it's fine

Yeah, there are definitely some edge cases where this will fail. This is why it is only used if required.

I generally a bit unhappy that the adaptor/score wrapper objects are reinitialized in each call of potential (which does require to hash as much as possible). I think it would be good the restructure the VF potential a bit that such quantities can be computed once at the beginning i.e. via an initalize_aux which needs to be called once berfore sample.

janfb · 2026-02-25T06:36:43Z

tests/score_samplers_test.py

 from sbi.utils import BoxUniform, MultipleIndependent


+def build_some_priors(num_dim: int):


rename to build_test_priors and add short docstring?

janfb · 2026-02-25T06:36:52Z

tests/score_samplers_test.py

+    # Bug in Independent introduced????
+    priors = [prior1, prior2, prior2_2, prior3, prior4, prior5]  # Something broke


what is this about? is there a bug?

Good catch will check, but it includes all the priors so these tests did atleast pass. (so might have to remove the comments)

manuelgloeckler and others added 11 commits March 18, 2025 10:13

Some progress

18512f4

Some progress

f5e8e7f

Merge branch '1446-add-api-for-guidance' of https://github.com/sbi-de…

7def57b

…v/sbi into 1446-add-api-for-guidance

Moving and renaming. Refactoring IID function to reduce code redundancy

ab6b7b4

basic api on potentials

52825d3

working minimal example

0237adf

Docs: Introduce RTD website

2e7d27b

Fix tests for new docs

9b7a14f

chore: pytest split for ci workflow (#1465)

45d9ca6

- add pytest-split plugin to dev dependencies - use split in ci workflow based on test durations in .test_durations

Prevent notebook execution upon doc build

6f891da

Fix broken links on website

5592682

manuelgloeckler linked an issue Mar 18, 2025 that may be closed by this pull request

Add API for guidance #1446

Open

5 tasks

internal dataclasses but external stringly...

bd2735f

manuelgloeckler marked this pull request as draft March 18, 2025 18:48

manuelgloeckler self-assigned this Mar 18, 2025

manuelgloeckler added 15 commits March 19, 2025 11:09

Other more general guidance methods

762ed43

universal guidance working

9653d0e

Specialized child class for interval constriants

fe23d43

Some progress

4789025

Moving and renaming. Refactoring IID function to reduce code redundancy

7b43583

basic api on potentials

ec7ae06

working minimal example

1469881

internal dataclasses but external stringly...

b5edda9

Other more general guidance methods

51cde80

universal guidance working

251fc88

Specialized child class for interval constriants

4a5f431

Merge branch '1446-add-api-for-guidance' of https://github.com/sbi-de…

7b5b2e1

…v/sbi into 1446-add-api-for-guidance

Formatting

3652eb5

typing

48906b1

tests

1eba715

manuelgloeckler added 3 commits March 19, 2025 18:53

check independet prior, problem

f29d75c

Fix hidden bug in score_utils

c38db7a

Format

68616d4

manuelgloeckler marked this pull request as ready for review March 20, 2025 07:05

Removing unncessary. Adding tests on basic API

a75a944

manuelgloeckler added 3 commits March 20, 2025 10:25

test lower upper, lower upper where switched

2efb7c7

Formats and tests

ff289b1

add documentation

5fdbb4c

manuelgloeckler added the blocked Something is in the way of fixing this. Refer to it in the issue label Mar 24, 2025

manuelgloeckler removed the blocked Something is in the way of fixing this. Refer to it in the issue label Sep 5, 2025

manuelgloeckler and others added 11 commits September 5, 2025 14:06

Merge remote-tracking branch 'origin/main' into 1446-add-api-for-guid…

3679f3c

…ance

Port changes

5972a29

Update to recent changes

50fde2f

Tutorial

5637db4

Merge remote-tracking branch 'origin/main' into 1446-add-api-for-guid…

e50f6a1

…ance

froamting

0476ab0

Merge branch 'main' into 1446-add-api-for-guidance

e05c456

Update the branch

Added new prior Guide method. Renamed and merged previous files toget…

651874e

…her i.e. score -> vector field

Ruff formatting

efbb8b3

improvement notes

f1799ad

Updated documentaiton.

672514b

manuelgloeckler requested a review from janfb February 2, 2026 16:49

janfb requested changes Feb 25, 2026

View reviewed changes

		guidance_method: Optional[str] = None,
		guidance_params: Optional[Dict] = None,

		return denoising_posterior_precision


		def compute_score(p: Distribution, inputs: Tensor):

		from sbi.utils import BoxUniform, MultipleIndependent


		def build_some_priors(num_dim: int):

		# Bug in Independent introduced????
		priors = [prior1, prior2, prior2_2, prior3, prior4, prior5] # Something broke

Comments

Conversation

manuelgloeckler commented Mar 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Mar 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

janfb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

manuelgloeckler Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

manuelgloeckler commented Mar 18, 2025 •

edited

Loading

codecov bot commented Mar 20, 2025 •

edited

Loading

janfb left a comment •

edited

Loading

manuelgloeckler Feb 25, 2026 •

edited

Loading