pypomp.functional.mif

pypomp.functional.mif(struct: PompStruct, thetas_array: Array, sigmas_array: Array, sigmas_init_array: Array, M: int, cooling_fn: Callable | float, J: int, keys: Array, thresh: float = 0.0, n_monitors: int = 0) → tuple[Array, Array, Array][source]

This is a pure functional implementation of the Iterated Filtering algorithm, intended for users who need to compose it within custom JAX loops or higher-order functions. For a more user-friendly (but impurely-functional) interface, see pypomp.core.pomp.Pomp.mif().

This implementation leverages JAX to efficiently vectorize the algorithm across multiple initial parameter sets simultaneously.

Parameters:

  • struct (PompStruct) – The compiled structural representation of the POMP model.

  • thetas_array (jax.Array) – Array of initial parameters. Shape (n_reps, J, n_params) on the natural scale. Must be aligned with the canonical order of struct.param_names (e.g. prepared via align_params).

  • sigmas_array (jax.Array) – Array of random walk sigmas. Shape (n_params,). Must be aligned with the canonical order of struct.param_names.

  • sigmas_init_array (jax.Array) – Array of initial random walk sigmas. Shape (n_params,). Must be aligned with the canonical order of struct.param_names.

  • M (int) – Number of iterations.

  • cooling_fn (Callable | float) – Cooling function taking (nt, m, ntimes) or float cooling factor.

  • J (int) – Number of particles.

  • keys (jax.Array) – Random keys. Shape (n_reps, …).

  • thresh (float) – Resampling threshold.

  • n_monitors (int) – Number of monitors for likelihood averaging.

Returns:

Negative log-likelihood history: Shape (n_reps, M). Parameter trace history: Shape (n_reps, M+1, n_params) on the natural scale. Final particle swarm: Shape (n_reps, J, n_params) on the natural scale.

Return type:

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

Note

To align and stack input parameter dictionaries/scalars into the correct canonical ordering required by these arrays, use pypomp.functional.align_params().