diff --git a/tests/numpyro/test_smc_logfns_space.py b/tests/numpyro/test_smc_logfns_space.py new file mode 100644 index 00000000..95aad518 --- /dev/null +++ b/tests/numpyro/test_smc_logfns_space.py @@ -0,0 +1,177 @@ +# Copyright 2026- The Blackjax Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests pinning build_smc_logfns / build_prior_sample_fn to a single, +consistent unconstrained-space convention. + +SMC consumes ``logprior_fn``/``loglikelihood_fn`` and particle positions in +the same **unconstrained** parameterization as ``build_logdensity_fn``'s +``init_position`` and ``postprocess_fn`` (unconstrained -> constrained). A +model with only real-supported latents can't catch a space mix-up (the two +spaces coincide), so these tests use a tiny model with one positive-support +latent, where constrained and unconstrained genuinely differ, and check +against an independently-derived analytic reference built from NumPyro's own +bijector rather than from the functions under test. +""" + +import jax +import jax.numpy as jnp +import numpyro +import numpyro.distributions as dist +import pytest +from numpyro.distributions import constraints +from numpyro.distributions.transforms import biject_to +from numpyro.infer.util import initialize_model + +from tuningfork.model._base import Posterior +from tuningfork.model._numpyro import build_prior_sample_fn, build_smc_logfns + +_Y_DATA = jnp.array([0.3, -0.2, 0.5, -0.1, 0.4]) + + +def _toy_halfnormal_model(y: jnp.ndarray) -> None: + """sigma ~ HalfNormal(2.0); y ~ Normal(0, sigma). + + One positive-support latent: constrained space (sigma > 0) and + unconstrained space (log sigma, any real) genuinely differ, so a + constrained/unconstrained mix-up is not invisible here. + """ + sigma = numpyro.sample("sigma", dist.HalfNormal(2.0)) + numpyro.sample("y", dist.Normal(0.0, sigma), obs=y) + + +_TOY_ENTRY = Posterior( + name="test_toy_halfnormal", + dim=1, + class_="test", + numpyro_model=_toy_halfnormal_model, + model_args=(_Y_DATA,), +) + +_SIGMA_TRANSFORM = biject_to(constraints.positive) + +# Dated, not searched (per project convention: seeds are stamped, never tuned). +_SEED = 20261002 +_Z_GRID = (-2.0, -1.0, -0.3, 0.0, 0.5, 1.0, 2.0) + + +def _analytic_logprior(z: float) -> float: + """log p(sigma) + log|det J| at unconstrained z, via NumPyro's own bijector.""" + sigma = _SIGMA_TRANSFORM(jnp.asarray(z)) + log_abs_det_jac = _SIGMA_TRANSFORM.log_abs_det_jacobian(jnp.asarray(z), sigma) + return float(dist.HalfNormal(2.0).log_prob(sigma) + log_abs_det_jac) + + +def _analytic_loglik(z: float) -> float: + """log p(y | sigma = constrain(z)) at unconstrained z.""" + sigma = _SIGMA_TRANSFORM(jnp.asarray(z)) + return float(dist.Normal(0.0, sigma).log_prob(_Y_DATA).sum()) + + +class TestBuildSmcLogfnsUnconstrainedSpace: + """logprior_fn / loglikelihood_fn must both operate in unconstrained space.""" + + @pytest.mark.fast + def test_logprior_fn_matches_analytic_prior(self) -> None: + key = jax.random.key(_SEED) + _, logprior_fn, _, _ = build_smc_logfns(key, _TOY_ENTRY) + for z in _Z_GRID: + got = float(logprior_fn({"sigma": jnp.asarray(z)})) + want = _analytic_logprior(z) + assert got == pytest.approx( + want, abs=1e-5 + ), f"logprior_fn({z}) = {got}, expected {want}" + + @pytest.mark.fast + def test_loglikelihood_fn_matches_analytic_likelihood(self) -> None: + key = jax.random.key(_SEED) + _, _, loglik_fn, _ = build_smc_logfns(key, _TOY_ENTRY) + for z in _Z_GRID: + got = float(loglik_fn({"sigma": jnp.asarray(z)})) + want = _analytic_loglik(z) + assert got == pytest.approx( + want, abs=1e-5 + ), f"loglikelihood_fn({z}) = {got}, expected {want}" + + @pytest.mark.fast + def test_logprior_plus_loglikelihood_equals_negative_potential(self) -> None: + """logprior_fn(z) + loglikelihood_fn(z) == -potential_fn(z) for every z. + + True by construction once each half is individually correct (the + Jacobian terms must cancel exactly) -- kept as an explicit assertion + since it's the invariant SMC's own tempered-target interpolation + relies on. + """ + key = jax.random.key(_SEED) + _, logprior_fn, loglik_fn, _ = build_smc_logfns(key, _TOY_ENTRY) + model_info = initialize_model( + key, + _TOY_ENTRY.numpyro_model, + model_args=_TOY_ENTRY.model_args, + model_kwargs=_TOY_ENTRY.model_kwargs, + dynamic_args=False, + ) + potential_fn = model_info.potential_fn + for z in _Z_GRID: + position = {"sigma": jnp.asarray(z)} + total = float(logprior_fn(position)) + float(loglik_fn(position)) + expected = float(-potential_fn(position)) + assert total == pytest.approx( + expected, abs=1e-5 + ), f"logprior+loglik={total}, -potential_fn={expected} at z={z}" + + @pytest.mark.fast + def test_blocked_model_latent_sites_match_joint_model(self) -> None: + """The prior-only (blocked) model's init_position has exactly the + same latent-site key set as the joint model's -- blocking only ever + hides is_observed sites, never adds or removes a latent one.""" + key = jax.random.key(_SEED) + init_position, _, _, _ = build_smc_logfns(key, _TOY_ENTRY) + assert set(init_position.keys()) == {"sigma"} + + +class TestBuildPriorSampleFnUnconstrainedSpace: + """build_prior_sample_fn's Predictive fallback must also return + unconstrained particles, matching postprocess_fn's own convention.""" + + @pytest.mark.fast + def test_predictive_fallback_particles_include_negative_values(self) -> None: + """sigma's constrained support is strictly positive; any negative + particle value proves these are not already-constrained draws.""" + prior_sample_fn = build_prior_sample_fn(_TOY_ENTRY) + particles = prior_sample_fn(jax.random.key(_SEED + 1), 2000) + assert jnp.any(particles["sigma"] < 0.0), ( + "particles look constrained (all non-negative); " + "build_prior_sample_fn's fallback must return unconstrained draws" + ) + + @pytest.mark.fast + def test_predictive_fallback_particles_recover_prior_mean_after_postprocess( + self, + ) -> None: + """Pushing particles through postprocess_fn (unconstrained -> + constrained) must recover the model's actual HalfNormal(2.0) prior + on sigma. Already-constrained particles would be double-transformed + (constrain an already-constrained value) and badly inflate the mean. + """ + _, _, _, postprocess_fn = build_smc_logfns(jax.random.key(_SEED), _TOY_ENTRY) + prior_sample_fn = build_prior_sample_fn(_TOY_ENTRY) + particles = prior_sample_fn(jax.random.key(_SEED + 1), 20_000) + + constrained = postprocess_fn(particles) + mean_sigma = float(jnp.mean(constrained["sigma"])) + analytic_mean = 2.0 * (2.0 / jnp.pi) ** 0.5 # E[HalfNormal(scale=2)] + assert mean_sigma == pytest.approx(analytic_mean, abs=0.05), ( + f"postprocess_fn(particles)['sigma'] mean = {mean_sigma}, " + f"expected ~{analytic_mean} (HalfNormal(2.0)'s own mean)" + ) diff --git a/tuningfork/model/_numpyro.py b/tuningfork/model/_numpyro.py index 74a3cd41..a6910793 100644 --- a/tuningfork/model/_numpyro.py +++ b/tuningfork/model/_numpyro.py @@ -36,7 +36,9 @@ from numpyro.handlers import block as _numpyro_block from numpyro.handlers import seed as _numpyro_seed from numpyro.handlers import trace as _numpyro_trace -from numpyro.infer.util import initialize_model, log_density +from numpyro.infer.util import get_transforms as _numpyro_get_transforms +from numpyro.infer.util import initialize_model +from numpyro.infer.util import transform_fn as _numpyro_transform_fn from tuningfork.model._base import Posterior @@ -134,18 +136,31 @@ def build_smc_logfns( the joint ``logdensity_fn`` into its prior and likelihood components by blocking the observed sites via NumPyro's ``block`` handler: - - ``logprior_fn(position)`` = log p(params) — runs the model with - observed sites blocked so only prior contributions are accumulated. - - ``loglikelihood_fn(position)`` = log p(data | params) — computed as - ``logposterior − logprior`` (numerically exact, no separate model call). - - Both functions are JAX-traceable and JIT-compatible. + - ``logprior_fn(position)`` = log p(constrain(position)) + log|det + J(position)| — the *negated potential energy of the blocked (prior-only) + model*, built the same way ``initialize_model`` builds the joint + ``potential_fn``, so both share the identical unconstrained-space + convention and Jacobian term. + - ``loglikelihood_fn(position)`` = log p(data | constrain(position)) — + computed as ``logposterior − logprior``; because both terms carry the + *same* ``log|det J|`` contribution, it cancels out exactly, leaving a + pure data-likelihood term with no Jacobian of its own (numerically + exact, no separate model call needed). + + ``position`` is always the **unconstrained** parameterization used by + ``initialize_model``/``find_valid_initial_params`` (the same space as + ``build_logdensity_fn``'s ``init_position`` and this function's own + ``init_position``) — the same space ``postprocess_fn`` maps *from*. + Both ``logprior_fn`` and ``loglikelihood_fn`` are JAX-traceable and + JIT-compatible. Parameters ---------- rng_key - JAX random key for ``initialize_model`` (used to draw the initial - unconstrained position). + JAX random key for ``initialize_model``, used verbatim for the joint + model (so ``init_position`` is unchanged from any prior version of + this function); the prior-only (blocked) model's own init search + uses a key folded in from ``rng_key``. entry Posterior registry entry describing the NumPyro model. @@ -156,16 +171,19 @@ def build_smc_logfns( logprior_fn ``(position: dict) -> float`` — log prior in unconstrained space. loglikelihood_fn - ``(position: dict) -> float`` — log likelihood (joint minus prior). + ``(position: dict) -> float`` — log likelihood (joint minus prior), + in the same unconstrained space. postprocess_fn Transforms unconstrained draws to constrained space. Notes ----- The ``block`` handler is applied once at function-build time to identify - observed sites; the resulting blocked model is then used to build a - ``log_density`` callable that is evaluated lazily at JAX-trace time. - This means no Python overhead per SMC step. + observed sites; the resulting blocked (prior-only) model is then run + through its own ``initialize_model`` call to obtain a genuine + unconstrained-space potential function — matching the joint model's own + ``inv_transforms`` for every latent site, since blocking only hides + observed sites and never touches which sites are latent. """ model_info = initialize_model( rng_key, @@ -181,21 +199,29 @@ def build_smc_logfns( obs_sites = _get_observed_site_names(entry) # Blocked model: same as the original model but with observed sites - # silenced → log_density gives prior only. + # silenced -- its own initialize_model gives a prior-only potential_fn + # in the identical unconstrained space/Jacobian convention as the joint + # model's potential_fn above, so the two compose correctly below. Uses a + # key folded in from rng_key (not a split of it) so init_position above + # is bit-identical to what build_smc_logfns returned before this fix -- + # the blocked model's own key only affects its own init search, not the + # resulting potential_fn. _blocked_model = _numpyro_block(entry.numpyro_model, hide=obs_sites) + blocked_model_info = initialize_model( + jax.random.fold_in(rng_key, 1), + _blocked_model, + model_args=entry.model_args, + model_kwargs=entry.model_kwargs, + dynamic_args=False, + ) + blocked_potential_fn = blocked_model_info.potential_fn def logprior_fn(position: dict) -> float: - """Log prior p(params) in unconstrained space.""" - logp, _ = log_density( - _blocked_model, - entry.model_args, - entry.model_kwargs, - position, - ) - return logp + """Log prior p(constrain(position)) + log|det J(position)|.""" + return -blocked_potential_fn(position) def loglikelihood_fn(position: dict) -> float: - """Log likelihood log p(data | params) = log p(data, params) − log p(params).""" + """Log likelihood log p(data | constrain(position)); Jacobian cancels.""" return -potential_fn(position) - logprior_fn(position) return init_position, logprior_fn, loglikelihood_fn, model_info.postprocess_fn @@ -210,10 +236,13 @@ def build_prior_sample_fn( - **Fast path** (``entry.analytic_sampler is not None``): calls ``entry.analytic_sampler(key, n)`` directly. Available for mvn_10, - ill_cond_50, banana, neals_funnel, gmm_25. + ill_cond_50, banana, neals_funnel, gmm_25. Per ``Posterior``'s own + contract, ``analytic_sampler`` already draws in unconstrained space. - **Fallback** (all other models): uses ``numpyro.infer.Predictive`` to - draw samples from the prior predictive and then returns only the - *latent* (non-observed) sites in unconstrained space. + draw samples from the prior predictive (in the model's *constrained* + space), keeps only the *latent* (non-observed) sites, then applies + each site's inverse constrain-to-unconstrain bijector so the returned + particles match the analytic path's unconstrained convention. Parameters ---------- @@ -224,19 +253,21 @@ def build_prior_sample_fn( ------- prior_sample_fn Callable ``(rng_key: jax.Array, n_particles: int) -> dict[str, Array]`` - where each array has shape ``(n_particles, *site_shape)``. + returning **unconstrained**-space particles, where each array has + shape ``(n_particles, *site_shape)`` — the same convention + ``build_smc_logfns``'s ``logprior_fn``/``loglikelihood_fn`` and + ``postprocess_fn`` (unconstrained → constrained) expect. Notes ----- - For the ``Predictive`` fallback, ``Predictive`` returns samples in the - model's *constrained* space. NumPyro automatically transforms these to - unconstrained space when constrained-to-unconstrained bijectors are - registered for the site's distribution. For most distributions used in - the benchmark suite (Normal, Bernoulli, etc.), the constrained and - unconstrained spaces coincide or the transformation is handled by - NumPyro's internal constrain/unconstrain utilities. Models that require - a non-trivial bijector (e.g. positive-constrained parameters) are handled - transparently by ``Predictive``. + ``Predictive`` itself returns samples in the model's constrained space + (e.g. a HalfNormal-prior site comes back strictly positive); the inverse + bijector (``biject_to(support).inv``, the same transform + ``initialize_model`` registers for that site) is applied explicitly + below — it is *not* automatic. The bijector set is derived once from a + single seeded trace (any one valid draw fixes every site's declared + support) and then applied batched over all ``n_particles`` at once, so + no second per-particle model trace is needed. Generated SMC programs call this function's output as ``prior_sample_fn(key, n_particles)`` when initializing particles. @@ -257,12 +288,24 @@ def _analytic_prior_sample_fn( _obs_sites = _get_observed_site_names(entry) from numpyro.infer import Predictive as _Predictive # noqa: PLC0415 + # Inverse (constrained -> unconstrained) bijectors, derived once from a + # single seeded trace. biject_to(support) depends only on each site's + # declared support, not on the realized value, so this is safe to reuse + # unbatched against the later batch of n_particles draws. + _seeded_model = _numpyro_seed(entry.numpyro_model, rng_seed=0) + _transforms = _numpyro_get_transforms( + _seeded_model, entry.model_args, entry.model_kwargs, {} + ) + def _predictive_prior_sample_fn( rng_key: jax.Array, n_particles: int ) -> dict[str, jax.Array]: pred = _Predictive(entry.numpyro_model, num_samples=n_particles) samples = pred(rng_key, *entry.model_args, **entry.model_kwargs) - # Return only latent sites (exclude observed data arrays). - return {k: v for k, v in samples.items() if k not in _obs_sites} + # Keep only latent sites (exclude observed data arrays), then map + # from the model's constrained space to the unconstrained space + # build_smc_logfns/postprocess_fn expect. + latent = {k: v for k, v in samples.items() if k not in _obs_sites} + return _numpyro_transform_fn(_transforms, latent, invert=True) return _predictive_prior_sample_fn