RejectionPosterior

class RejectionPosterior(potential_fn, proposal, theta_transform=None, max_sampling_batch_size=10000, num_samples_to_find_max=10000, num_iter_to_find_max=100, m=1.2, device=None, x_shape=None)

Bases: NeuralPosterior

Provides rejection sampling to sample from the posterior.

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

Parameters:

• potential_fn (BasePotential | CustomPotential)

• proposal (Any)

• theta_transform (Transform | None)

• max_sampling_batch_size (int)

• num_samples_to_find_max (int)

• num_iter_to_find_max (int)

• m (float)

• device (str | device | None)

• x_shape (Size | None)

to(device)

Move potential fucntion, proposal and x_o to the device.

This method reinstantiates the posterior and resets the default x_o

Parameters:

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

Return type:

None

log_prob(theta, x=None, track_gradients=False)

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, max_sampling_batch_size=None, num_samples_to_find_max=None, num_iter_to_find_max=None, m=None, show_progress_bars=True, reject_outside_prior=True, max_sampling_time=None, return_partial_on_timeout=False)

Draw samples from the approximate posterior via rejection sampling.

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().

• max_sampling_batch_size (int | None) – Maximum batch size for rejection sampling. If not provided, uses the value specified at initialization.

• num_samples_to_find_max (int | None) – Number of samples to find the maximum of the potential function. If not provided, uses the value from initialization.

• num_iter_to_find_max (int | None) – Number of optimization iterations to find the maximum. If not provided, uses the value from initialization.

• m (float | None) – Multiplier for the proposal distribution. If not provided, uses the value from initialization.

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

• reject_outside_prior (bool) – If True (default), rejection sampling is used to ensure samples lie within the prior support. If False, samples are drawn directly from the proposal without rejection, which is faster but may include samples outside the prior support.

• max_sampling_time (float | None) – Optional maximum allowed sampling time in seconds. If exceeded, sampling is aborted and a RuntimeError is raised. Only applies when reject_outside_prior=True (no effect otherwise since direct sampling from the proposal is fast).

• return_partial_on_timeout (bool) – If True and max_sampling_time is exceeded, return the samples collected so far instead of raising a RuntimeError. A warning will be issued. Only applies when reject_outside_prior=True (default).

Returns:

Samples from posterior.

sample_batched(sample_shape, x, max_sampling_batch_size=10000, show_progress_bars=True)

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.

Parameters:

• sample_shape (Size | Tuple[int, ...]) – Shape of samples to draw for each observation.

• x (Tensor) – Batch of observations with shape (batch_dim, *event_shape_x).

• show_progress_bars (bool) – Whether to show a progress bar during sampling.

• **kwargs – Additional keyword arguments passed to the specific posterior’s sampling method.

• max_sampling_batch_size (int)

Returns:

Samples with shape (*sample_shape, batch_dim, *theta_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)

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