From 7c020f873d6f89196b918ad921bb2adcc6db0647 Mon Sep 17 00:00:00 2001 From: Jan Date: Wed, 1 Jul 2026 10:59:01 +0200 Subject: [PATCH 1/4] fix: make pymc HMC c2st test reliable via target_accept The hmc parametrizations of test_c2st_pymc_sampler_on_Gaussian were failing in CD (c2st ~0.62). Root cause is PyMC's HamiltonianMC default target_accept=0.65, which mixes poorly on this peaked target regardless of seed, warmup, or number of draws (NUTS is unaffected). Expose seed and target_accept on PyMCSampler and set target_accept=0.95 (c2st ~0.51-0.54 across seeds and chain counts) plus a fixed seed in the test for reproducibility. --- sbi/samplers/mcmc/pymc_wrapper.py | 19 ++++++++++++++++++- tests/mcmc_test.py | 4 ++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/sbi/samplers/mcmc/pymc_wrapper.py b/sbi/samplers/mcmc/pymc_wrapper.py index 78e8ef57f..2933e31f9 100644 --- a/sbi/samplers/mcmc/pymc_wrapper.py +++ b/sbi/samplers/mcmc/pymc_wrapper.py @@ -114,6 +114,8 @@ def __init__( progressbar: bool = True, param_name: str = "theta", device: str = "cpu", + seed: Optional[int] = None, + target_accept: Optional[float] = None, ): """Interface for PyMC samplers @@ -128,6 +130,13 @@ def __init__( progressbar: Whether to show/hide progress bars. param_name: Name for parameter variable, for PyMC and ArviZ structures device: The device to which to move the parameters for potential_fn. + seed: Random seed passed to `pymc.sample` for reproducible sampling. + If None (default), PyMC seeds from system entropy. + target_accept: Target acceptance probability for the `"hmc"` and + `"nuts"` step methods. If None (default), PyMC's own default is + used. PyMC's HMC default (0.65) mixes poorly on peaked targets; + a higher value (e.g. 0.9-0.95) gives markedly less biased samples. + Ignored for the `"slice"` step. """ self.param_name = param_name self._step = step @@ -138,6 +147,8 @@ def __init__( self._mp_ctx = mp_ctx self._progressbar = progressbar self._device = device + self._seed = seed + self._target_accept = target_accept # create PyMC model object track_gradients = step in ("nuts", "hmc") @@ -157,15 +168,21 @@ def run(self) -> np.ndarray: MCMC samples """ step_class = dict(slice=pymc.Slice, hmc=pymc.HamiltonianMC, nuts=pymc.NUTS) + # target_accept only applies to the gradient-based samplers; Slice + # does not take it. + step_kwargs = {} + if self._target_accept is not None and self._step in ("hmc", "nuts"): + step_kwargs["target_accept"] = self._target_accept with self._model: inference_data = pymc.sample( - step=step_class[self._step](), + step=step_class[self._step](**step_kwargs), tune=self._tune, draws=self._draws, initvals=self._initvals, # type: ignore chains=self._chains, progressbar=self._progressbar, mp_ctx=self._mp_ctx, + random_seed=self._seed, ) self._inference_data = inference_data traces = inference_data.posterior # type: ignore diff --git a/tests/mcmc_test.py b/tests/mcmc_test.py index 13fa40264..dc4153d5f 100644 --- a/tests/mcmc_test.py +++ b/tests/mcmc_test.py @@ -158,6 +158,10 @@ def lp_f(x, track_gradients=True): draws=(int(num_samples / num_chains)), # PyMC does not use thinning tune=warmup, chains=num_chains, + seed=0, # reproducible sampling; PyMC is not covered by the global seed + # PyMC's HMC default (target_accept=0.65) is biased on this peaked + # target (c2st ~0.63); 0.95 samples it accurately. See #1894 follow-up. + target_accept=0.95, ) samples = sampler.run() assert samples.shape == ( From f738493eb3af093d0c19257c689afaada4c238dd Mon Sep 17 00:00:00 2001 From: Jan Date: Wed, 1 Jul 2026 11:01:40 +0200 Subject: [PATCH 2/4] fix: raise pymc HMC target_accept in production sampling MCMCPosterior._pymc_mcmc left PyMC's HamiltonianMC at its default target_accept=0.65, which produces biased hmc_pymc posteriors (c2st ~0.62 on a Gaussian). Default it to 0.9 for the HMC step, leaving nuts_pymc and slice_pymc unchanged. hmc_pymc now samples accurately (c2st ~0.51). --- sbi/inference/posteriors/mcmc_posterior.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py index f90db67d6..b9810c053 100644 --- a/sbi/inference/posteriors/mcmc_posterior.py +++ b/sbi/inference/posteriors/mcmc_posterior.py @@ -894,10 +894,16 @@ def _pymc_mcmc( thin = _process_thin_default(thin) num_chains = mp.cpu_count() - 1 if num_chains is None else num_chains steps = dict(slice_pymc="slice", hmc_pymc="hmc", nuts_pymc="nuts") + step = steps[mcmc_method] + + # PyMC's HamiltonianMC default (target_accept=0.65) mixes poorly on + # peaked targets and yields biased samples. Use a higher value for HMC; + # NUTS keeps PyMC's default (0.8), which already samples well. + target_accept = 0.9 if step == "hmc" else None sampler = PyMCSampler( potential_fn=potential_function, - step=steps[mcmc_method], + step=step, initvals=tensor2numpy(initial_params), draws=ceil((thin * num_samples) / num_chains), tune=warmup_steps, @@ -906,6 +912,7 @@ def _pymc_mcmc( progressbar=show_progress_bars, param_name=self.param_name, device=self._device, + target_accept=target_accept, ) samples = sampler.run() samples = torch.from_numpy(samples).to(dtype=torch.float32, device=self._device) From 2d5a99f3aa14eac87852e6ce5523b2b1860cdc7f Mon Sep 17 00:00:00 2001 From: Jan Date: Thu, 2 Jul 2026 06:33:05 +0200 Subject: [PATCH 3/4] feat: expose target_accept for HMC/NUTS MCMC samplers --- sbi/inference/posteriors/mcmc_posterior.py | 30 +++++++++++++++---- .../posteriors/posterior_parameters.py | 11 +++++++ tests/mcmc_test.py | 30 +++++++++++++++++++ tests/posterior_parameters_test.py | 14 +++++++++ 4 files changed, 80 insertions(+), 5 deletions(-) diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py index b9810c053..db23fb25d 100644 --- a/sbi/inference/posteriors/mcmc_posterior.py +++ b/sbi/inference/posteriors/mcmc_posterior.py @@ -64,6 +64,7 @@ def __init__( mp_context: Literal["fork", "spawn"] = "spawn", device: Optional[Union[str, torch.device]] = None, x_shape: Optional[torch.Size] = None, + target_accept: Optional[float] = None, ): """ Args: @@ -107,6 +108,9 @@ def __init__( device: Training device, e.g., "cpu", "cuda" or "cuda:0". If None, `potential_fn.device` is used. x_shape: Deprecated, should not be passed. + target_accept: Target acceptance probability for the gradient-based + samplers (HMC/NUTS). See `MCMCPosteriorParameters` for details. If + `None`, HMC uses `0.9` and NUTS/Pyro keep their backend defaults. """ if method == "slice": warn( @@ -136,6 +140,7 @@ def __init__( self.init_strategy_parameters = init_strategy_parameters or {} self.num_workers = num_workers self.mp_context = mp_context + self.target_accept = target_accept self._posterior_sampler = None # Hardcode parameter name to reduce clutter kwargs. @@ -257,6 +262,7 @@ def sample( num_workers: Optional[int] = None, mp_context: Optional[str] = None, show_progress_bars: bool = True, + target_accept: Optional[float] = None, ) -> Tensor: r"""Draw samples from the approximate posterior distribution $p(\theta|x)$. @@ -285,6 +291,9 @@ def sample( mp_context: Multiprocessing context (`fork` or `spawn`). If not provided, uses the value specified at initialization. show_progress_bars: Whether to show sampling progress monitor. + target_accept: Target acceptance probability for the gradient-based + samplers (HMC/NUTS). If not provided, uses the value specified at + initialization. Returns: Samples from posterior. @@ -301,6 +310,7 @@ def sample( init_strategy = self.init_strategy if init_strategy is None else init_strategy num_workers = self.num_workers if num_workers is None else num_workers mp_context = self.mp_context if mp_context is None else mp_context + target_accept = self.target_accept if target_accept is None else target_accept init_strategy_parameters = ( self.init_strategy_parameters if init_strategy_parameters is None @@ -342,6 +352,7 @@ def sample( num_chains=num_chains, show_progress_bars=show_progress_bars, mp_context=mp_context, + target_accept=target_accept, ) elif method in ("hmc_pymc", "nuts_pymc", "slice_pymc"): transformed_samples = self._pymc_mcmc( @@ -354,6 +365,7 @@ def sample( num_chains=num_chains, show_progress_bars=show_progress_bars, mp_context=mp_context, + target_accept=target_accept, ) else: raise NameError(f"The sampling method {method} is not implemented!") @@ -794,6 +806,7 @@ def _pyro_mcmc( num_chains: Optional[int] = 1, show_progress_bars: bool = True, mp_context: str = "spawn", + target_accept: Optional[float] = None, ) -> Tensor: r"""Return samples obtained using Pyro's HMC or NUTS sampler. @@ -809,6 +822,8 @@ def _pyro_mcmc( warmup_steps: Initial number of samples to discard. num_chains: Whether to sample in parallel. If None, use all but one CPU. show_progress_bars: Whether to show a progressbar during sampling. + target_accept: Target acceptance probability for step-size adaptation. + If None, Pyro's default is used. Returns: Tensor of shape (num_samples, shape_of_single_theta). @@ -827,9 +842,12 @@ def _pyro_mcmc( thin = _process_thin_default(thin) num_chains = mp.cpu_count() - 1 if num_chains is None else num_chains kernels = dict(hmc_pyro=HMC, nuts_pyro=NUTS) + kernel_kwargs = {"potential_fn": potential_function} + if target_accept is not None: + kernel_kwargs["target_accept_prob"] = target_accept sampler = MCMC( - kernel=kernels[mcmc_method](potential_fn=potential_function), + kernel=kernels[mcmc_method](**kernel_kwargs), num_samples=ceil((thin * num_samples) / num_chains), warmup_steps=warmup_steps, initial_params={self.param_name: initial_params}, @@ -862,6 +880,7 @@ def _pymc_mcmc( num_chains: Optional[int] = 1, show_progress_bars: bool = True, mp_context: str = "spawn", + target_accept: Optional[float] = None, ) -> Tensor: r"""Return samples obtained using PyMC's HMC, NUTS or slice samplers. @@ -877,6 +896,9 @@ def _pymc_mcmc( warmup_steps: Initial number of samples to discard. num_chains: Whether to sample in parallel. If None, use all but one CPU. show_progress_bars: Whether to show a progressbar during sampling. + target_accept: Target acceptance probability for HMC/NUTS. If None, + HMC uses 0.9 and NUTS keeps PyMC's default. Ignored by the slice + sampler. Returns: Tensor of shape (num_samples, shape_of_single_theta). @@ -896,10 +918,8 @@ def _pymc_mcmc( steps = dict(slice_pymc="slice", hmc_pymc="hmc", nuts_pymc="nuts") step = steps[mcmc_method] - # PyMC's HamiltonianMC default (target_accept=0.65) mixes poorly on - # peaked targets and yields biased samples. Use a higher value for HMC; - # NUTS keeps PyMC's default (0.8), which already samples well. - target_accept = 0.9 if step == "hmc" else None + if target_accept is None and step == "hmc": + target_accept = 0.9 sampler = PyMCSampler( potential_fn=potential_function, diff --git a/sbi/inference/posteriors/posterior_parameters.py b/sbi/inference/posteriors/posterior_parameters.py index 09e3a90ad..336a721e1 100644 --- a/sbi/inference/posteriors/posterior_parameters.py +++ b/sbi/inference/posteriors/posterior_parameters.py @@ -244,6 +244,13 @@ class MCMCPosteriorParameters(PosteriorParameters): (default), used by Pyro and PyMC samplers. `"fork"` can be significantly faster than `"spawn"` but is only supported on POSIX-based systems (e.g. Linux and macOS, not Windows). + target_accept: Target acceptance probability for the gradient-based + samplers (`hmc_pymc`, `nuts_pymc`, `hmc_pyro`, `nuts_pyro`), controlling + the step-size adaptation during warmup. Higher values take smaller, + more accurate steps at the cost of speed. Ignored by the slice samplers. + If `None`, HMC uses `0.9` (PyMC's default of `0.65` mixes poorly on + peaked targets and yields biased samples), while NUTS and the Pyro + samplers keep their respective backend defaults. """ method: Literal[ @@ -262,10 +269,14 @@ class MCMCPosteriorParameters(PosteriorParameters): init_strategy_parameters: Optional[Dict[str, Any]] = None num_workers: int = 1 mp_context: Literal["fork", "spawn"] = "spawn" + target_accept: Optional[float] = None def validate(self): """Validate MCMCPosteriorParameters fields.""" + if self.target_accept is not None and not (0.0 < self.target_accept < 1.0): + raise ValueError("target_accept must be between 0 and 1, or None.") + if not ( self.init_strategy_parameters is None or isinstance(self.init_strategy_parameters, Dict) diff --git a/tests/mcmc_test.py b/tests/mcmc_test.py index dc4153d5f..4afa44ddf 100644 --- a/tests/mcmc_test.py +++ b/tests/mcmc_test.py @@ -177,6 +177,36 @@ def lp_f(x, track_gradients=True): check_c2st(samples, target_samples, alg=alg) +@pytest.mark.mcmc +@pytest.mark.parametrize( + ("method", "target_accept", "expected"), + [ + ("hmc_pymc", None, 0.9), # HMC default is raised (PyMC's 0.65 is biased) + ("hmc_pymc", 0.99, 0.99), # user override wins + ("nuts_pymc", None, None), # NUTS keeps PyMC's own default + ], +) +def test_pymc_target_accept_forwarded(method, target_accept, expected): + """target_accept from the user reaches the PyMC sampler with the right default.""" + theta_dim = 2 + prior = BoxUniform(low=-2 * ones(theta_dim), high=2 * ones(theta_dim)) + + def potential_fn(theta): + return -0.5 * (theta**2).sum(axis=-1) + + mcmc_posterior = build_from_potential(potential_fn, prior) + mcmc_posterior.sample( + (10,), + method=method, + target_accept=target_accept, + num_chains=1, + warmup_steps=10, + thin=1, + show_progress_bars=False, + ) + assert mcmc_posterior._posterior_sampler._target_accept == expected + + @pytest.mark.mcmc def test_direct_mcmc_unconditional(): "Test MCMCPosterior from user defined potential (unconditional)" diff --git a/tests/posterior_parameters_test.py b/tests/posterior_parameters_test.py index d3585d9d0..adf13608d 100644 --- a/tests/posterior_parameters_test.py +++ b/tests/posterior_parameters_test.py @@ -396,6 +396,20 @@ def test_invalid_literal_field_values(): MCMCPosteriorParameters(method="invalid") +@pytest.mark.parametrize("target_accept", [None, 0.5, 0.9, 0.99]) +def test_mcmc_target_accept_accepts_valid_values(target_accept): + """target_accept accepts None or a probability in (0, 1).""" + params = MCMCPosteriorParameters(target_accept=target_accept) + assert params.target_accept == target_accept + + +@pytest.mark.parametrize("target_accept", [0.0, 1.0, 1.5, -0.1]) +def test_mcmc_target_accept_rejects_invalid_values(target_accept): + """target_accept outside (0, 1) is rejected.""" + with pytest.raises(ValueError, match="target_accept"): + MCMCPosteriorParameters(target_accept=target_accept) + + @pytest.mark.parametrize( "params", [ From 52266baaf551a002e2734145300c64dfdb29ddef Mon Sep 17 00:00:00 2001 From: Jan Date: Mon, 20 Jul 2026 22:27:42 +0200 Subject: [PATCH 4/4] fix: refine PyMC target acceptance handling --- CHANGELOG.md | 1 + sbi/inference/posteriors/mcmc_posterior.py | 47 +++++++++-------- .../posteriors/posterior_parameters.py | 24 +++++---- sbi/samplers/mcmc/pymc_wrapper.py | 21 ++++++-- tests/mcmc_test.py | 50 ++++++++++++------- 5 files changed, 90 insertions(+), 53 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b37e00e15..e3eb4f17d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ ### 🐛 Bug Fixes +* **Improve PyMC HMC sampling** ([#1908](https://github.com/sbi-dev/sbi/pull/1908)): Use a target acceptance probability of `0.9` by default for `hmc_pymc`, avoiding poor finite-chain mixing observed with PyMC's `0.65` default. The value can be customized with `MCMCPosteriorParameters.target_accept` or `MCMCPosterior.sample(target_accept=...)`. * **Fix TARP z-scoring bug** ([#1832](https://github.com/sbi-dev/sbi/issues/1832)): Reference points are now z-scored alongside `thetas` and `posterior_samples` when `z_score_theta=True`, fixing incorrect distance calculations that masked bias detection. * **Fix broken `biased_toy_gaussian` test helper**: Rewrote to create actual location bias (posterior mean shifted from truth) instead of the previous NaN-producing formula. * **Raise on unsupported `transform_to_unconstrained` z-scoring**: `z_score_x="transform_to_unconstrained"` was silently ignored by the nflows builders (`maf`, `nsf`, `maf_rqs`, `made`), the MDN builder, the unconditional Zuko builder, the ratio-based classifier builders (`linear`, `mlp`, `resnet`), and the vector field builders (FMPE / NPSE), producing a model with no reparametrization. These builders now raise a clear `ValueError`. The mixed builders (MNLE / MNPE) raise transitively when their continuous flow is a non-Zuko model. The option remains supported by the conditional `zuko_*` models. diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py index db23fb25d..f9e8c0be0 100644 --- a/sbi/inference/posteriors/mcmc_posterior.py +++ b/sbi/inference/posteriors/mcmc_posterior.py @@ -31,6 +31,7 @@ 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 +from sbi.utils.typechecks import validate_float_range class MCMCPosterior(NeuralPosterior): @@ -108,9 +109,10 @@ def __init__( device: Training device, e.g., "cpu", "cuda" or "cuda:0". If None, `potential_fn.device` is used. x_shape: Deprecated, should not be passed. - target_accept: Target acceptance probability for the gradient-based - samplers (HMC/NUTS). See `MCMCPosteriorParameters` for details. If - `None`, HMC uses `0.9` and NUTS/Pyro keep their backend defaults. + target_accept: Target acceptance probability for PyMC's HMC and NUTS + samplers. See `MCMCPosteriorParameters` for details. If `None`, PyMC + HMC uses `0.9` and PyMC NUTS keeps its backend default. Ignored by + other samplers. """ if method == "slice": warn( @@ -140,6 +142,14 @@ def __init__( self.init_strategy_parameters = init_strategy_parameters or {} self.num_workers = num_workers self.mp_context = mp_context + if target_accept is not None: + validate_float_range( + target_accept, + "target_accept", + min_val=0.0, + max_val=1.0, + range_inclusive=False, + ) self.target_accept = target_accept self._posterior_sampler = None @@ -291,9 +301,9 @@ def sample( mp_context: Multiprocessing context (`fork` or `spawn`). If not provided, uses the value specified at initialization. show_progress_bars: Whether to show sampling progress monitor. - target_accept: Target acceptance probability for the gradient-based - samplers (HMC/NUTS). If not provided, uses the value specified at - initialization. + target_accept: Target acceptance probability for PyMC's HMC and NUTS + samplers. If not provided, uses the value specified at initialization. + Ignored by other samplers. Returns: Samples from posterior. @@ -311,6 +321,14 @@ def sample( num_workers = self.num_workers if num_workers is None else num_workers mp_context = self.mp_context if mp_context is None else mp_context target_accept = self.target_accept if target_accept is None else target_accept + if target_accept is not None: + validate_float_range( + target_accept, + "target_accept", + min_val=0.0, + max_val=1.0, + range_inclusive=False, + ) init_strategy_parameters = ( self.init_strategy_parameters if init_strategy_parameters is None @@ -352,7 +370,6 @@ def sample( num_chains=num_chains, show_progress_bars=show_progress_bars, mp_context=mp_context, - target_accept=target_accept, ) elif method in ("hmc_pymc", "nuts_pymc", "slice_pymc"): transformed_samples = self._pymc_mcmc( @@ -806,7 +823,6 @@ def _pyro_mcmc( num_chains: Optional[int] = 1, show_progress_bars: bool = True, mp_context: str = "spawn", - target_accept: Optional[float] = None, ) -> Tensor: r"""Return samples obtained using Pyro's HMC or NUTS sampler. @@ -822,8 +838,6 @@ def _pyro_mcmc( warmup_steps: Initial number of samples to discard. num_chains: Whether to sample in parallel. If None, use all but one CPU. show_progress_bars: Whether to show a progressbar during sampling. - target_accept: Target acceptance probability for step-size adaptation. - If None, Pyro's default is used. Returns: Tensor of shape (num_samples, shape_of_single_theta). @@ -842,12 +856,8 @@ def _pyro_mcmc( thin = _process_thin_default(thin) num_chains = mp.cpu_count() - 1 if num_chains is None else num_chains kernels = dict(hmc_pyro=HMC, nuts_pyro=NUTS) - kernel_kwargs = {"potential_fn": potential_function} - if target_accept is not None: - kernel_kwargs["target_accept_prob"] = target_accept - sampler = MCMC( - kernel=kernels[mcmc_method](**kernel_kwargs), + kernel=kernels[mcmc_method](potential_fn=potential_function), num_samples=ceil((thin * num_samples) / num_chains), warmup_steps=warmup_steps, initial_params={self.param_name: initial_params}, @@ -896,8 +906,8 @@ def _pymc_mcmc( warmup_steps: Initial number of samples to discard. num_chains: Whether to sample in parallel. If None, use all but one CPU. show_progress_bars: Whether to show a progressbar during sampling. - target_accept: Target acceptance probability for HMC/NUTS. If None, - HMC uses 0.9 and NUTS keeps PyMC's default. Ignored by the slice + target_accept: Target acceptance probability for HMC/NUTS. If `None`, + HMC uses `0.9` and NUTS keeps PyMC's default. Ignored by the slice sampler. Returns: @@ -918,9 +928,6 @@ def _pymc_mcmc( steps = dict(slice_pymc="slice", hmc_pymc="hmc", nuts_pymc="nuts") step = steps[mcmc_method] - if target_accept is None and step == "hmc": - target_accept = 0.9 - sampler = PyMCSampler( potential_fn=potential_function, step=step, diff --git a/sbi/inference/posteriors/posterior_parameters.py b/sbi/inference/posteriors/posterior_parameters.py index 336a721e1..bd3abe9d3 100644 --- a/sbi/inference/posteriors/posterior_parameters.py +++ b/sbi/inference/posteriors/posterior_parameters.py @@ -23,6 +23,7 @@ is_nonnegative_int, is_positive_float, is_positive_int, + validate_float_range, ) @@ -244,13 +245,12 @@ class MCMCPosteriorParameters(PosteriorParameters): (default), used by Pyro and PyMC samplers. `"fork"` can be significantly faster than `"spawn"` but is only supported on POSIX-based systems (e.g. Linux and macOS, not Windows). - target_accept: Target acceptance probability for the gradient-based - samplers (`hmc_pymc`, `nuts_pymc`, `hmc_pyro`, `nuts_pyro`), controlling - the step-size adaptation during warmup. Higher values take smaller, - more accurate steps at the cost of speed. Ignored by the slice samplers. - If `None`, HMC uses `0.9` (PyMC's default of `0.65` mixes poorly on - peaked targets and yields biased samples), while NUTS and the Pyro - samplers keep their respective backend defaults. + target_accept: Target acceptance probability for `hmc_pymc` and `nuts_pymc`, + controlling step-size adaptation during warmup. Higher values generally + result in smaller steps and fewer rejections, at the cost of speed. If + `None`, HMC uses `0.9` because PyMC's `0.65` default mixed poorly in sbi's + Gaussian regression test; NUTS keeps PyMC's backend default. Ignored by + other samplers. """ method: Literal[ @@ -274,8 +274,14 @@ class MCMCPosteriorParameters(PosteriorParameters): def validate(self): """Validate MCMCPosteriorParameters fields.""" - if self.target_accept is not None and not (0.0 < self.target_accept < 1.0): - raise ValueError("target_accept must be between 0 and 1, or None.") + if self.target_accept is not None: + validate_float_range( + self.target_accept, + "target_accept", + min_val=0.0, + max_val=1.0, + range_inclusive=False, + ) if not ( self.init_strategy_parameters is None diff --git a/sbi/samplers/mcmc/pymc_wrapper.py b/sbi/samplers/mcmc/pymc_wrapper.py index 2933e31f9..9d4e7ba71 100644 --- a/sbi/samplers/mcmc/pymc_wrapper.py +++ b/sbi/samplers/mcmc/pymc_wrapper.py @@ -9,6 +9,7 @@ import torch from sbi.utils.torchutils import tensor2numpy +from sbi.utils.typechecks import validate_float_range class PyMCPotential(pt.Op): # type: ignore @@ -133,10 +134,10 @@ def __init__( seed: Random seed passed to `pymc.sample` for reproducible sampling. If None (default), PyMC seeds from system entropy. target_accept: Target acceptance probability for the `"hmc"` and - `"nuts"` step methods. If None (default), PyMC's own default is - used. PyMC's HMC default (0.65) mixes poorly on peaked targets; - a higher value (e.g. 0.9-0.95) gives markedly less biased samples. - Ignored for the `"slice"` step. + `"nuts"` step methods. If `None`, HMC uses sbi's default of `0.9`, + while NUTS keeps PyMC's backend default. The HMC default avoids poor + finite-chain mixing observed with PyMC's `0.65` default in sbi's + Gaussian regression test. Ignored for the `"slice"` step. """ self.param_name = param_name self._step = step @@ -148,7 +149,17 @@ def __init__( self._progressbar = progressbar self._device = device self._seed = seed - self._target_accept = target_accept + if target_accept is not None: + validate_float_range( + target_accept, + "target_accept", + min_val=0.0, + max_val=1.0, + range_inclusive=False, + ) + self._target_accept = ( + 0.9 if target_accept is None and step == "hmc" else target_accept + ) # create PyMC model object track_gradients = step in ("nuts", "hmc") diff --git a/tests/mcmc_test.py b/tests/mcmc_test.py index 4afa44ddf..788030db0 100644 --- a/tests/mcmc_test.py +++ b/tests/mcmc_test.py @@ -159,9 +159,6 @@ def lp_f(x, track_gradients=True): tune=warmup, chains=num_chains, seed=0, # reproducible sampling; PyMC is not covered by the global seed - # PyMC's HMC default (target_accept=0.65) is biased on this peaked - # target (c2st ~0.63); 0.95 samples it accurately. See #1894 follow-up. - target_accept=0.95, ) samples = sampler.run() assert samples.shape == ( @@ -179,32 +176,47 @@ def lp_f(x, track_gradients=True): @pytest.mark.mcmc @pytest.mark.parametrize( - ("method", "target_accept", "expected"), + ("step", "target_accept", "expected"), [ - ("hmc_pymc", None, 0.9), # HMC default is raised (PyMC's 0.65 is biased) - ("hmc_pymc", 0.99, 0.99), # user override wins - ("nuts_pymc", None, None), # NUTS keeps PyMC's own default + ("hmc", None, 0.9), + ("hmc", 0.99, 0.99), + ("nuts", None, None), ], ) -def test_pymc_target_accept_forwarded(method, target_accept, expected): - """target_accept from the user reaches the PyMC sampler with the right default.""" +def test_pymc_target_accept_default_and_override(step, target_accept, expected): + """PyMC HMC uses sbi's default and explicit user values take precedence.""" + from sbi.samplers.mcmc.pymc_wrapper import PyMCSampler + theta_dim = 2 - prior = BoxUniform(low=-2 * ones(theta_dim), high=2 * ones(theta_dim)) def potential_fn(theta): return -0.5 * (theta**2).sum(axis=-1) - mcmc_posterior = build_from_potential(potential_fn, prior) - mcmc_posterior.sample( - (10,), - method=method, + sampler = PyMCSampler( + potential_fn=potential_fn, + initvals=np.zeros((1, theta_dim)), + step=step, target_accept=target_accept, - num_chains=1, - warmup_steps=10, - thin=1, - show_progress_bars=False, + draws=1, + tune=1, + chains=1, ) - assert mcmc_posterior._posterior_sampler._target_accept == expected + assert sampler._target_accept == expected + + +@pytest.mark.mcmc +def test_mcmc_posterior_rejects_invalid_target_accept(): + """Per-call overrides validate target_accept at the public API boundary.""" + theta_dim = 2 + prior = BoxUniform(low=-2 * ones(theta_dim), high=2 * ones(theta_dim)) + + def potential_fn(theta): + return -0.5 * (theta**2).sum(axis=-1) + + with pytest.warns(UserWarning, match="unconditional potential"): + posterior = build_from_potential(potential_fn, prior) + with pytest.raises(ValueError, match="target_accept"): + posterior.sample((1,), method="hmc_pymc", target_accept=0.0) @pytest.mark.mcmc