#1405 add util to generate mcmc samples from user defined potential (#1483)

* add util to generate mcmc samples directly from potential

* refactor to leverage MCMCPosterior class and return sampler

* move from utils to posteriors and return full posterior

* add tests for mcmc posterior from user defined potential

* fix formatting

Author: hayden-johnson

sbi/inference/posteriors/mcmc_posterior.py

@@ -1,5 +1,6 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+import inspect
 import warnings
 from copy import deepcopy
 from functools import partial
@@ -32,6 +33,7 @@
     sir_init,
 )
 from sbi.sbi_types import Shape, TorchTransform
+from sbi.utils import mcmc_transform
 from sbi.utils.potentialutils import pyro_potential_wrapper, transformed_potential
 from sbi.utils.torchutils import ensure_theta_batched, tensor2numpy
 
@@ -1110,3 +1112,80 @@ def _maybe_use_dict_entry(default: Any, key: str, dict_to_check: Dict) -> Any:
     """
     attribute = dict_to_check.get(key, default)
     return attribute
+
+
+def _num_required_args(func):
+    """
+    Utility for counting the number of positional args in a function.
+
+    This function counts each parameter in the signature that are positional -- ie.
+        (1) cannot only be passed in as keyword arguments
+        (2) do not have a default value
+
+    Args:
+        func: A callable function.
+
+    Returns:
+        Number of required positional arguments.
+    """
+    sig = inspect.signature(func)
+    return sum(
+        1
+        for param in sig.parameters.values()
+        if param.kind
+        in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        and param.default is inspect._empty
+    )
+
+
+def build_from_potential(
+    potential_fn: Callable, prior: Any, x: Optional[Tensor] = None, **kwargs
+) -> MCMCPosterior:
+    """
+    Returns a sampler from a MCMCPosterior object, given user-defined potential
+    function and prior.
+
+    The user-defined potential can be conditional (accepts theta and x as positional
+    arguments) or unconditional (accepting only theta).
+
+    Args:
+        potential_fn: User defined potential function. Must be of type Callable.
+        prior: Prior distribution for parameter transformation and initialization.
+        x: Conditional x value. Provided if using a conditional potential function.
+
+    Returns:
+        Callable sampling function from MCMCPosterior object.
+    """
+    # build transformation to unrestricted space for sampling
+    transform = mcmc_transform(prior)
+
+    # potential_fn must take 1 or 2 required arguments: (theta) or (theta, x)
+    num_args = _num_required_args(potential_fn)
+    assert num_args > 0 and num_args < 3, (
+        "potential_fn must take 1-2 required arguments"
+    )
+    is_conditional = num_args == 2
+
+    if is_conditional:
+        # you could remove this and require use to set x before calling sample
+        assert x is not None, "x must be provided if potential_fn is conditional"
+        posterior = MCMCPosterior(potential_fn, prior, theta_transform=transform)
+        posterior.set_default_x(x)
+
+    else:
+        warn(
+            "x has not been provided. Using unconditional potential function.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+        # define an unconditional potential function (ignores x)
+        def unconditional_potential_fn(theta, x):
+            return potential_fn(theta)
+
+        posterior = MCMCPosterior(
+            unconditional_potential_fn, prior, theta_transform=transform, **kwargs
+        )
+        posterior.set_default_x(torch.zeros(1))  # set default_x to dummy value
+
+    return posterior

tests/mcmc_test.py

@@ -17,6 +17,7 @@
     MCMCPosterior,
     likelihood_estimator_based_potential,
 )
+from sbi.inference.posteriors.mcmc_posterior import build_from_potential
 from sbi.neural_nets import likelihood_nn
 from sbi.samplers.mcmc.pymc_wrapper import PyMCSampler
 from sbi.samplers.mcmc.slice_numpy import (
@@ -28,6 +29,7 @@
     diagonal_linear_gaussian,
     true_posterior_linear_gaussian_mvn_prior,
 )
+from sbi.utils import BoxUniform
 from sbi.utils.user_input_checks import process_prior
 from tests.test_utils import check_c2st
 
@@ -252,3 +254,93 @@ def test_getting_inference_diagnostics(method, mcmc_params_fast: dict):
         f"MCMC samples for method {method} have incorrect shape (n_samples, n_dims). "
         f"Expected {(num_samples, num_dim)}, got {samples.shape}"
     )
+
+
+@pytest.mark.mcmc
+def test_direct_mcmc_unconditional():
+    "Test MCMCPosterior from user defined potential (unconditional)"
+    num_samples = 100
+    theta_dim = 2
+
+    prior = BoxUniform(low=-2 * torch.ones(theta_dim), high=2 * torch.ones(theta_dim))
+
+    def potential_fn(theta: np.ndarray) -> np.ndarray:
+        # Example: a 2D Gaussian with mean=[0,0], identity covariance
+        return -0.5 * (theta**2).sum(axis=-1)
+
+    mcmc_posterior = build_from_potential(potential_fn, prior)
+
+    # test sampling
+    samples = mcmc_posterior.sample(
+        (num_samples,), num_chains=10, warmup_steps=50, thin=10
+    )
+
+    assert samples.shape == (num_samples, theta_dim), (
+        f"MCMC samples have incorrect shape (n_samples, n_dims). "
+        f"Expected {(num_samples, theta_dim)}, got {samples.shape}"
+    )
+
+    # test potential evaluation
+    dist = torch.distributions.MultivariateNormal(
+        torch.zeros(theta_dim), torch.eye(theta_dim)
+    )
+    samples = dist.sample((num_samples,))
+    log_p = mcmc_posterior.potential(samples)
+
+    assert log_p.shape == (num_samples,), (
+        f"Potential evals have incorrect shape. "
+        f"Expected ({num_samples}), got {log_p.shape}"
+    )
+
+
+@pytest.mark.mcmc
+def test_direct_mcmc_conditional():
+    "Test MCMCPosterior from user defined potential (conditional)"
+    theta_dim = 2
+    num_samples = 100
+    num_batches = 5
+    num_samples_batch = num_samples // num_batches
+
+    prior = BoxUniform(low=-2 * torch.ones(theta_dim), high=2 * torch.ones(theta_dim))
+
+    def potential_fn(theta: np.ndarray, x: np.ndarray) -> np.ndarray:
+        # Example: a 2D Gaussian with mean=[0,0], variance conditioned on x
+        return -x * (theta**2).sum(axis=-1)
+
+    # test sampling
+    x = torch.tensor([0.5])
+    mcmc_posterior = build_from_potential(potential_fn, prior, x=x)
+    samples = mcmc_posterior.sample(
+        (num_samples,), num_chains=10, warmup_steps=50, thin=10
+    )
+
+    assert samples.shape == (num_samples, theta_dim), (
+        f"MCMC samples have incorrect shape (n_chains, n_samples, n_dims). "
+        f"Expected {(num_samples, theta_dim)}, got {samples.shape}"
+    )
+
+    # test batched sampling
+    x_batch = torch.linspace(0.1, 0.9, num_batches).unsqueeze(1)
+    samples_batched = mcmc_posterior.sample_batched(
+        (num_samples_batch,), x=x_batch, num_chains=10, warmup_steps=50, thin=10
+    )
+    assert samples_batched.shape == (num_samples_batch, num_batches, theta_dim), (
+        f"MCMC samples have incorrect shape (n_samples, n_batches, n_dims). "
+        f"Expected {(num_samples, num_batches, theta_dim)}, got {samples.shape}"
+    )
+
+    # test potential evaluation
+    dist = torch.distributions.MultivariateNormal(
+        torch.zeros(theta_dim), torch.eye(theta_dim)
+    )
+    theta_samples = dist.sample((num_samples,))
+    x_samples = torch.rand((num_samples,))
+    log_p = mcmc_posterior.potential(theta_samples, x_samples)
+
+    assert log_p.shape == (
+        1,
+        num_samples,
+    ), (
+        f"Potential evals have incorrect shape. "
+        f"Expected (1, {num_samples}), got {log_p.shape}"
+    )