MCMCPosterior

MCMCPosterior#

class MCMCPosterior(potential_fn, proposal, theta_transform=None, method='slice_np_vectorized', thin=-1, warmup_steps=200, num_chains=20, init_strategy='resample', init_strategy_parameters=None, init_strategy_num_candidates=None, num_workers=1, mp_context='spawn', device=None, x_shape=None)[source]#

Bases: NeuralPosterior

Provides MCMC to sample from the posterior.

SNLE or SNRE train neural networks to approximate the likelihood(-ratios). MCMCPosterior allows to sample from the posterior with MCMC.

Parameters:

  • potential_fn (Callable | BasePotential)

  • proposal (Any)

  • theta_transform (Transform | None)

  • method (Literal['slice_np', 'slice_np_vectorized', 'hmc_pyro', 'nuts_pyro', 'slice_pymc', 'hmc_pymc', 'nuts_pymc'])

  • thin (int)

  • warmup_steps (int)

  • num_chains (int)

  • init_strategy (Literal['proposal', 'sir', 'resample'])

  • init_strategy_parameters (Dict[str, Any] | None)

  • init_strategy_num_candidates (int | None)

  • num_workers (int)

  • mp_context (Literal['fork', 'spawn'])

  • device (str | device | None)

  • x_shape (Size | None)

to(device)[source]#

Moves potential_fn, proposal, x_o and theta_transform to the

specified device. Reinstantiates the posterior and resets the default x_o.

Parameters:

device (str | device) – Device to move the posterior to.

Return type:

None

property mcmc_method: str#

Returns MCMC method.

property posterior_sampler#

Returns sampler created by sample.

set_mcmc_method(method)[source]#

Sets sampling method to for MCMC and returns NeuralPosterior.

Parameters:

method (str) – Method to use.

Returns:

NeuralPosterior for chainable calls.

Return type:

NeuralPosterior

log_prob(theta, x=None, track_gradients=False)[source]#

Returns the log-probability of theta under the posterior.

Parameters:

  • theta (Tensor) – Parameters \(\theta\).

  • track_gradients (bool) – Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption.

  • x (Tensor | None)

Returns:

len(\(\theta\))-shaped log-probability.

Return type:

Tensor

sample(sample_shape=(), x=None, method=None, thin=None, warmup_steps=None, num_chains=None, init_strategy=None, init_strategy_parameters=None, num_workers=None, mp_context=None, show_progress_bars=True)[source]#

Draw samples from the approximate posterior distribution \(p(\theta|x)\).

Parameters:

  • sample_shape (Size | Tuple[int, ...]) – Desired shape of samples that are drawn from posterior. If sample_shape is multidimensional we simply draw sample_shape.numel() samples and then reshape into the desired shape.

  • x (Tensor | None) – Conditioning observation \(x_o\). If not provided, uses the default x set via .set_default_x().

  • method (str | None) – MCMC method to use. One of slice_np, slice_np_vectorized, hmc_pyro, nuts_pyro, slice_pymc, hmc_pymc, nuts_pymc. If not provided, uses the method specified at initialization.

  • thin (int | None) – Thinning factor for the chain. If not provided, uses the value specified at initialization.

  • warmup_steps (int | None) – Number of warmup steps to discard. If not provided, uses the value specified at initialization.

  • num_chains (int | None) – Number of MCMC chains to run. If not provided, uses the value specified at initialization.

  • init_strategy (str | None) – Initialization strategy for chains (proposal, sir, or resample). If not provided, uses the value specified at initialization.

  • init_strategy_parameters (Dict[str, Any] | None) – Parameters for the initialization strategy. If not provided, uses the value specified at initialization.

  • num_workers (int | None) – Number of CPU cores for parallelization. If not provided, uses the value specified at initialization.

  • mp_context (str | None) – Multiprocessing context (fork or spawn). If not provided, uses the value specified at initialization.

  • show_progress_bars (bool) – Whether to show sampling progress monitor.

Returns:

Samples from posterior.

Return type:

Tensor

sample_batched(sample_shape, x, method=None, thin=None, warmup_steps=None, num_chains=None, init_strategy=None, init_strategy_parameters=None, num_workers=None, mp_context=None, show_progress_bars=True)[source]#

Draw samples from the posteriors for a batch of different xs.

Given a batch of observations [x_1, …, x_B], this method samples from posteriors \(p(\theta|x_1), \ldots, p(\theta|x_B)\) in a vectorized manner.

Check the __init__() method for a description of all arguments as well as their default values.

Parameters:

  • sample_shape (Size | Tuple[int, ...]) – Desired shape of samples that are drawn from the posterior given every observation.

  • x (Tensor) – A batch of observations, of shape (batch_dim, event_shape_x). batch_dim corresponds to the number of observations to be drawn.

  • method (str | None) – Method used for MCMC sampling, e.g., “slice_np_vectorized”.

  • thin (int | None) – The thinning factor for the chain, default 1 (no thinning).

  • warmup_steps (int | None) – The initial number of samples to discard.

  • num_chains (int | None) – The number of chains used for each x passed in the batch.

  • init_strategy (str | None) – The initialisation strategy for chains.

  • init_strategy_parameters (Dict[str, Any] | None) – Dictionary of keyword arguments passed to the init strategy.

  • num_workers (int | None) – number of cpu cores used to parallelize initial parameter generation and mcmc sampling.

  • mp_context (str | None) – Multiprocessing start method, either “fork” or “spawn”

  • show_progress_bars (bool) – Whether to show sampling progress monitor.

Returns:

Samples from the posteriors of shape (*sample_shape, B, *input_shape)

Return type:

Tensor

map(x=None, num_iter=1000, num_to_optimize=100, learning_rate=0.01, init_method='proposal', num_init_samples=1000, save_best_every=10, show_progress_bars=False, force_update=False)[source]#

Returns the maximum-a-posteriori estimate (MAP).

The method can be interrupted (Ctrl-C) when the user sees that the log-probability converges. The best estimate will be saved in self._map and can be accessed with self.map(). The MAP is obtained by running gradient ascent from a given number of starting positions (samples from the posterior with the highest log-probability). After the optimization is done, we select the parameter set that has the highest log-probability after the optimization.

Warning: The default values used by this function are not well-tested. They might require hand-tuning for the problem at hand.

For developers: if the prior is a BoxUniform, we carry out the optimization in unbounded space and transform the result back into bounded space.

Parameters:

  • x (Tensor | None) – Deprecated - use .set_default_x() prior to .map().

  • num_iter (int) – Number of optimization steps that the algorithm takes to find the MAP.

  • learning_rate (float) – Learning rate of the optimizer.

  • init_method (str | Tensor) – How to select the starting parameters for the optimization. If it is a string, it can be either [posterior, prior], which samples the respective distribution num_init_samples times. If it is a tensor, the tensor will be used as init locations.

  • num_init_samples (int) – Draw this number of samples from the posterior and evaluate the log-probability of all of them.

  • num_to_optimize (int) – From the drawn num_init_samples, use the num_to_optimize with highest log-probability as the initial points for the optimization.

  • save_best_every (int) – The best log-probability is computed, saved in the map-attribute, and printed every save_best_every-th iteration. Computing the best log-probability creates a significant overhead (thus, the default is 10.)

  • show_progress_bars (bool) – Whether to show a progressbar during sampling from the posterior.

  • force_update (bool) – Whether to re-calculate the MAP when x is unchanged and have a cached value.

  • log_prob_kwargs – Will be empty for SNLE and SNRE. Will contain {‘norm_posterior’: True} for SNPE.

Returns:

The MAP estimate.

Return type:

Tensor

property default_x: Tensor | None#

Return default x used by .sample(), .log_prob as conditioning context.

potential(theta, x=None, track_gradients=False)#

Evaluates \(\theta\) under the potential that is used to sample the posterior.

The potential is the unnormalized log-probability of \(\theta\) under the posterior.

Parameters:

  • theta (Tensor) – Parameters \(\theta\).

  • track_gradients (bool) – Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption.

  • x (Tensor | None)

Return type:

Tensor

set_default_x(x)#

Set new default x for .sample(), .log_prob to use as conditioning context.

Reset the MAP stored for the old default x if applicable.

This is a pure convenience to avoid having to repeatedly specify x in calls to .sample() and .log_prob() - only $ heta$ needs to be passed.

This convenience is particularly useful when the posterior is focused, i.e. has been trained over multiple rounds to be accurate in the vicinity of a particular x=x_o (you can check if your posterior object is focused by printing it).

NOTE: this method is chainable, i.e. will return the NeuralPosterior object so that calls like posterior.set_default_x(my_x).sample(mytheta) are possible.

Parameters:

x (Tensor) – The default observation to set for the posterior \(p( heta|x)\).

Returns:

NeuralPosterior that will use a default x when not explicitly passed.

Return type:

NeuralPosterior

Navigation