diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 8d181012b..ede6ee394 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -51,7 +51,7 @@ jobs: run: uv run pytest -v -x -n auto -m "not gpu" --cov=sbi --cov-report=xml --junitxml=junit.xml -o junit_family=legacy tests/ - name: Upload coverage to Codecov - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: env_vars: OS,PYTHON files: ./coverage.xml @@ -62,7 +62,7 @@ jobs: - name: Upload test results to Codecov if: ${{ !cancelled() }} - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: report_type: test_results env: diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c937915ba..43054217c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,7 +64,7 @@ jobs: run: uv run pytest -v -n auto -m "not slow and not gpu" --cov=sbi --cov-report=xml --junitxml=junit.xml -o junit_family=legacy --splitting-algorithm least_duration --splits ${{ matrix.split_size }} --group ${{ matrix.group_number }} tests/ - name: Upload coverage to Codecov - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: env_vars: OS,PYTHON file: ./coverage.xml @@ -75,7 +75,7 @@ jobs: - name: Upload test results to Codecov if: ${{ !cancelled() }} - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: report_type: test_results env: diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index fd27f72ba..9ae4fe341 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -77,7 +77,7 @@ jobs: name: python-package-distributions path: dist/ - name: Sign the dists with Sigstore - uses: sigstore/gh-action-sigstore-python@04cffa1d795717b140764e8b640de88853c92acc # v3.3.0 + uses: sigstore/gh-action-sigstore-python@5b79a39c381910c090341a2c9b0bf022c8b387e1 # v3.4.0 with: inputs: >- ./dist/*.tar.gz diff --git a/docs/notebooks/sbi_old_vs_new_api.ipynb b/docs/notebooks/sbi_old_vs_new_api.ipynb new file mode 100644 index 000000000..43cb43830 --- /dev/null +++ b/docs/notebooks/sbi_old_vs_new_api.ipynb @@ -0,0 +1,264 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Old API vs New API Demo" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " Neural network successfully converged after 77 epochs." + ] + } + ], + "source": [ + "import warnings\n", + "\n", + "import torch\n", + "\n", + "from sbi.inference import NLE, MCMCPosterior\n", + "from sbi.utils import BoxUniform\n", + "\n", + "warnings.filterwarnings('default')\n", + "\n", + "\n", + "# Setup\n", + "prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2))\n", + "def simulator(θ):\n", + " return θ + torch.randn_like(θ) * 0.1\n", + "\n", + "torch.manual_seed(42)\n", + "theta_train = prior.sample((1000,))\n", + "x_train = simulator(theta_train)\n", + "x_o = torch.randn(1, 2)\n", + "\n", + "# Train\n", + "trainer = NLE()\n", + "likelihood_estimator = trainer.append_simulations(theta_train, x_train).train()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OLD API (Deprecated)\n", + "\n", + "The old API uses stateful classes that bind `x_o` at construction time." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:02<00:00, 7.25it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:03<00:00, 594.07it/s] " + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OLD API - samples mean: tensor([0.9753, 0.0945])\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "# OLD API: stateful, emits DeprecationWarning\n", + "from sbi.inference.potentials import likelihood_estimator_based_potential\n", + "\n", + "potential_fn, parameter_transform = likelihood_estimator_based_potential(\n", + " likelihood_estimator, prior, x_o\n", + ")\n", + "posterior_old = MCMCPosterior(\n", + " potential_fn, proposal=prior, theta_transform=parameter_transform, warmup_steps=10\n", + ")\n", + "samples_old = posterior_old.sample((1000,))\n", + "print(f\"OLD API - samples mean: {samples_old.mean(dim=0)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## NEW API (Recommended)\n", + "\n", + "The new API uses stateless functions that follow the `PotentialFunction` protocol:\n", + "- `(theta, x) -> log_prob`\n", + "- No hidden state, all inputs explicit\n", + "- Combine with `bind_observation()`" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 46.07it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:05<00:00, 425.35it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "NEW API - samples mean: tensor([0.9765, 0.0895])\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "# NEW API: stateless + bind_observation\n", + "from sbi.inference.potentials import bind_observation, likelihood_potential\n", + "from sbi.utils import mcmc_transform\n", + "\n", + "# Step 1: Create stateless potential (satisfies PotentialFunction protocol)\n", + "potential_fn = likelihood_potential(likelihood_estimator, prior)\n", + "\n", + "# Step 2: Get transform from prior (for MCMC in unconstrained space)\n", + "theta_transform = mcmc_transform(prior, enable_transform=True)\n", + "\n", + "# Step 3: Bind observation to create sampler-ready function\n", + "bound_fn = bind_observation(potential_fn, x_o, sum_iid=True)\n", + "\n", + "# Step 4: Use with MCMC sampler (pass theta_transform!)\n", + "posterior_new = MCMCPosterior(\n", + " bound_fn, proposal=prior, theta_transform=theta_transform, warmup_steps=10\n", + ")\n", + "samples_new = posterior_new.sample((1000,))\n", + "print(f\"NEW API - samples mean: {samples_new.mean(dim=0)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Lower-Level: Using the Protocol Directly\n", + "\n", + "You can also implement your own potential function that satisfies the protocol." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 45.43it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:02<00:00, 998.14it/s] " + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Custom potential - samples mean: tensor([0.9757, 0.0981])\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "# Custom potential function satisfying PotentialFunction protocol\n", + "def my_custom_potential(theta: torch.Tensor, x: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"Custom potential: log p(x|theta) + log p(theta)\"\"\"\n", + " log_likelihood = likelihood_estimator.log_prob(x, condition=theta)\n", + " log_prior = prior.log_prob(theta)\n", + " return log_likelihood + log_prior\n", + "\n", + "# Add device attribute for bind_observation compatibility\n", + "my_custom_potential.device = next(likelihood_estimator.parameters()).device\n", + "\n", + "# Bind observation\n", + "bound_custom = bind_observation(my_custom_potential, x_o, sum_iid=True)\n", + "\n", + "# Use with sampler\n", + "posterior_custom = MCMCPosterior(\n", + " bound_custom, proposal=prior, warmup_steps=10\n", + ")\n", + "samples_custom = posterior_custom.sample((1000,))\n", + "print(f\"Custom potential - samples mean: {samples_custom.mean(dim=0)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OLD API: tensor([0.9753, 0.0945])\n", + "NEW API: tensor([0.9765, 0.0895])\n", + "Custom: tensor([0.9757, 0.0981])\n" + ] + } + ], + "source": [ + "print(f\"OLD API: {samples_old.mean(dim=0)}\")\n", + "print(f\"NEW API: {samples_new.mean(dim=0)}\")\n", + "print(f\"Custom: {samples_custom.mean(dim=0)}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "sbi (3.12.12)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/sbi/inference/potentials/__init__.py b/sbi/inference/potentials/__init__.py index c88d6301c..a2fb09991 100644 --- a/sbi/inference/potentials/__init__.py +++ b/sbi/inference/potentials/__init__.py @@ -1,13 +1,23 @@ # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed # under the Apache License Version 2.0, see +from sbi.inference.potentials.binding import ( + BoundPotential, + bind_observation, +) from sbi.inference.potentials.likelihood_based_potential import ( likelihood_estimator_based_potential, + likelihood_potential, mixed_likelihood_estimator_based_potential, ) from sbi.inference.potentials.posterior_based_potential import ( posterior_estimator_based_potential, ) +from sbi.inference.potentials.protocol import ( + PotentialFunction, + PotentialFunctionWithDevice, + validate_potential, +) from sbi.inference.potentials.ratio_based_potential import ( ratio_estimator_based_potential, ) @@ -16,9 +26,15 @@ ) __all__ = [ + "bind_observation", + "BoundPotential", "likelihood_estimator_based_potential", + "likelihood_potential", "mixed_likelihood_estimator_based_potential", "posterior_estimator_based_potential", + "PotentialFunction", + "PotentialFunctionWithDevice", "ratio_estimator_based_potential", + "validate_potential", "vector_field_estimator_based_potential", ] diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index d24b05b94..f640d6e80 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_potential.py @@ -1,6 +1,7 @@ # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed # under the Apache License Version 2.0, see +import inspect from abc import ABCMeta, abstractmethod from typing import Optional, Protocol, Union @@ -147,10 +148,18 @@ def to(self, device: Union[str, torch.device]) -> "CustomPotentialWrapper": super().to(device) return self + def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): + super().set_x(x_o, x_is_iid) + def __call__(self, theta, track_gradients: bool = True): """Calls the custom potential function on given theta. Note, x_o is re-used from the initialization of the potential function. """ + sig = inspect.signature(self.potential_fn) + num_params = len(sig.parameters) with torch.set_grad_enabled(track_gradients): - return self.potential_fn(theta, self.x_o) + if num_params == 1: + return self.potential_fn(theta) + else: + return self.potential_fn(theta, self.x_o) diff --git a/sbi/inference/potentials/binding.py b/sbi/inference/potentials/binding.py new file mode 100644 index 000000000..3e92ddb03 --- /dev/null +++ b/sbi/inference/potentials/binding.py @@ -0,0 +1,127 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +"""Binding utilities for creating sampler-ready potential functions. + +This module provides Layer 2 utilities for binding observations to stateless +potential functions, making them ready for MCMC/VI samplers. +""" + +from typing import Union + +import torch +from torch import Tensor + +from sbi.inference.potentials.base_potential import BasePotential +from sbi.inference.potentials.protocol import PotentialFunction + + +def bind_observation( + potential: PotentialFunction, + x_o: Tensor, + sum_iid: bool = True, +) -> "BoundPotential": + """Create a sampler-compatible theta -> log_prob function. + + This utility binds observed data to a stateless potential function, producing + a callable ready for use with MCMC, VI, or rejection sampling algorithms. + + The function infers the device from x_o and validates compatibility with + the potential's device requirements (if the potential has a device attribute). + + Args: + potential: Stateless (theta, x) -> log_prob function satisfying the + PotentialFunction protocol. + x_o: Observed data tensor. Shape should be (obs_batch, *obs_shape) for + IID observations, or (*obs_shape) for single observation. + sum_iid: If True, sum log probabilities over the observation batch dimension. + This is the typical behavior for IID observations where we want to + compute log p(x_o | theta) = sum_i log p(x_i | theta). + + Returns: + A BoundPotential instance that takes theta (parameters) and returns the + log-probability log p(theta | x_o) + log p(theta). The returned object + is compatible with samplers that expect set_x() method. + + Example: + >>> # Create stateless potential + >>> def likelihood_potential(theta: Tensor, x: Tensor) -> Tensor: + ... return estimator.log_prob(x, context=theta) + prior.log_prob(theta) + >>> + >>> # Bind observation for sampling + >>> bound_fn = bind_observation(likelihood_potential, x_o, sum_iid=True) + >>> log_probs = bound_fn(theta_samples) # Ready for MCMC/VI + """ + device = x_o.device + + if hasattr(potential, "device"): + potential_device = potential.device + if isinstance(potential_device, str): + potential_device = torch.device(potential_device) + if potential_device != device: + raise ValueError( + f"Device mismatch: x_o is on {device}, but potential expects " + f"{potential_device}. Use x_o.to('{potential_device}') before " + f"binding, or ensure the potential was created on the same device " + f"as x_o." + ) + + x_o = x_o.to(device) + + return BoundPotential(potential, x_o, sum_iid) + + +class BoundPotential(BasePotential): + """Wrapper class for binding observation to a potential function. + + This provides an alternative to the functional bind_observation, maintaining + compatibility with existing code patterns while supporting the new stateless + protocol internally. + """ + + def __init__( + self, + potential: PotentialFunction, + x_o: Tensor, + sum_iid: bool = True, + ): + self._potential = potential + self._x_o = x_o.to(x_o.device) + self._sum_iid = sum_iid + self._device = x_o.device + self._x_is_iid = True + + @property + def device(self) -> torch.device: + return self._device + + def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: + theta = theta.to(self._device) + + with torch.set_grad_enabled(track_gradients): + log_prob = self._potential(theta, self._x_o) + + if self._sum_iid and self._x_o.ndim > 1: + log_prob = log_prob.sum(dim=0) + + if self._sum_iid: + return log_prob.reshape(-1) if log_prob.ndim > 0 else log_prob.unsqueeze(0) + else: + return log_prob + + def set_x(self, x_o: Tensor, x_is_iid: bool = True) -> None: + """Set observation for the potential function. + + For BoundPotential, x_o is already bound in __init__, so this is a no-op + for backward compatibility with samplers. + """ + pass + + def return_x_o(self) -> Tensor: + """Return the bound observation.""" + return self._x_o + + def to(self, device: Union[str, torch.device]) -> "BoundPotential": + self._device = torch.device(device) + self._x_o = self._x_o.to(device) + return self diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index d2907430c..411742e61 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/sbi/inference/potentials/likelihood_based_potential.py @@ -21,6 +21,73 @@ from sbi.utils.sbiutils import mcmc_transform +def likelihood_potential( + likelihood_estimator: ConditionalDensityEstimator, + prior: Distribution, +) -> Callable[[Tensor, Tensor], Tensor]: + r"""Create a stateless likelihood potential function. + + This returns a function that satisfies the PotentialFunction protocol: + (theta, x) -> log p(x|theta) + log p(theta) + + This is the new recommended API. For use with samplers, combine with + bind_observation(). + + Args: + likelihood_estimator: The density estimator modelling the likelihood. + prior: The prior distribution. + + Returns: + A callable (theta, x) -> log_prob that computes the likelihood potential. + The returned function has a `device` attribute for compatibility with + bind_observation. + """ + device = str(next(likelihood_estimator.parameters()).device) + + def _potential(theta: Tensor, x: Tensor) -> Tensor: + theta = theta.to(device) + x = x.to(device) + + x_batch_size = x.shape[0] + theta_batch_size = theta.shape[0] + + with torch.set_grad_enabled(True): + if x_batch_size == 1: + x_for_eval = x + log_likelihood = likelihood_estimator.log_prob( + x_for_eval, condition=theta + ) + log_likelihood = log_likelihood.squeeze(1) + elif x_batch_size == theta_batch_size: + log_likelihood = likelihood_estimator.log_prob(x, condition=theta) + log_likelihood = log_likelihood.reshape(-1) + else: + x_expanded = reshape_to_sample_batch_event( + x, event_shape=x.shape[1:], leading_is_sample=True + ) + trailing_minus_ones = [-1 for _ in range(x_expanded.dim() - 2)] + x_expanded = x_expanded.expand( + -1, theta_batch_size, *trailing_minus_ones + ) + + theta_batch = reshape_to_batch_event(theta, event_shape=theta.shape[1:]) + + log_likelihood_batch = likelihood_estimator.log_prob( + x_expanded, condition=theta_batch + ) + log_likelihood = log_likelihood_batch.reshape( + x_batch_size, theta_batch_size + ) + + prior_log_prob = prior.log_prob(theta) + + return log_likelihood + prior_log_prob + + _potential.device = torch.device(device) + + return _potential + + def likelihood_estimator_based_potential( likelihood_estimator: ConditionalDensityEstimator, prior: Distribution, # type: ignore diff --git a/sbi/inference/potentials/protocol.py b/sbi/inference/potentials/protocol.py new file mode 100644 index 000000000..414ceffe2 --- /dev/null +++ b/sbi/inference/potentials/protocol.py @@ -0,0 +1,106 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +"""Protocol definitions for stateless potential functions. + +This module provides the core protocol (Layer 1) for potential functions, +enabling stateless, composable potential evaluation. +""" + +from typing import Protocol + +import torch +from torch import Tensor + + +class PotentialFunction(Protocol): + """Stateless callable protocol for potential functions. + + A PotentialFunction takes parameters theta and observed data x, and returns + the unnormalized log-probability log p(theta, x). This is the core interface + that all potential functions must satisfy. + + The protocol is stateless - all inputs (theta, x) are passed as arguments, + with no hidden state dependencies. This enables: + - Easy testing (same inputs -> same outputs) + - Simple serialization + - Composable inference workflows + - Integration with external samplers (PyMC, custom PyTorch) + + Example: + >>> def my_potential(theta: Tensor, x: Tensor) -> Tensor: + ... ll = density_estimator.log_prob(x, context=theta) + ... return ll + prior.log_prob(theta) + >>> # Satisfies PotentialFunction protocol + """ + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + """Evaluate the potential function at theta given observed data x. + + Args: + theta: Parameter samples with shape (batch_size, *param_shape). + x: Observed data with shape (obs_batch, *obs_shape). + + Returns: + Log-probability with shape (batch_size,) or (obs_batch, batch_size) + depending on the potential implementation. + """ + ... + + +class PotentialFunctionWithDevice(Protocol): + """Extended protocol that reports device requirements. + + This protocol adds a device property to enable automatic device validation + and inference in bind_observation. + """ + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + """Evaluate the potential function at theta given observed data x.""" + ... + + @property + def device(self) -> torch.device: + """Return the expected device for inputs. + + This property allows bind_observation to validate device compatibility + and provide clear error messages for device mismatches. + """ + ... + + +def validate_potential(potential: PotentialFunction) -> bool: + """Validate that a potential function satisfies the protocol. + + This helper checks that the potential is callable and has the correct + signature. It performs runtime validation since Protocol typing + doesn't catch errors at definition time. + + Args: + potential: The potential function to validate. + + Returns: + True if the potential satisfies the protocol. + + Raises: + TypeError: If potential doesn't satisfy the protocol. + ValueError: If potential is callable but has wrong signature. + """ + if not callable(potential): + raise TypeError( + "Potential must be callable. " + "Expected a function or object with __call__ method." + ) + + import inspect + + sig = inspect.signature(potential.__call__) + params = list(sig.parameters.keys()) + + if len(params) < 2: + raise ValueError( + f"Potential __call__ must have at least 2 parameters (theta, x), " + f"got signature: {sig}" + ) + + return True diff --git a/tests/potential_protocol_test.py b/tests/potential_protocol_test.py new file mode 100644 index 000000000..6c125998b --- /dev/null +++ b/tests/potential_protocol_test.py @@ -0,0 +1,357 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +"""Unit tests for the stateless PotentialFunction protocol and bind_observation. + +These tests verify: +1. Protocol compliance for potential functions +2. bind_observation correctly binds x_o +3. IID handling (sum_iid parameter) +4. Device validation +5. Statelessness (repeated calls yield identical outputs) +""" + +import pytest +import torch +from torch import Tensor + +from sbi.inference.potentials.binding import BoundPotential, bind_observation +from sbi.inference.potentials.protocol import ( + validate_potential, +) + + +class SimplePotential: + """Simple potential function for testing.""" + + def __init__(self): + self.device = torch.device("cpu") + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + return -0.5 * ((theta - x.mean()).pow(2)).sum(dim=-1) + + +class DeviceAwarePotential: + """Potential with explicit device attribute.""" + + def __init__(self, device: torch.device): + self.device = device + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + return -0.5 * ((theta - x.mean()).pow(2)).sum(dim=-1) + + +def stateless_potential(theta: Tensor, x: Tensor) -> Tensor: + """Simple stateless potential function.""" + return -0.5 * ((theta - x.mean()).pow(2)).sum(dim=-1) + + +class TestPotentialFunctionProtocol: + """Tests for PotentialFunction protocol.""" + + def test_validate_callable_function(self): + """Test that validate_potential accepts valid functions.""" + assert validate_potential(stateless_potential) is True + + def test_validate_callable_class(self): + """Test that validate_potential accepts callable classes.""" + assert validate_potential(SimplePotential()) is True + + def test_validate_non_callable_raises(self): + """Test that validate_potential rejects non-callables.""" + with pytest.raises(TypeError, match="must be callable"): + validate_potential("not a potential") + + def test_validate_missing_call_raises(self): + """Test that validate_potential catches missing __call__.""" + + class NoCall: + pass + + with pytest.raises(TypeError, match="must be callable"): + validate_potential(NoCall()) + + +class TestBindObservation: + """Tests for bind_observation utility.""" + + def test_bind_simple_potential(self): + """Test binding a simple stateless potential.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0], [2.0]]) + result = bound_fn(theta) + + assert result.shape == (3,) + assert torch.isfinite(result).all() + + def test_bind_preserves_device(self): + """Test that bound function works on correct device.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0]]) + result = bound_fn(theta) + + assert result.device == x_o.device + + def test_iid_summing(self): + """Test that sum_iid=True correctly sums over IID observations.""" + x_o = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + def potential(theta: Tensor, x: Tensor) -> Tensor: + return -((theta.unsqueeze(0) - x.unsqueeze(1)).pow(2)).sum(dim=-1) + + bound_fn_sum = bind_observation(potential, x_o, sum_iid=True) + bound_fn_no_sum = bind_observation(potential, x_o, sum_iid=False) + + theta = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) + + result_sum = bound_fn_sum(theta) + result_no_sum = bound_fn_no_sum(theta) + + assert result_sum.shape == (2,) + assert result_no_sum.shape == (2, 2) + + def test_single_observation(self): + """Test binding with single observation (no IID dimension).""" + x_o = torch.tensor([1.0, 2.0]) + + def potential(theta: Tensor, x: Tensor) -> Tensor: + return -((theta - x).pow(2)).sum(dim=-1) + + bound_fn = bind_observation(potential, x_o, sum_iid=True) + theta = torch.tensor([[0.0, 0.0]]) + result = bound_fn(theta) + + assert result.shape == (1,) + assert torch.isfinite(result).all() + + def test_device_mismatch_raises(self): + """Test that device mismatch raises clear error.""" + potential = DeviceAwarePotential(torch.device("cpu")) + x_o = torch.tensor([1.0, 2.0, 3.0]) + + if torch.cuda.is_available(): + x_o_cuda = x_o.to("cuda") + with pytest.raises(ValueError, match="Device mismatch"): + bind_observation(potential, x_o_cuda, sum_iid=True) + + +class TestBoundPotential: + """Tests for BoundPotential class.""" + + def test_bound_potential_callable(self): + """Test that BoundPotential is callable.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound = BoundPotential(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0], [2.0]]) + result = bound(theta) + + assert result.shape == (3,) + assert torch.isfinite(result).all() + + def test_bound_potential_device_property(self): + """Test device property returns correct device.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound = BoundPotential(stateless_potential, x_o, sum_iid=True) + + assert bound.device == x_o.device + + def test_bound_potential_to_method(self): + """Test .to() method changes device.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound = BoundPotential(stateless_potential, x_o, sum_iid=True) + + if torch.cuda.is_available(): + bound_cuda = bound.to("cuda") + assert bound_cuda.device.type == "cuda" + + +class TestStatelessness: + """Tests verifying stateless behavior - key property of the new design.""" + + def test_repeated_calls_same_output(self): + """Test that repeated calls with same inputs yield identical outputs.""" + x_o = torch.tensor([1.0, 2.0, 3.0, 4.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0], [2.0]]) + + result1 = bound_fn(theta.clone()) + result2 = bound_fn(theta.clone()) + result3 = bound_fn(theta.clone()) + + assert torch.allclose(result1, result2) + assert torch.allclose(result2, result3) + + def test_different_theta_different_output(self): + """Test that different theta values yield different outputs.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta1 = torch.tensor([[0.0]]) + theta2 = torch.tensor([[1.0]]) + theta3 = torch.tensor([[2.0]]) + + result1 = bound_fn(theta1) + result2 = bound_fn(theta2) + result3 = bound_fn(theta3) + + assert not torch.allclose(result1, result2) + assert not torch.allclose(result2, result3) + + +class TestIntegration: + """Integration tests with realistic potential functions.""" + + def test_with_prior_and_likelihood(self): + """Test binding a potential combining prior and likelihood.""" + prior = torch.distributions.Normal(0.0, 1.0) + + class MockEstimator: + def log_prob(self, x: Tensor, context: Tensor) -> Tensor: + return -((x - context).pow(2)).sum(dim=-1) + + estimator = MockEstimator() + + def potential(theta: Tensor, x: Tensor) -> Tensor: + likelihood = estimator.log_prob(x, context=theta) + prior_log_prob = prior.log_prob(theta).sum(dim=-1) + return likelihood + prior_log_prob + + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(potential, x_o, sum_iid=True) + + theta = torch.randn(10, 1) + result = bound_fn(theta) + + assert result.shape == (10,) + assert torch.isfinite(result).all() + + def test_with_gradient_tracking(self): + """Test that gradients can be tracked through bound potential.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.randn(5, 1, requires_grad=True) + result = bound_fn(theta) + + assert result.requires_grad + result.sum().backward() + assert theta.grad is not None + + +class TestNewStatelessAPI: + """Tests for the new stateless potential functions (likelihood_potential, etc.).""" + + def test_likelihood_potential_creates_valid_function(self): + """Test that likelihood_potential returns a valid PotentialFunction.""" + from sbi.inference import NLE + from sbi.inference.potentials import bind_observation, likelihood_potential + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) + theta = prior.sample((50,)) + x = theta + torch.randn_like(theta) * 0.1 + x_o = torch.randn(1, 2) + + trainer = NLE() + estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + stateless_fn = likelihood_potential(estimator, prior) + + assert hasattr(stateless_fn, "device") + assert callable(stateless_fn) + + bound_fn = bind_observation(stateless_fn, x_o) + theta_test = torch.rand(5, 2) + result = bound_fn(theta_test) + + assert result.shape == (5,) + assert torch.isfinite(result).all() + + def test_new_stateless_functions_satisfy_protocol(self): + """Test that new stateless functions satisfy PotentialFunction protocol.""" + from sbi.inference import NLE + from sbi.inference.potentials import likelihood_potential, validate_potential + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) + theta = prior.sample((30,)) + x = theta + torch.randn_like(theta) * 0.1 + + trainer = NLE() + estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + stateless_fn = likelihood_potential(estimator, prior) + assert validate_potential(stateless_fn) is True + + +class TestBackwardCompatibility: + """Tests verifying backward compatibility with existing sbi workflows.""" + + def test_existing_likelihood_potential_still_works(self): + """Test that existing likelihood_estimator_based_potential works.""" + from sbi.inference import NLE + from sbi.inference.potentials.likelihood_based_potential import ( + likelihood_estimator_based_potential, + ) + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3)) + theta = prior.sample((50,)) + x = theta + torch.randn_like(theta) * 0.1 + x_o = torch.randn(1, 3) + + trainer = NLE() + likelihood_estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + potential_fn, parameter_transform = likelihood_estimator_based_potential( + likelihood_estimator, prior, x_o + ) + + assert hasattr(potential_fn, "x_o") + assert hasattr(potential_fn, "set_x") + + theta_test = torch.randn(5, 3) + result = potential_fn(theta_test) + assert result.shape == (5,) + + def test_validate_potential_with_real_potential(self): + """Test that validate_potential works with existing potentials.""" + from sbi.inference import NLE + from sbi.inference.potentials.likelihood_based_potential import ( + LikelihoodBasedPotential, + ) + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) + theta = prior.sample((30,)) + x = theta + torch.randn_like(theta) * 0.1 + + trainer = NLE() + likelihood_estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + potential = LikelihoodBasedPotential(likelihood_estimator, prior) + potential.set_x(torch.randn(1, 2)) + + assert validate_potential(potential) is True