pespace.inference.sampler.flowmc

FlowMC sampler interface for Bayesian parameter estimation.

flowMC is a JAX-based normalizing-flow-enhanced MCMC sampler (Wong et al. 2023, Gabrié et al. 2022). A local MCMC kernel (Gaussian random walk / MALA / HMC) walks the chains while a rational-quadratic-spline normalizing flow is trained online on the accumulated chain history and used as an independent-Metropolis global proposal, which dramatically accelerates the mixing of multi-modal or high-dimensional posteriors.

This module provides FlowMCSampler, which follows the same design as the other samplers in this package: priors are declared through numpyro distributions (NumpyroPriorSamplerBase), sampling is performed in the unbounded (unconstrained) space prepared by the prior bijections, and the output is transformed back to physical parameters.

Six built-in flowMC resource-strategy bundles are available under their flowMC names (_buildin_bundles), with "RQSpline_MALA" as the default. Bundle hyperparameters are supplied as a plain bundle_config dictionary and merged on top of the built-in default configuration (_default_config). n_chains and checkpoint_interval are sampler-level settings managed directly by the constructor. Fully custom setups are supported by passing a pre-built ResourceStrategyBundle instance (bundle_config is then ignored). flowMC is an optional dependency and is imported lazily.

Notes

  • The likelihood only needs to provide log_likelihood(params) and the likelihood_param_type namedtuple type — exactly the interface WhittleLikelihoodJax exposes.

  • MCMC provides no evidence estimate; log_evidence is reported as NaN.

  • For the parallel-tempering bundles the logpdf passed to flowMC is the log-likelihood, and the log-prior (including the Jacobian of the prior bijections, which must not be rescaled by the temperature) is passed as the untempered logprior.

  • flowMC has not reached v1.0 yet and its API is unstable; this interface is developed and tested against flowMC 0.6.x.

Classes

FlowMCSampler(likelihood, prior_dict[, ...])

Normalizing-flow-enhanced MCMC (flowMC) sampler.
class pespace.inference.sampler.flowmc.FlowMCSampler(likelihood, prior_dict, label=None, outdir=None, rkey_seed=0, resource_strategy_bundles='RQSpline_MALA', bundle_config={'adapt_step_size': True, 'adapt_step_size_per_dim': True, 'batch_size': 10000, 'chain_batch_size': 0, 'early_stopping': False, 'early_stopping_min_acceptance': 0.1, 'early_stopping_patience': 3, 'early_stopping_tolerance': 0.05, 'global_thinning': 1, 'history_window': 100, 'learning_rate': 0.001, 'local_thinning': 1, 'mala_step_size': 0.1, 'n_NFproposal_batch_size': 10000, 'n_epochs': 100, 'n_global_steps': 100, 'n_local_steps': 100, 'n_max_examples': 10000, 'n_production_loops': 20, 'n_training_loops': 20, 'rq_spline_hidden_units': (32, 32), 'rq_spline_n_bins': 8, 'rq_spline_n_layers': 4, 'verbose': False}, n_chains=64, likelihood_batch_size=0, checkpoint_interval=600.0)[source]

Bases: NumpyroPriorSamplerBase

Normalizing-flow-enhanced MCMC (flowMC) sampler.

The chains evolve in the unbounded space prepared by NumpyroPriorSamplerBase (the prior bijections and the transform back to physical parameters are handled by the base class). Each training loop runs n_local_steps steps of the local kernel (GRW / MALA / HMC), trains the RQ-spline normalizing flow for n_epochs epochs on the accumulated chain history, and then proposes n_global_steps independent-Metropolis moves from the flow. After n_training_loops such loops (the warmup, during which the step size is adapted and the flow is fitted), the flow is frozen and n_production_loops loops of local + global steps are collected as the posterior sample. Flow proposals are corrected by their proposal density, so the frozen-flow production chain satisfies detailed balance.

The flowMC sampler (bundle and backend) is constructed in __init__; run() only executes — and times — the sample call.

Parameters:

  • likelihood (object) – Likelihood instance providing a pure log_likelihood(params) -> scalar method (JAX-transformable) and the likelihood_param_type namedtuple type of the likelihood parameters.

  • prior_dict (dict) – Prior declaration, see NumpyroPriorSamplerBase.

  • label (str, optional) – Label of the run, converted to str; also the JSON filename (outdir/{label}_results.json). Default is None (-> "None").

  • outdir (str, optional) – Output directory for the saved result and the flowMC checkpoint. It must not exist yet: the base class creates it and refuses to reuse an existing path, so that the results of an earlier run are never overwritten. Default is None, which resolves to an outdir_{label} directory under the current working directory.

  • rkey_seed (int, optional) – Seed for the JAX random number generator key. Default is 0.

  • resource_strategy_bundles (str or ResourceStrategyBundle, optional) – Either the name of a built-in bundle (_buildin_bundles) or a pre-constructed ResourceStrategyBundle instance. Default is "RQSpline_MALA". For pre-built instances bundle_config is ignored and the number of chains is inferred from the bundle’s buffers.

  • bundle_config (dict, optional) – Bundle hyperparameters, merged on top of the built-in default configuration _default_config. Defaults to _default_config["RQSpline_MALA"]. Unknown keys are discarded with a printed warning; settings that are not explicitly given fall back to their defaults (also with a printed warning). Ignored when a pre-built bundle instance is provided.

  • n_chains (int, optional) – Number of parallel chains. Default is 64. With a pre-built bundle instance, passing None infers the value from the bundle’s buffers; a concrete value is checked against the bundle’s buffers and a mismatch raises an error.

  • likelihood_batch_size (int, optional) – Evaluate the final log-likelihoods in sequential batches of this size (jax.lax.map); 0 keeps a single jax.vmap. A positive value bounds the peak memory of the likelihood intermediates, at the cost of cross-batch parallelism. Default is 0.

  • checkpoint_interval (float, optional) – Minimum wall-clock seconds between flowMC checkpoint writes; set 0 to disable checkpointing. Default is 600.0.

Raises:

  • ImportError – If flowMC is not installed.

  • ValueError –

    • If resource_strategy_bundles names an unknown bundle or is not a string / bundle instance; - If the number of chains cannot be inferred from a pre-built bundle instance.

_z_to_array(z)[source]

Stack a sampled-parameter namedtuple into an (..., n_dim) array.

Handles both scalar leaves (shape ()) and batched leaves (shape (n_chains,) or (n_samples,)).

Return type:

Array

_build_logpdfs(is_pt)[source]

Return the flowMC logpdf (and logprior for PT bundles).

The flowMC callables take f(x, data) with x a 1-D array; they rewrap x into the sampled-parameter namedtuple and evaluate the unbounded-space functions of the base class. PT bundles temper the likelihood, so logpdf is the log-likelihood there and the log-prior (including the bijection Jacobian, which must not be rescaled by the temperature) is returned as logprior.

Parameters:

is_pt (bool)

_resolve_bundle(bundle_arg, bundle_config)[source]

Resolve the resource_strategy_bundles / bundle_config inputs into a ready-to-use ResourceStrategyBundle.

Sets self.resource_strategy_bundles (the bundle class name), self.bundle_config (the merged configuration, or None for a pre-built instance) and self.n_chains.

Unknown bundle_config keys are discarded with a printed warning; settings that are not explicitly given fall back to their defaults (also with a printed warning).

Return type:

ResourceStrategyBundle

_build_sampler()[source]

Construct the flowMC Sampler from the resolved bundle.

Called from __init__ so that run only executes and times sample.

Return type:

Sampler

_get_production_samples()[source]

Collect the production samples from the bundle resources.

Built-in bundles store them in the "positions_production" buffer; custom bundles are expected to contain a State resource whose "target_positions" entry points to the output buffer (the same convention the built-in bundles use).

Return type:

Array

_finalize(flat_samples, runing_time)[source]

Assemble the chain output into the run-result dictionary.

The chain positions are restored to the physical space and reported as sampled_params; the fixed_params are broadcast to the sample count, and the derived_params are taken from the likelihood-input assembly. MCMC samples are equally weighted, so weights is left as None and log_evidence is NaN.

Return type:

dict

run()[source]

Run the flowMC sampler.

Only the sample call is executed (and timed) here; the flowMC sampler and its bundle were constructed in __init__.

Returns:

Run result containing label, outdir, rkey_seed, prior_dict, likelihood_class, runing_time, sampler_class, sampler_config, posteriors (physical sampled/derived/fixed parameters) and log_likelihoods; the log-evidence log_evidence is NaN and weights is an empty array (equally weighted MCMC samples). The result is also saved to outdir/{label}_results.json.

Return type:

dict

_abc_impl = <_abc._abc_data object>