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],
x[sortidx][:num_particles],
)
def set_default_x(self, x: Tensor) -> "NeuralPosterior":
"""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 $\theta$ 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.
"""
self._x = process_x(x,
x_shape=self._x_shape,
allow_iid_x=self.potential_fn.allow_iid_x).to(
self._device)
self._map = None
return self
def set_x(self, x_o: Optional[Tensor]):
"""Check the shape of the observed data and, if valid, set it."""
if x_o is not None:
x_o = process_x(x_o, allow_iid_x=False).to(self.device)
self._x_o = x_o
for comp_potential in self.potential_fns:
comp_potential.set_x(x_o)
def _x_else_default_x(self, x: Optional[Array]) -> Tensor:
if x is not None:
return process_x(x,
x_shape=self._x_shape,
allow_iid_x=self.potential_fn.allow_iid_x)
elif self.default_x is None:
raise ValueError(
"Context `x` needed when a default has not been set."
"If you'd like to have a default, use the `.set_default_x()` method."
)
else:
return self.default_x
def test_process_x(x, x_shape):
process_x(x, x_shape)
def set_x(self, x_o: Optional[Tensor]):
"""Check the shape of the observed data and, if valid, set it."""
if x_o is not None:
x_o = process_x(x_o, allow_iid_x=self.allow_iid_x).to(self.device)
self._x_o = x_o
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