VectorFieldPosterior#
- class VectorFieldPosterior(vector_field_estimator, prior, max_sampling_batch_size=10000, device=None, enable_transform=True, sample_with='sde', **kwargs)[source]#
Bases:
NeuralPosteriorPosterior based on flow- or score-matching estimators.
This posterior samples from the vector field model - typically a score-based or a flow matching model - given the vector_field_estimator and rejects samples that lie outside of the prior bounds.
The posterior is defined by a vector field estimator and a prior. The vector field estimator defines a continuous transformation from a base distribution to the approximated posterior distribution. Sampling is done by running either an ordinary differential equation (ODE) or a stochastic differential equation (SDE) defined by the vector field estimator with the starting points sampled from the base distribution.
Log probabilities are obtained by calling the potential function, which in turn uses the ODE to compute the log-probability.
- Parameters:
- sample(sample_shape=(), x=None, predictor='euler_maruyama', corrector=None, predictor_params=None, corrector_params=None, steps=500, ts=None, iid_method=None, iid_params=None, guidance_method=None, guidance_params=None, max_sampling_batch_size=10000, sample_with=None, show_progress_bars=True, reject_outside_prior=True, max_sampling_time=None, return_partial_on_timeout=False)[source]#
Return samples from posterior distribution \(p(\theta|x)\).
- Parameters:
sample_shape (Size | Tuple[int, ...]) – Shape of the samples to be drawn.
predictor (str | Predictor) – The predictor for the vector field sampler. Can be a string or a custom predictor following the API in sbi.samplers.score.predictors. Currently, only euler_maruyama is implemented.
corrector (str | Corrector | None) – The corrector for the vector field sampler. Either of [None].
predictor_params (Dict | None) – Additional parameters passed to predictor.
corrector_params (Dict | None) – Additional parameters passed to corrector.
steps (int) – Number of steps to take for the Euler-Maruyama method. If sample_with is “ode”, this is ignored.
ts (Tensor | None) – Time points at which to evaluate the vector field process. If None, a linear grid between t_max and t_min is used. If sample_with is “ode”, this is ignored.
iid_method (Literal['fnpe', 'gauss', 'auto_gauss', 'jac_gauss'] | None) – Which method to use for computing the score in the iid setting. We currently support “fnpe”, “gauss”, “auto_gauss”, “jac_gauss”. The fnpe method is simple and generally applicable. However, it can become inaccurate already for quite a few iid samples (as it based on heuristic approximations), and should be used at best only with a corrector. The “gauss” methods are more accurate, by aiming for an efficient approximation of the correct marginal score in the iid case. This however requires estimating some hyperparamters, which is done in a systematic way in the “auto_gauss” (initial overhead) and “jac_gauss” (iterative jacobian computations are expensive). We default to “auto_gauss” for these reasons. Note that in order to use the iid method, the vector field estimator must support it and have SCORE_DEFINED and MARGINALS_DEFINED class attributes set to True.
iid_params (Dict | None) – Additional parameters passed to the iid method. See the specific IIDScoreFunction child class for details.
guidance_method (str | None) – Method to guide the diffusion process. If None, no guidance is used. Currently we support affine_classifier_free, which allows to scale and shift the “likelihood” or “prior” score contribution. This can be used to perform “super” conditioning i.e. shrink the variance of the likelihood. Universal can be used to guide the diffusion process with a general guidance function. Interval is an instance of that where the guidance function constraints the diffusion process to a given interval.
guidance_params (Dict | None) – Additional parameters passed to the guidance method. See the specific ScoreAdaptation child class for details, specifically AffineClassifierFreeCfg, UniversalCfg, and IntervalCfg.
max_sampling_batch_size (int) – Maximum batch size for sampling.
sample_with (str | None) – Sampling method to use - ‘ode’ or ‘sde’. Note that in order to use the ‘sde’ sampling method, the vector field estimator must support it and have the SCORE_DEFINED class attribute set to True.
show_progress_bars (bool) – Whether to show a progress bar during sampling.
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 ODE/SDE sampler 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 does not use rejection).
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).
x (Tensor | None)
- Return type:
- sample_via_ode(sample_shape=(), **kwargs)[source]#
Return samples from posterior distribution with probability flow ODE.
This builds the probability flow ODE and then samples from the corresponding flow.
- log_prob(theta, x=None, track_gradients=False, ode_kwargs=None)[source]#
Returns the log-probability of the posterior \(p(\theta|x)\).
This requires building and evaluating the probability flow ODE.
- Parameters:
theta (Tensor) – Parameters \(\theta\).
x (Tensor | None) – Observed data \(x_o\). If None, the default \(x_o\) is used.
track_gradients (bool) – Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption.
ode_kwargs (Dict | None) – Additional keyword arguments for the ODE solver.
- Returns:
(len(θ),)-shaped log posterior probability \(\log p(\theta|x)\) for θ in the support of the prior, -∞ (corresponding to 0 probability) outside.
- Return type:
- sample_batched(sample_shape, x, predictor='euler_maruyama', corrector=None, predictor_params=None, corrector_params=None, steps=500, ts=None, max_sampling_batch_size=10000, show_progress_bars=True, reject_outside_prior=True, max_sampling_time=None, return_partial_on_timeout=False)[source]#
Given a batch of observations [x_1, …, x_B] this function samples from posteriors \(p(\theta|x_1)\), … ,\(p(\theta|x_B)\), in a batched (i.e. vectorized) manner.
- Parameters:
sample_shape (Size) – 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.
predictor (str | Predictor) – The predictor for the diffusion-based sampler. Can be a string or a custom predictor following the API in sbi.samplers.score.predictors. Currently, only euler_maruyama is implemented.
corrector (str | Corrector | None) – The corrector for the diffusion-based sampler.
predictor_params (Dict | None) – Additional parameters passed to predictor.
corrector_params (Dict | None) – Additional parameters passed to corrector.
steps (int) – Number of steps to take for the Euler-Maruyama method.
ts (Tensor | None) – Time points at which to evaluate the diffusion process. If None, a linear grid between t_max and t_min is used.
max_sampling_batch_size (int) – Maximum batch size for sampling.
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 ODE/SDE sampler 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.
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.
- Returns:
Samples from the posteriors of shape (*sample_shape, B, *input_shape)
- Return type:
- map(x=None, num_iter=1000, num_to_optimize=1000, learning_rate=0.01, init_method='posterior', num_init_samples=1000, save_best_every=1000, 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.
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.
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.
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.
- Returns:
The MAP estimate.
- Return type:
- 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.
- 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