Source code for pypomp.functional.pmcmc

"""
Pure-functional PMCMC entry point.

This is a thin wrapper around :func:`pypomp.core.algorithms.pmcmc._vmapped_pmcmc_internal`
that takes a compiled :class:`pypomp.functional.structs.PompStruct` and runs
``n_chains`` independent PMCMC chains in parallel via ``jax.vmap``.
"""

from typing import Callable

import jax

from .structs import PompStruct
from ..core.algorithms.pmcmc import _vmapped_pmcmc_internal
from ..core.algorithms.types import PmcmcConfig, PmcmcInputs




def pmcmc(
    struct: PompStruct,
    thetas_array: jax.Array,
    proposal,
    dprior: Callable,
    M: int,
    J: int,
    thresh: float,
    keys: jax.Array,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    """
    Functional PMCMC entry point.

    Runs ``n_chains`` independent particle-MCMC chains in parallel.  Each chain
    starts at the corresponding row of ``thetas_array`` using the corresponding
    PRNG key in ``keys``.  Intended for users who need to compose PMCMC inside
    larger JAX programs; see :meth:`pypomp.core.pomp.Pomp.pmcmc` for a
    higher-level interface.

    Parameters
    ----------
    struct : PompStruct
        Compiled POMP model (see :class:`PompStruct`).
    thetas_array : jax.Array
        Starting parameter vectors, shape ``(n_chains, d)``.
    proposal
        Proposal object (see :mod:`pypomp.proposals`).  Shared across
        chains; per-chain state is initialised internally.
    dprior : Callable
        Log-prior density.  Pure JAX function with signature
        ``dprior(theta_arr) -> scalar``.
    M : int
        Number of MCMC iterations per chain.
    J : int
        Number of particles per filter evaluation.
    thresh : float
        Adaptive resampling threshold for the particle filter.
    keys : jax.Array
        PRNG keys, shape ``(n_chains, ...)``.

    Returns
    -------
    tuple[jax.Array, jax.Array, jax.Array, jax.Array]
        Tuple ``(loglik_traces, log_prior_traces, theta_traces, accepts)``:

        * ``loglik_traces``: shape ``(n_chains, M + 1)``.
        * ``log_prior_traces``: shape ``(n_chains, M + 1)``.
        * ``theta_traces``: shape ``(n_chains, M + 1, d)``.
        * ``accepts``: shape ``(n_chains,)`` -- count of accepted proposals per chain.
    """
    if struct.dmeas_pf is None:
        raise ValueError("PMCMC requires struct.dmeas_pf to be non-None.")
    config = PmcmcConfig.from_pmcmc_struct(
        struct=struct,
        M=M,
        J=J,
        dprior=dprior,
        thresh=thresh,
        should_trans=False,
    )
    inputs = PmcmcInputs.from_pmcmc_struct(struct)
    return _vmapped_pmcmc_internal(
        thetas_array,
        proposal,
        config,
        inputs,
        keys,
    )