def ratio_estimator_based_potential(
ratio_estimator: nn.Module,
prior: Distribution,
x_o: Optional[Tensor],
) -> Tuple[Callable, TorchTransform]:
r"""Returns the potential for ratio-based methods.
It also returns a transformation that can be used to transform the potential into
unconstrained space.
Args:
ratio_estimator: The neural network modelling likelihood-to-evidence ratio.
prior: The prior distribution.
x_o: The observed data at which to evaluate the likelihood-to-evidence ratio.
Returns:
The potential function and a transformation that maps
to unconstrained space.
"""
device = str(next(ratio_estimator.parameters()).device)
potential_fn = RatioBasedPotential(ratio_estimator,
prior,
x_o,
device=device)
theta_transform = mcmc_transform(prior, device=device)
return potential_fn, theta_transform
def posterior_estimator_based_potential(
posterior_estimator: nn.Module,
prior: Distribution,
x_o: Optional[Tensor],
) -> Tuple[Callable, TorchTransform]:
r"""Returns the potential for posterior-based methods.
It also returns a transformation that can be used to transform the potential into
unconstrained space.
The potential is the same as the log-probability of the `posterior_estimator`, but
it is set to $-\inf$ outside of the prior bounds.
Args:
posterior_estimator: The neural network modelling the posterior.
prior: The prior distribution.
x_o: The observed data at which to evaluate the posterior.
Returns:
The potential function and a transformation that maps
to unconstrained space.
"""
device = str(next(posterior_estimator.parameters()).device)
potential_fn = PosteriorBasedPotential(posterior_estimator,
prior,
x_o,
device=device)
theta_transform = mcmc_transform(prior, device=device)
return potential_fn, theta_transform
def test_prior_wrappers(wrapper, prior, kwargs):
"""Test prior wrappers to pytorch distributions."""
prior = wrapper(prior, **kwargs)
# use 2 here to test for minimal case >1
batch_size = 2
theta = prior.sample((batch_size,))
assert isinstance(theta, Tensor)
assert theta.shape[0] == batch_size
# Test log prob on batch of thetas.
log_probs = prior.log_prob(theta)
assert isinstance(log_probs, Tensor)
assert log_probs.shape[0] == batch_size
# Test return type
assert prior.sample().dtype == torch.float32
# Test support check.
within_support(prior, prior.sample((2,)))
# Test transform
mcmc_transform(prior)
def test_mcmc_transform(prior, enable_transform):
"""
Test whether the transform for MCMC returns the log_abs_det in the correct shape.
"""
num_samples = 1000
prior, _, _ = process_prior(prior)
tf = mcmc_transform(prior, enable_transform=enable_transform)
samples = prior.sample((num_samples, ))
unconstrained_samples = tf(samples)
samples_original = tf.inv(unconstrained_samples)
log_abs_det = tf.log_abs_det_jacobian(samples_original,
unconstrained_samples)
assert log_abs_det.shape == torch.Size([num_samples])
def test_transforms(prior, target_transform):
if isinstance(prior, UserNumpyUniform):
prior, *_ = process_prior(
prior,
dict(lower_bound=torch.zeros(2), upper_bound=torch.ones(2)),
)
transform = mcmc_transform(prior)
core_transform = transform._inv
if isinstance(core_transform, IndependentTransform):
core_transform = core_transform.base_transform
if hasattr(core_transform, "parts"):
transform_to_inspect = core_transform.parts[0]
else:
transform_to_inspect = core_transform
assert isinstance(transform_to_inspect, target_transform)
samples = prior.sample((2, ))
transformed_samples = transform(samples)
assert torch.allclose(samples, transform.inv(transformed_samples))
def __init__(
self,
potential_fn: Callable,
prior: Optional[TorchDistribution] = None,
q: Union[str, PyroTransformedDistribution, "VIPosterior",
Callable] = "maf",
theta_transform: Optional[TorchTransform] = None,
vi_method: str = "rKL",
device: str = "cpu",
x_shape: Optional[torch.Size] = None,
parameters: Iterable = [],
modules: Iterable = [],
):
"""
Args:
potential_fn: The potential function from which to draw samples.
prior: This is the prior distribution. Note that this is only
used to check/construct the variational distribution or within some
quality metrics. Please make sure that this matches with the prior
within the potential_fn. If `None` is given, we will try to infer it
from potential_fn or q, if this fails we raise an Error.
q: Variational distribution, either string, `TransformedDistribution`, or a
`VIPosterior` object. This specifies a parametric class of distribution
over which the best possible posterior approximation is searched. For
string input, we currently support [nsf, scf, maf, mcf, gaussian,
gaussian_diag]. You can also specify your own variational family by
passing a pyro `TransformedDistribution`.
Additionally, we allow a `Callable`, which allows you the pass a
`builder` function, which if called returns a distribution. This may be
useful for setting the hyperparameters e.g. `num_transfroms` within the
`get_flow_builder` method specifying the number of transformations
within a normalizing flow. If q is already a `VIPosterior`, then the
arguments will be copied from it (relevant for multi-round training).
theta_transform: Maps form prior support to unconstrained space. The
inverse is used here to ensure that the posterior support is equal to
that of the prior.
vi_method: This specifies the variational methods which are used to fit q to
the posterior. We currently support [rKL, fKL, IW, alpha]. Note that
some of the divergences are `mode seeking` i.e. they underestimate
variance and collapse on multimodal targets (`rKL`, `alpha` for alpha >
1) and some are `mass covering` i.e. they overestimate variance but
typically cover all modes (`fKL`, `IW`, `alpha` for alpha < 1).
device: Training device, e.g., `cpu`, `cuda` or `cuda:0`. We will ensure
that all other objects are also on this device.
x_shape: Shape of a single simulator output. If passed, it is used to check
the shape of the observed data and give a descriptive error.
parameters: List of parameters of the variational posterior. This is only
required for user-defined q i.e. if q does not have a `parameters`
attribute.
modules: List of modules of the variational posterior. This is only
required for user-defined q i.e. if q does not have a `modules`
attribute.
"""
super().__init__(potential_fn,
theta_transform,
device,
x_shape=x_shape)
# Especially the prior may be on another device -> move it...
self._device = device
self.potential_fn.device = device
move_all_tensor_to_device(self.potential_fn, device)
# Get prior and previous builds
if prior is not None:
self._prior = prior
elif hasattr(self.potential_fn, "prior") and isinstance(
self.potential_fn.prior, Distribution):
self._prior = self.potential_fn.prior
elif isinstance(q, VIPosterior) and isinstance(q._prior, Distribution):
self._prior = q._prior
else:
raise ValueError(
"We could not find a suitable prior distribution within `potential_fn`"
"or `q` (if a VIPosterior is given). Please explicitly specify a prior."
)
move_all_tensor_to_device(self._prior, device)
self._optimizer = None
# In contrast to MCMC we want to project into constrained space.
if theta_transform is None:
self.link_transform = mcmc_transform(self._prior).inv
else:
self.link_transform = theta_transform.inv
# This will set the variational distribution and VI method
self.set_q(q, parameters=parameters, modules=modules)
self.set_vi_method(vi_method)
self._purpose = (
"It provides Variational inference to .sample() from the posterior and "
"can evaluate the _normalized_ posterior density with .log_prob()."
)
def test_mnle_accuracy(sampler):
def mixed_simulator(theta):
# Extract parameters
beta, ps = theta[:, :1], theta[:, 1:]
# Sample choices and rts independently.
choices = Binomial(probs=ps).sample()
rts = InverseGamma(concentration=2 * torch.ones_like(beta),
rate=beta).sample()
return torch.cat((rts, choices), dim=1)
prior = MultipleIndependent(
[
Gamma(torch.tensor([1.0]), torch.tensor([0.5])),
Beta(torch.tensor([2.0]), torch.tensor([2.0])),
],
validate_args=False,
)
num_simulations = 2000
num_samples = 1000
theta = prior.sample((num_simulations, ))
x = mixed_simulator(theta)
# MNLE
trainer = MNLE(prior)
trainer.append_simulations(theta, x).train()
posterior = trainer.build_posterior()
mcmc_kwargs = dict(
num_chains=10,
warmup_steps=100,
method="slice_np_vectorized",
init_strategy="proposal",
)
for num_trials in [10]:
theta_o = prior.sample((1, ))
x_o = mixed_simulator(theta_o.repeat(num_trials, 1))
# True posterior samples
transform = mcmc_transform(prior)
true_posterior_samples = MCMCPosterior(
PotentialFunctionProvider(prior, atleast_2d(x_o)),
theta_transform=transform,
proposal=prior,
**mcmc_kwargs,
).sample((num_samples, ), show_progress_bars=False)
posterior = trainer.build_posterior(prior=prior, sample_with=sampler)
posterior.set_default_x(x_o)
if sampler == "vi":
posterior.train()
mnle_posterior_samples = posterior.sample(
sample_shape=(num_samples, ),
show_progress_bars=False,
**mcmc_kwargs if sampler == "mcmc" else {},
)
check_c2st(
mnle_posterior_samples,
true_posterior_samples,
alg=f"MNLE with {sampler}",
)