pypomp.functional.pmcmc

pypomp.functional.pmcmc(struct: PompStruct, thetas_array: Array, proposal, dprior: Callable, M: int, J: int, thresh: float, keys: Array) tuple[Array, Array, Array, Array][source]

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 pypomp.core.pomp.Pomp.pmcmc() for a higher-level interface.

Parameters:
  • struct (PompStruct) – Compiled POMP model (see PompStruct).

  • thetas_array (jax.Array) – Starting parameter vectors, shape (n_chains, d).

  • proposal – Proposal object (see 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 (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.

Return type:

tuple[jax.Array, jax.Array, jax.Array, jax.Array]