nifty.re.evi module

class Samples(*, pos: P = None, samples: P, keys=None)

Bases: object

Storage class for samples (relative to some expansion point) that is fully compatible with JAX transformations like vmap, pmap, etc.

This class is used to store samples for the Variational Inference schemes MGVI and geoVI where samples are defined relative to some expansion point (a.k.a. latent mean or offset).

See also

Geometric Variational Inference, Philipp Frank, Reimar Leike, Torsten A. Enßlin, https://arxiv.org/abs/2105.10470 https://doi.org/10.3390/e23070853

Metric Gaussian Variational Inference, Jakob Knollmüller, Torsten A. Enßlin, https://arxiv.org/abs/1901.11033

__init__(*, pos: P = None, samples: P, keys=None)

at(pos, old_pos=None)

Update the offset (usually the latent mean) of all samples and optionally subtracts old_pos from all samples before.

property keys

property pos

property samples

squeeze()

Convenience method to merge the two leading axis of stacked samples (e.g. from batching).

tree_flatten()

classmethod tree_unflatten(aux, children)

concatenate_zip(*arrays)

draw_linear_residual(likelihood: ~nifty.re.likelihood.Likelihood, pos: ~nifty.re.evi.P, key, *, from_inverse: bool = True, point_estimates: ~nifty.re.evi.P | ~typing.Tuple[str] = (), cg: ~typing.Callable = <function cg>, cg_name: str | None = None, cg_kwargs: dict | None = None, jit_metric=False, _raise_nonposdef: bool = False) → tuple[P, int]

draw_residual(likelihood: ~nifty.re.likelihood.Likelihood, pos: ~nifty.re.evi.P, key, *, point_estimates: ~nifty.re.evi.P | ~typing.Tuple[str] = (), cg: ~typing.Callable = <function static_cg>, cg_name: str | None = None, cg_kwargs: dict | None = None, minimize: ~typing.Callable[[...], ~nifty.re.optimize.OptimizeResults] = <function _newton_cg>, minimize_kwargs={}, _raise_nonposdef: bool = False, _raise_notconverged: bool = False) → tuple[P, OptimizeResults]

nonlinearly_update_residual(likelihood=None, pos: ~nifty.re.evi.P = None, residual_sample=None, metric_sample_key=None, metric_sample_sign=1.0, *, point_estimates=(), minimize: ~typing.Callable[[...], ~nifty.re.optimize.OptimizeResults] = <function _newton_cg>, minimize_kwargs={}, jit_residual_funcs=False, _raise_notconverged=False) → tuple[P, OptimizeResults]

sample_likelihood(likelihood: Likelihood, point_estimates, primals, key)

wiener_filter_posterior(likelihood: LikelihoodWithModel, position: P | None = None, *, key, n_samples: int = 0, residual_map='lmap', draw_linear_kwargs: dict | None = None, jit=True, model_is_linear: bool | None = True, signal_space: bool | None = True, noise_covariance: callable | None = None) → Tuple[Samples, Tuple]

Computes Wiener filter solution for a standardized model. For non-linear models, the Wiener filter solution is computed for a linearized model.

Parameters:

• likelihood (LikelihoodWithModel) – Likelihood to be used for the Wiener filter.

• position (tree-like) – Position around which to linearize (if the model is non-linear).

• key (jax random number generation key)

• n_samples (int) – Number of samples to draw.

• residual_map (callable) – Map function used for the residual sample drawing.

• draw_linear_kwargs (dict) – Optional parameters for the conjugate gradient used to compute the posterior mean and to draw samples.

• jit (bool or callable, default=True) – Whether to JIT-compile the Wiener filter covariance.

• model_is_linear (bool) – Whether the model is linear. If the model is non-linear, you must specify a position around which to linearize it. For non-linear models, consider applying variational inference via optimize_kl() instead.

• signal_space (bool) – Wheter the Wiener filter should be solved in signal or data space. The result should be equal, up to numerical precision, thus depends on the condition number of the forward model.

• noise_covariance (callable) – Noise covariance of the data. This callable is only needed for the Wiener filter in data space.