Merge remote-tracking branch 'upstream/main' into simformer

Author: nMaax

.github/workflows/publish.yml

@@ -76,14 +76,13 @@ jobs:
         inputs: >-
           ./dist/*.tar.gz
           ./dist/*.whl
-    - name: Create GitHub Release
+    - name: Ensure GitHub Release exists (no-op if already exists)
       env:
         GITHUB_TOKEN: ${{ github.token }}
-      run: >-
-        gh release create
-        '${{ github.ref_name }}'
-        --repo '${{ github.repository }}'
-        --notes ""
+      run: |
+        # If a release for this tag already exists (e.g., created via GH UI), skip creation.
+        gh release view '${{ github.ref_name }}' --repo '${{ github.repository }}' >/dev/null 2>&1 || \
+          gh release create '${{ github.ref_name }}' --repo '${{ github.repository }}' --notes ""
     - name: Upload artifact signatures to GitHub Release
       env:
         GITHUB_TOKEN: ${{ github.token }}

CHANGELOG.md

@@ -1,19 +1,17 @@
 # Changelog
 
-# Changelog
-
 ## v0.25.0
 
 ### ✨ Highlights
 
 #### 🚀 New Inference Methods
 
-* **MNPE class similar to MNLE** by @dgedon in https://github.com/sbi-dev/sbi/pull/1362
-* **Implementing SNPE-B (#199)** by @etouron1 in https://github.com/sbi-dev/sbi/pull/1471
+* **MNPE class mixed parameter (similar to MNLE)** by @dgedon in https://github.com/sbi-dev/sbi/pull/1362
+* **Implementation of SNPE-B (#199)** by @etouron1 in https://github.com/sbi-dev/sbi/pull/1471
 
 #### 🧠 Neural Network Architectures & Embedding Networks
 
-* **Simple transformer implementation** by @NicolasRR in https://github.com/sbi-dev/sbi/pull/1494
+* **Add transformer embedding net** by @NicolasRR in https://github.com/sbi-dev/sbi/pull/1494
 * **Add embedding net that uses 1D causal convolutions (#1459)** by @Aranka-S in https://github.com/sbi-dev/sbi/pull/1499
 * **Add LRU-backed embedding networks** by @famura in https://github.com/sbi-dev/sbi/pull/1512
 * **Add ResNet as embedding model** by @StefanWahl in https://github.com/sbi-dev/sbi/pull/1472
@@ -22,12 +20,13 @@
 #### ⭐ Major Features & Capabilities
 
 * **Unify flow matching and score-based models** by @StarostinV in https://github.com/sbi-dev/sbi/pull/1497
-* **Model misspecification based on MMD** by @coschroeder in https://github.com/sbi-dev/sbi/pull/1502
+* **Model misspecification detection based on MMD** by @coschroeder in https://github.com/sbi-dev/sbi/pull/1502
 * **Marginal estimator log-prob based test for misspecification** by @swag2198 in https://github.com/sbi-dev/sbi/pull/1522
 * **Adding interface for unconditional flow training** by @plcrodrigues in https://github.com/sbi-dev/sbi/pull/1470
 * **Support using trained estimators in Pyro models** by @sethaxen in https://github.com/sbi-dev/sbi/pull/1491
 * **Add util to generate mcmc samples from user defined potential (#1405)** by @hayden-johnson in https://github.com/sbi-dev/sbi/pull/1483
 * **Logit transform** by @anastasiakrouglova in https://github.com/sbi-dev/sbi/pull/1485
+* **Log-prob for iid data for score estimators** by @Kartik-Sama in https://github.com/sbi-dev/sbi/pull/1508
 
 #### 📚 Documentation & Tutorials
 
@@ -46,6 +45,8 @@
 * fix: cap max_sampling_batch_size to prevent excessive memory by @janfb in https://github.com/sbi-dev/sbi/pull/1624
 * 1561 computation of denoising posterior precision matrix in jac method score fn iid by @manuelgloeckler in https://github.com/sbi-dev/sbi/pull/1636
 * fix xfail test, fix deprecation warnings by @janfb in https://github.com/sbi-dev/sbi/pull/1642
+* fix: iid-score device handling by @janfb in https://github.com/sbi-dev/sbi/pull/1650
+* fix: fmpe singularity on sde sampling by @manuelgloeckler in https://github.com/sbi-dev/sbi/pull/1661
 
 ### 🛠️ Maintenance & Improvements
 
@@ -60,6 +61,7 @@
 * Use TypeAlias and consistent naming for sbi types by @janfb in https://github.com/sbi-dev/sbi/pull/1637
 * Add protocol for estimator builder by @abelaba in https://github.com/sbi-dev/sbi/pull/1633
 * Improve abc implementation by @janfb in https://github.com/sbi-dev/sbi/pull/1615
+* Refactor RatioEstimator to subclass ConditionalEstimator @abelaba in https://github.com/sbi-dev/sbi/pull/1652
 
 #### 🏷️ Type Hints & API Improvements
 
@@ -86,6 +88,8 @@
 * chore: reorder setup steps for Python and uv in CI/CD workflows by @janfb in https://github.com/sbi-dev/sbi/pull/1601
 * Fix/lc2st numpy type fixes by @janfb in https://github.com/sbi-dev/sbi/pull/1613
 * Fix failing CI on main. by @janfb in https://github.com/sbi-dev/sbi/pull/1618
+* Fix slow vector field tests by @janfb in https://github.com/sbi-dev/sbi/pull/1657
+* Add tests for sensitivity analysis by @janfb in https://github.com/sbi-dev/sbi/pull/1662
 
 #### 📖 Documentation & Website
 
@@ -113,6 +117,9 @@
 * fixed misrendered bullet list, tested locally by @psteinb in https://github.com/sbi-dev/sbi/pull/1594
 * Improvements to L-C2ST tutorial by @michaeldeistler in https://github.com/sbi-dev/sbi/pull/1588
 * docs: Change colortheme in light mode by @michaeldeistler in https://github.com/sbi-dev/sbi/pull/1638
+* Posterior parameters doc by @abelaba in https://github.com/sbi-dev/sbi/pull/1644
+* fix contributing links by @janfb in https://github.com/sbi-dev/sbi/pull/1647
+* docs: add posterior parameters dataclass how to guide by @abelaba in https://github.com/sbi-dev/sbi/pull/1654
 
 #### 🏗️ Infrastructure & Dependencies
 

README.md

@@ -1,7 +1,7 @@
 [![PyPI version](https://badge.fury.io/py/sbi.svg)](https://badge.fury.io/py/sbi)
 [![Conda Version](https://img.shields.io/conda/vn/conda-forge/sbi.svg)](https://github.com/conda-forge/sbi-feedstock)
 [![Contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://sbi.readthedocs.io/en/latest/contributing.html)
-[![Tests](https://github.com/sbi-dev/sbi/actions/workflows/ci.yml/badge.svg)](https://github.com/sbi-dev/sbi/actions)
+[![Tests](https://github.com/sbi-dev/sbi/actions/workflows/cd.yml/badge.svg)](https://github.com/sbi-dev/sbi/actions)
 [![codecov](https://codecov.io/gh/sbi-dev/sbi/branch/main/graph/badge.svg)](https://codecov.io/gh/sbi-dev/sbi)
 [![GitHub license](https://img.shields.io/github/license/sbi-dev/sbi)](https://github.com/sbi-dev/sbi/blob/master/LICENSE.txt)
 [![DOI](https://joss.theoj.org/papers/10.21105/joss.07754/status.svg)](https://doi.org/10.21105/joss.07754)

docs/how_to_guide.rst

@@ -79,3 +79,12 @@ Visualization
    :maxdepth: 1
 
    how_to_guide/05_conditional_distributions.ipynb
+
+
+Posterior Parameters
+--------------------
+
+.. toctree::
+   :maxdepth: 1
+
+   how_to_guide/19_posterior_parameters.ipynb

docs/how_to_guide/19_posterior_parameters.ipynb

@@ -29,7 +29,7 @@
    "id": "8c4df01a",
    "metadata": {},
    "source": [
-    "# Usage"
+    "## Usage"
    ]
   },
   {

sbi/analysis/sensitivity_analysis.py

@@ -47,7 +47,9 @@ def destandardizing_net(batch_t: Tensor, min_std: float = 1e-7) -> nn.Module:
     is_valid_t, *_ = handle_invalid_x(batch_t, True)
 
     t_mean = torch.mean(batch_t[is_valid_t], dim=0)
-    if len(batch_t > 1):
+
+    # Use batch size to decide whether a reliable std can be computed.
+    if len(batch_t) > 1:
         t_std = torch.std(batch_t[is_valid_t], dim=0)
         t_std[t_std < min_std] = min_std
     else:

sbi/diagnostics/misspecification.py

@@ -8,7 +8,6 @@
 from typing import Optional
 
 import torch
-import torch.nn as nn
 from torch import Tensor
 
 from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer
@@ -149,25 +148,13 @@ def calc_misspecification_mmd(
                 "no neural net found,"
                 "neural_net should not be None when mode is 'embedding'"
             )
-        neural_net = inference._neural_net
-        if neural_net is None:
-            raise ValueError(
-                "no neural net found,"
-                "neural_net should not be None when mode is 'embedding'"
-            )
-        if neural_net.embedding_net is None:
-            raise ValueError(
-                "no embedding net found,"
-                "embedding_net should not be None when mode is 'embedding'"
-            )
-        if isinstance(neural_net.embedding_net, nn.modules.linear.Identity):
-            warnings.warn(
-                "The embedding net might be the identity function,"
-                "in that case the MMD is computed in the x-space.",
-                stacklevel=2,
+        if inference._neural_net.embedding_net is None:
+            raise AttributeError(
+                "embedding_net attribute is None but is required for misspecification "
+                "detection."
             )
-        z_obs = neural_net.embedding_net(x_obs).detach()
-        z = neural_net.embedding_net(x).detach()
+        z_obs = inference._neural_net.embedding_net(x_obs).detach()
+        z = inference._neural_net.embedding_net(x).detach()
     else:
         raise ValueError("mode should be either 'x_space' or 'embedding'")
 

sbi/inference/posteriors/vector_field_posterior.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 math
 import warnings
 from typing import Dict, Literal, Optional, Union
 
@@ -150,7 +151,9 @@ def sample(
         corrector_params: Optional[Dict] = None,
         steps: int = 500,
         ts: Optional[Tensor] = None,
-        iid_method: Literal["fnpe", "gauss", "auto_gauss", "jac_gauss"] = "auto_gauss",
+        iid_method: Optional[
+            Literal["fnpe", "gauss", "auto_gauss", "jac_gauss"]
+        ] = None,
         iid_params: Optional[Dict] = None,
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
@@ -210,19 +213,22 @@ def sample(
 
         is_iid = x.shape[0] > 1
         self.potential_fn.set_x(
-            x, x_is_iid=is_iid, iid_method=iid_method, iid_params=iid_params
+            x,
+            x_is_iid=is_iid,
+            iid_method=iid_method or self.potential_fn.iid_method,
+            iid_params=iid_params,
         )
 
         num_samples = torch.Size(sample_shape).numel()
 
         if sample_with == "ode":
-            samples = rejection.accept_reject_sample(
+            samples, _ = rejection.accept_reject_sample(
                 proposal=self.sample_via_ode,
                 accept_reject_fn=lambda theta: within_support(self.prior, theta),
                 num_samples=num_samples,
                 show_progress_bars=show_progress_bars,
                 max_sampling_batch_size=max_sampling_batch_size,
-            )[0]
+            )
         elif sample_with == "sde":
             proposal_sampling_kwargs = {
                 "predictor": predictor,
@@ -234,14 +240,14 @@ def sample(
                 "max_sampling_batch_size": max_sampling_batch_size,
                 "show_progress_bars": show_progress_bars,
             }
-            samples = rejection.accept_reject_sample(
+            samples, _ = rejection.accept_reject_sample(
                 proposal=self._sample_via_diffusion,
                 accept_reject_fn=lambda theta: within_support(self.prior, theta),
                 num_samples=num_samples,
                 show_progress_bars=show_progress_bars,
                 max_sampling_batch_size=max_sampling_batch_size,
                 proposal_sampling_kwargs=proposal_sampling_kwargs,
-            )[0]
+            )
         else:
             raise ValueError(
                 f"Expected sample_with to be 'ode' or 'sde', but got {sample_with}."
@@ -263,6 +269,7 @@ def _sample_via_diffusion(
         ts: Optional[Tensor] = None,
         max_sampling_batch_size: int = 10_000,
         show_progress_bars: bool = True,
+        save_intermediate: bool = False,
     ) -> Tensor:
         r"""Return samples from posterior distribution $p(\theta|x)$.
 
@@ -284,20 +291,26 @@ def _sample_via_diffusion(
             sample_with: Deprecated - use `.build_posterior(sample_with=...)` prior to
                 `.sample()`.
             show_progress_bars: Whether to show a progress bar during sampling.
+            save_intermediate: Whether to save intermediate results of the diffusion
+                process. If True, the returned tensor has shape
+                `(*sample_shape, steps, *input_shape)`.
         """
 
         if not self.vector_field_estimator.SCORE_DEFINED:
             raise ValueError(
                 "The vector field estimator does not support the 'sde' sampling method."
             )
 
-        num_samples = torch.Size(sample_shape).numel()
+        total_samples_needed = torch.Size(sample_shape).numel()
 
-        max_sampling_batch_size = (
+        # Determine effective batch size for sampling
+        effective_batch_size = (
             self.max_sampling_batch_size
             if max_sampling_batch_size is None
             else max_sampling_batch_size
         )
+        # Ensure we don't use larger batches than total samples needed
+        effective_batch_size = min(effective_batch_size, total_samples_needed)
 
         # TODO: the time schedule should be provided by the estimator, see issue #1437
         if ts is None:
@@ -306,28 +319,46 @@ def _sample_via_diffusion(
             ts = torch.linspace(t_max, t_min, steps)
         ts = ts.to(self.device)
 
+        # Initialize the diffusion sampler
         diffuser = Diffuser(
             self.potential_fn,
             predictor=predictor,
             corrector=corrector,
             predictor_params=predictor_params,
             corrector_params=corrector_params,
         )
-        max_sampling_batch_size = min(max_sampling_batch_size, num_samples)
-        samples = []
-        num_iter = num_samples // max_sampling_batch_size
-        num_iter = (
-            num_iter + 1 if (num_samples % max_sampling_batch_size) != 0 else num_iter
-        )
-        for _ in range(num_iter):
-            samples.append(
-                diffuser.run(
-                    num_samples=max_sampling_batch_size,
-                    ts=ts,
-                    show_progress_bars=show_progress_bars,
-                )
+
+        # Calculate how many batches we need
+        num_batches = math.ceil(total_samples_needed / effective_batch_size)
+
+        # Generate samples in batches
+        all_samples = []
+        samples_generated = 0
+
+        for _ in range(num_batches):
+            # Calculate how many samples to generate in this batch
+            remaining_samples = total_samples_needed - samples_generated
+            current_batch_size = min(effective_batch_size, remaining_samples)
+
+            # Generate samples for this batch
+            batch_samples = diffuser.run(
+                num_samples=current_batch_size,
+                ts=ts,
+                show_progress_bars=show_progress_bars,
+                save_intermediate=save_intermediate,
+            )
+
+            all_samples.append(batch_samples)
+            samples_generated += current_batch_size
+
+        # Concatenate all batches and ensure we return exactly the requested number
+        samples = torch.cat(all_samples, dim=0)[:total_samples_needed]
+
+        if torch.isnan(samples).all():
+            raise RuntimeError(
+                "All samples NaN after diffusion sampling. "
+                "This may indicate numerical instability in the vector field."
             )
-        samples = torch.cat(samples, dim=0)[:num_samples]
 
         return samples
 
@@ -382,7 +413,10 @@ def log_prob(
             `(len(θ),)`-shaped log posterior probability $\log p(\theta|x)$ for θ in the
             support of the prior, -∞ (corresponding to 0 probability) outside.
         """
-        self.potential_fn.set_x(self._x_else_default_x(x), **(ode_kwargs or {}))
+        x = self._x_else_default_x(x)
+        x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
+        is_iid = x.shape[0] > 1
+        self.potential_fn.set_x(x, x_is_iid=is_iid, **(ode_kwargs or {}))
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(
@@ -461,14 +495,14 @@ def sample_batched(
             max_sampling_batch_size = capped
 
         if self.sample_with == "ode":
-            samples = rejection.accept_reject_sample(
+            samples, _ = rejection.accept_reject_sample(
                 proposal=self.sample_via_ode,
                 accept_reject_fn=lambda theta: within_support(self.prior, theta),
                 num_samples=num_samples,
                 num_xos=batch_size,
                 show_progress_bars=show_progress_bars,
                 max_sampling_batch_size=max_sampling_batch_size,
-            )[0]
+            )
             samples = samples.reshape(
                 sample_shape + batch_shape + self.vector_field_estimator.input_shape
             )
@@ -483,15 +517,15 @@ def sample_batched(
                 "max_sampling_batch_size": max_sampling_batch_size,
                 "show_progress_bars": show_progress_bars,
             }
-            samples = rejection.accept_reject_sample(
+            samples, _ = rejection.accept_reject_sample(
                 proposal=self._sample_via_diffusion,
                 accept_reject_fn=lambda theta: within_support(self.prior, theta),
                 num_samples=num_samples,
                 num_xos=batch_size,
                 show_progress_bars=show_progress_bars,
                 max_sampling_batch_size=max_sampling_batch_size,
                 proposal_sampling_kwargs=proposal_sampling_kwargs,
-            )[0]
+            )
             samples = samples.reshape(
                 sample_shape + batch_shape + self.vector_field_estimator.input_shape
             )