jaxili.posterior.mcmc_posterior module

jaxili.posterior.mcmc_posterior module#

MCMC Posterior.

This module contains the MCMCPosterior class that wraps the NeuralPosterior class to perform MCMC sampling.

class jaxili.posterior.mcmc_posterior.MCMCPosterior(model: NDENetwork, state: TrainState, prior_distr: Distribution, verbose: bool | None = False, x: Array | None = None, mcmc_method: str | None = 'nuts_numpyro', mcmc_kwargs: dict | None = {})[source]#

Bases: NeuralPosterior

Likelihood $p(x|theta)$ with log_prob() and sample() methods.

The class wraps the trained neural network using Neural Likelihood Estimation (NLE). Sampling is performed using Markov Chain Monte Carlo (MCMC) methods to get samples from the posterior.

Methods

log_likelihood(x, theta)

Compute the log-likelihood learned by the neural density estimator.

log_prior(theta)

Compute the log prior of the parameters.

print_implemented_methods()

Print the implemented MCMC methods.

sample(num_samples, key[, x])

Sample from the posterior using MCMC.

set_default_x(x)

Set the default data for the posterior.

set_mcmc_kwargs(mcmc_kwargs)

Set the keyword arguments for the MCMC method.

set_mcmc_method(mcmc_method)

Set the MCMC method to use.

set_prior(prior_distr)

Set the prior distribution for the parameters.

unnormalized_log_prob(theta[, x])

Compute the unnormalized log probability of the posterior.
log_likelihood(x: Array, theta: Array)[source]#

Compute the log-likelihood learned by the neural density estimator.

Parameters:

  • theta (Array) – The parameters to evaluate the log probability.

  • x (Array) – The data used to condition the posterior.

Returns:

The unnormalized log probability.

Return type:

Array

log_prior(theta)[source]#

Compute the log prior of the parameters.

Parameters:

theta (Array) – The parameters to evaluate the log prior.

Returns:

The log prior of the parameters.

Return type:

Array

print_implemented_methods()[source]#

Print the implemented MCMC methods.

sample(num_samples: int, key: Array, x: Array | None = None, **kwargs)[source]#

Sample from the posterior using MCMC.

Parameters:

  • x (Array) – The data used to sample the parameters.

  • num_samples (int) – The number of samples to draw.

  • key (Array) – The random key used to generate the samples.

Returns:

The samples from the posterior.

Return type:

Array

set_default_x(x: Array)[source]#

Set the default data for the posterior.

set_mcmc_kwargs(mcmc_kwargs: dict)[source]#

Set the keyword arguments for the MCMC method.

set_mcmc_method(mcmc_method: str)[source]#

Set the MCMC method to use.

set_prior(prior_distr: Distribution)[source]#

Set the prior distribution for the parameters.

unnormalized_log_prob(theta: Array, x: Array | None = None)[source]#

Compute the unnormalized log probability of the posterior.

Parameters:

  • theta (Array) – The parameters to evaluate the log probability.

  • x (Array) – The data used to condition the posterior.

Returns:

The unnormalized log probability.

Return type:

Array

Contents