PR #1-2, Weeks 3-5: implement core protocol + bind_observation + minimal notebook example

.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:

.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:

.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

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
+}

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 <https://www.apache.org/licenses/>
 
+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",
 ]

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 <https://www.apache.org/licenses/>
 
+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)