jesterTOV.inference.samplers.flowmc.FlowMCSampler

jesterTOV.inference.samplers.flowmc.FlowMCSampler#

class FlowMCSampler(
likelihood,
prior,
config,
sample_transforms=None,
likelihood_transforms=None,
seed=0,
local_sampler_name='GaussianRandomWalk',
local_sampler_arg=None,
num_layers=10,
hidden_size=None,
num_bins=8,
)[source]#

Bases: JesterSampler

FlowMC-specific sampler implementation.

This class inherits from JesterSampler and adds flowMC-specific initialization and configuration. It creates a flowMC Sampler with: - Local sampler (MALA or GaussianRandomWalk) - Normalizing flow model (MaskedCouplingRQSpline) - Training and production sampling loops

Parameters:

  • likelihood (LikelihoodBase) – Likelihood object with evaluate(params, data) method

  • prior (Prior) – Prior object with sample() and log_prob() methods

  • config (FlowMCSamplerConfig) – Configuration object from YAML (contains n_chains, n_loop_training, learning_rate, etc.)

  • sample_transforms (list[BijectiveTransform], optional) – Bijective transforms applied during sampling (with Jacobians)

  • likelihood_transforms (list[NtoMTransform], optional) – N-to-M transforms applied before likelihood evaluation

  • seed (int, optional) – Random seed (default: 0)

  • local_sampler_name (str, optional) – Name of the local sampler: “MALA” or “GaussianRandomWalk” (default: “GaussianRandomWalk”)

  • local_sampler_arg (dict[str, Any], optional) – Arguments for local sampler (e.g., {“step_size”: …})

  • num_layers (int, optional) – Number of coupling layers in normalizing flow (default: 10)

  • hidden_size (list[int], optional) – Hidden layer sizes for normalizing flow (default: [128, 128])

  • num_bins (int, optional) – Number of bins for rational quadratic splines (default: 8)

Variables:

  • sampler (Sampler) – FlowMC sampler instance

  • config (FlowMCSamplerConfig) – Configuration object

__init__(
likelihood,
prior,
config,
sample_transforms=None,
likelihood_transforms=None,
seed=0,
local_sampler_name='GaussianRandomWalk',
local_sampler_arg=None,
num_layers=10,
hidden_size=None,
num_bins=8,
)[source]#

Methods

__init__(likelihood, prior, config[, ...])

add_name(x)

Turn an array into a dictionary.

get_log_prob()

Get log probabilities from flowMC sampler (production samples only).

get_n_samples()

Get number of production samples from flowMC sampler.

get_n_training_samples()

Get number of training samples from flowMC sampler.

get_sampler_output()

Get standardized sampler output (production samples only).

get_samples()

Get production samples from the flowMC sampler.

get_training_sampler_output()

Get standardized sampler output for training samples.

posterior(params, data)

Evaluate posterior log probability from flat array.

posterior_from_dict(named_params, data)

Evaluate posterior log probability from parameter dict.

print_summary([transform])

Generate summary of the flowMC run.

sample(key)

Run flowMC sampling.

Attributes

sampler

likelihood

prior

sample_transforms

likelihood_transforms

parameter_names

get_log_prob()[source]#

Get log probabilities from flowMC sampler (production samples only).

Return type:

Array

Returns:

Array – Log posterior probability values (1D array, flattened across chains)

get_n_samples()[source]#

Get number of production samples from flowMC sampler.

Return type:

int

Returns:

int – Number of production samples (total across all chains)

get_n_training_samples()[source]#

Get number of training samples from flowMC sampler.

This is a FlowMC-specific method for diagnostic purposes.

Return type:

int

Returns:

int – Number of training samples (total across all chains)

get_sampler_output()[source]#

Get standardized sampler output (production samples only).

Return type:

SamplerOutput

Returns:

SamplerOutput –

  • samples: Parameter samples (dict of arrays)

  • log_prob: Log posterior probability

  • metadata: {} (empty, MCMC has equal weights)

get_samples()[source]#

Get production samples from the flowMC sampler.

Return type:

dict

Returns:

dict – Dictionary of production samples

get_training_sampler_output()[source]#

Get standardized sampler output for training samples.

This is a FlowMC-specific method for diagnostic purposes.

Return type:

SamplerOutput

Returns:

SamplerOutput –

  • samples: Parameter samples from training phase (dict of arrays)

  • log_prob: Log posterior probability from training phase

  • metadata: {} (empty, MCMC has equal weights)

print_summary(transform=True)[source]#

Generate summary of the flowMC run.

Parameters:

transform (bool, optional) – Whether to apply inverse sample transforms to results (default: True)

sample(key)[source]#

Run flowMC sampling.

Parameters:

key (PRNGKeyArray) – JAX random key

Notes

This method includes a critical bug fix: parameter ordering is preserved when converting from dictionary to array using a list comprehension instead of jax.tree.leaves().

sampler: flowMC.Sampler.Sampler#
Contents