def _set_xo_and_sample_initial_population(
self, x_o, num_particles: int, num_initial_pop: int,
) -> Tuple[Tensor, float, Tensor]:
"""Return particles, epsilon and distances of initial population."""
assert (
num_particles <= num_initial_pop
), "number of initial round simulations must be greater than population size"
theta = self.prior.sample((num_initial_pop,))
x = self._simulate_with_budget(theta)
# Infer x shape to test and set x_o.
self.x_shape = x[0].unsqueeze(0).shape
self.x_o = process_x(x_o, self.x_shape)
distances = self.distance(self.x_o, x)
sortidx = torch.argsort(distances)
particles = theta[sortidx][:num_particles]
# Take last accepted distance as epsilon.
initial_epsilon = distances[sortidx][num_particles - 1]
if not torch.isfinite(initial_epsilon):
initial_epsilon = 1e8
return particles, initial_epsilon, distances[sortidx][:num_particles]
def _handle_x_o_wrt_amortization(self, x_o: Optional[Tensor],
x_shape: torch.Size,
num_rounds: int) -> None:
"""
Check whether provided `x_o` is consistent with `num_rounds`. For multi-round
inference, set `x_o` as the default observation in the posterior.
Warns:
When `x_o` is provided, yet `num_rounds=1`, i.e. inference is amortized.
Raises:
When `x_o` is absent, yet `num_rounds>1`, i.e. inference is focused.
"""
if num_rounds == 1 and x_o is not None:
warn("Providing `x_o` in a single-round scenario has no effect.")
elif num_rounds > 1 and x_o is None:
raise ValueError(
"Observed data `x_o` is required if `num_rounds>1` since "
"samples in all rounds after the first one are drawn from "
"the posterior `p(theta|x_o)`.")
if num_rounds > 1:
processed_x = process_x(x_o, x_shape)
self._posterior._x_o_training_focused_on = processed_x
self._posterior.default_x = processed_x
def set_default_x(self, x: Tensor) -> "NeuralPosterior":
"""Set new default x for `.sample(), .log_prob` to use as conditioning context.
This is a pure convenience to avoid having to repeatedly specify `x` in calls to
`.sample()` and `.log_prob()` - only θ 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.
Args:
x: The default observation to set for the posterior $p(theta|x)$.
Returns:
`NeuralPosterior` that will use a default `x` when not explicitly passed.
"""
processed_x = process_x(x, self._x_shape)
self._x = processed_x
return self
def test_process_x(x, x_shape):
process_x(x, x_shape)
def __call__(
self,
x_o: Union[Tensor, ndarray],
num_simulations: int,
eps: Optional[float] = None,
quantile: Optional[float] = None,
return_distances: bool = False,
) -> Union[Distribution, Tuple[Distribution, Tensor]]:
r"""Run MCABC.
Args:
x_o: Observed data.
num_simulations: Number of simulations to run.
eps: Acceptance threshold $\epsilon$ for distance between observed and
simulated data.
quantile: Upper quantile of smallest distances for which the corresponding
parameters are returned, e.g, q=0.01 will return the top 1%. Exactly
one of quantile or `eps` have to be passed.
return_distances: Whether to return the distances corresponding to the
selected parameters.
Returns:
posterior: Empirical distribution based on selected parameters.
distances: Tensor of distances of the selected parameters.
"""
# Exactly one of eps or quantile need to be passed.
assert (eps is not None) ^ (
quantile
is not None), "Eps or quantile must be passed, but not both."
# Simulate and calculate distances.
theta = self.prior.sample((num_simulations, ))
x = self._batched_simulator(theta)
# Infer shape of x to test and set x_o.
self.x_shape = x[0].unsqueeze(0).shape
self.x_o = process_x(x_o, self.x_shape)
distances = self.distance(self.x_o, x)
# Select based on acceptance threshold epsilon.
if eps is not None:
is_accepted = distances < eps
num_accepted = is_accepted.sum().item()
assert num_accepted > 0, f"No parameters accepted, eps={eps} too small"
theta_accepted = theta[is_accepted]
distances_accepted = distances[is_accepted]
# Select based on quantile on sorted distances.
elif quantile is not None:
num_top_samples = int(num_simulations * quantile)
sort_idx = torch.argsort(distances)
theta_accepted = theta[sort_idx][:num_top_samples]
distances_accepted = distances[sort_idx][:num_top_samples]
else:
raise ValueError("One of epsilon or quantile has to be passed.")
posterior = Empirical(theta_accepted,
log_weights=ones(theta_accepted.shape[0]))
if return_distances:
return posterior, distances_accepted
else:
return posterior
def __call__(
self,
x_o: Union[Tensor, ndarray],
num_simulations: int,
eps: Optional[float] = None,
quantile: Optional[float] = None,
return_distances: bool = False,
return_x_accepted: bool = False,
lra: bool = False,
sass: bool = False,
sass_fraction: float = 0.25,
sass_expansion_degree: int = 1,
) -> Union[Distribution, Tuple[Distribution, Tensor]]:
r"""Run MCABC.
Args:
x_o: Observed data.
num_simulations: Number of simulations to run.
eps: Acceptance threshold $\epsilon$ for distance between observed and
simulated data.
quantile: Upper quantile of smallest distances for which the corresponding
parameters are returned, e.g, q=0.01 will return the top 1%. Exactly
one of quantile or `eps` have to be passed.
return_distances: Whether to return the distances corresponding to
the accepted parameters.
return_distances: Whether to return the simulated data corresponding to
the accepted parameters.
lra: Whether to run linear regression adjustment as in Beaumont et al. 2002
sass: Whether to determine semi-automatic summary statistics as in
Fearnhead & Prangle 2012.
sass_fraction: Fraction of simulation budget used for the initial sass run.
sass_expansion_degree: Degree of the polynomial feature expansion for the
sass regression, default 1 - no expansion.
Returns:
posterior: Empirical distribution based on selected parameters.
distances: Tensor of distances of the selected parameters.
"""
# Exactly one of eps or quantile need to be passed.
assert (eps is not None) ^ (
quantile
is not None), "Eps or quantile must be passed, but not both."
# Run SASS and change the simulator and x_o accordingly.
if sass:
num_pilot_simulations = int(sass_fraction * num_simulations)
self.logger.info(
f"Running SASS with {num_pilot_simulations} pilot samples.")
num_simulations -= num_pilot_simulations
pilot_theta = self.prior.sample((num_pilot_simulations, ))
pilot_x = self._batched_simulator(pilot_theta)
sass_transform = self.get_sass_transform(pilot_theta, pilot_x,
sass_expansion_degree)
simulator = lambda theta: sass_transform(
self._batched_simulator(theta))
x_o = sass_transform(x_o)
else:
simulator = self._batched_simulator
# Simulate and calculate distances.
theta = self.prior.sample((num_simulations, ))
x = simulator(theta)
# Infer shape of x to test and set x_o.
self.x_shape = x[0].unsqueeze(0).shape
self.x_o = process_x(x_o, self.x_shape)
distances = self.distance(self.x_o, x)
# Select based on acceptance threshold epsilon.
if eps is not None:
is_accepted = distances < eps
num_accepted = is_accepted.sum().item()
assert num_accepted > 0, f"No parameters accepted, eps={eps} too small"
theta_accepted = theta[is_accepted]
distances_accepted = distances[is_accepted]
x_accepted = x[is_accepted]
# Select based on quantile on sorted distances.
elif quantile is not None:
num_top_samples = int(num_simulations * quantile)
sort_idx = torch.argsort(distances)
theta_accepted = theta[sort_idx][:num_top_samples]
distances_accepted = distances[sort_idx][:num_top_samples]
x_accepted = x[sort_idx][:num_top_samples]
else:
raise ValueError("One of epsilon or quantile has to be passed.")
# Maybe adjust theta with LRA.
if lra:
self.logger.info("Running Linear regression adjustment.")
theta_adjusted = self.run_lra(theta_accepted,
x_accepted,
observation=self.x_o)
else:
theta_adjusted = theta_accepted
posterior = Empirical(theta_adjusted,
log_weights=ones(theta_accepted.shape[0]))
if return_distances and return_x_accepted:
return posterior, distances_accepted, x_accepted
if return_distances:
return posterior, distances_accepted
if return_x_accepted:
return posterior, x_accepted
else:
return posterior