def slice_np_mcmc(
self,
num_samples: int,
potential_function: Callable,
context: torch.Tensor,
thin: int = 10,
warmup_steps: int = 20,
) -> Tensor:
# go into eval mode for evaluating during sampling
# XXX set eval mode outside of calls to sample
self.neural_net.eval()
posterior_sampler = SliceSampler(
utils.tensor2numpy(self._prior.sample((1,))).reshape(-1),
lp_f=potential_function,
thin=thin,
)
posterior_sampler.gen(warmup_steps)
samples = posterior_sampler.gen(num_samples)
# back to training mode
# XXX train exited in log_prob, entered here?
self.neural_net.train(True)
return torch.tensor(samples, dtype=torch.float32)
def run_slice_np(inits, seed):
# Seed current job.
np.random.seed(seed)
posterior_sampler = SliceSampler(
tensor2numpy(inits).reshape(-1),
lp_f=potential_function,
thin=thin,
# Show pbars of workers only for single worker
verbose=show_progress_bars and num_workers == 1,
)
if warmup_steps > 0:
posterior_sampler.gen(int(warmup_steps))
return posterior_sampler.gen(ceil(num_samples / num_chains))
def _slice_np_mcmc(
self,
num_samples: int,
potential_function: Callable,
thin: int,
warmup_steps: int,
) -> Tensor:
"""
Custom implementation of slice sampling using Numpy.
Args:
num_samples: Desired number of samples.
potential_function: A callable **class**.
thin: Thinning (subsampling) factor.
warmup_steps: Initial number of samples to discard.
Returns: Tensor of shape (num_samples, shape_of_single_theta).
"""
# Go into eval mode for evaluating during sampling
self.net.eval()
num_chains = self._mcmc_init_params.shape[0]
dim_samples = self._mcmc_init_params.shape[1]
all_samples = []
for c in range(num_chains):
posterior_sampler = SliceSampler(
utils.tensor2numpy(self._mcmc_init_params[c, :]).reshape(-1),
lp_f=potential_function,
thin=thin,
)
if warmup_steps > 0:
posterior_sampler.gen(int(warmup_steps))
all_samples.append(
posterior_sampler.gen(int(num_samples / num_chains)))
all_samples = np.stack(all_samples).astype(np.float32)
samples = torch.from_numpy(all_samples) # chains x samples x dim
# Final sample will be next init location
self._mcmc_init_params = samples[:, -1, :].reshape(
num_chains, dim_samples)
samples = samples.reshape(-1, dim_samples)[:num_samples, :]
assert samples.shape[0] == num_samples
# Back to training mode
self.net.train(True)
return samples.type(torch.float32)
def _slice_np_mcmc(
self,
num_samples: int,
potential_function: Callable,
initial_params: Tensor,
thin: int,
warmup_steps: int,
) -> Tensor:
"""
Custom implementation of slice sampling using Numpy.
Args:
num_samples: Desired number of samples.
potential_function: A callable **class**.
initial_params: Initial parameters for MCMC chain.
thin: Thinning (subsampling) factor.
warmup_steps: Initial number of samples to discard.
Returns: Tensor of shape (num_samples, shape_of_single_theta).
"""
num_chains = initial_params.shape[0]
dim_samples = initial_params.shape[1]
all_samples = []
for c in range(num_chains):
posterior_sampler = SliceSampler(
utils.tensor2numpy(initial_params[c, :]).reshape(-1),
lp_f=potential_function,
thin=thin,
)
if warmup_steps > 0:
posterior_sampler.gen(int(warmup_steps))
all_samples.append(
posterior_sampler.gen(int(num_samples / num_chains)))
all_samples = np.stack(all_samples).astype(np.float32)
samples = torch.from_numpy(all_samples) # chains x samples x dim
# Save sample as potential next init (if init_strategy == 'latest_sample').
self._mcmc_init_params = samples[:, -1, :].reshape(
num_chains, dim_samples)
samples = samples.reshape(-1, dim_samples)[:num_samples, :]
assert samples.shape[0] == num_samples
return samples.type(torch.float32)
def run_slice_np_vectorized(inits, seed):
# Seed current job.
np.random.seed(seed)
posterior_sampler = SliceSamplerVectorized(
init_params=tensor2numpy(inits),
log_prob_fn=potential_function,
num_chains=inits.shape[0],
# Show pbars of workers only for single worker
verbose=show_progress_bars and num_workers == 1,
)
warmup_ = warmup_steps * thin
num_samples_ = ceil((num_samples * thin) / num_chains)
samples = posterior_sampler.run(warmup_ + num_samples_)
samples = samples[:, warmup_:, :] # discard warmup steps
samples = samples[:, ::thin, :] # thin chains
samples = torch.from_numpy(samples) # chains x samples x dim
return samples
def np_potential(self, theta: np.ndarray) -> ScalarFloat:
r"""
Return conditional posterior log-probability or $-\infty$ if outside prior.
Args:
theta: Free parameters $\theta_i$, batch dimension 1.
Returns:
Conditional posterior log-probability $\log(p(\theta_i|\theta_j, x))$,
masked outside of prior.
"""
theta = torch.as_tensor(theta, dtype=torch.float32)
theta_condition = deepcopy(self.condition)
theta_condition[:, self.dims_to_sample] = theta
return self.potential_fn_provider.np_potential(
utils.tensor2numpy(theta_condition))
def _slice_np_mcmc(
self,
num_samples: int,
potential_function: Callable,
initial_params: Tensor,
thin: int,
warmup_steps: int,
vectorized: bool = False,
show_progress_bars: bool = True,
) -> Tensor:
"""
Custom implementation of slice sampling using Numpy.
Args:
num_samples: Desired number of samples.
potential_function: A callable **class**.
initial_params: Initial parameters for MCMC chain.
thin: Thinning (subsampling) factor.
warmup_steps: Initial number of samples to discard.
vectorized: Whether to use a vectorized implementation of
the Slice sampler (still experimental).
show_progress_bars: Whether to show a progressbar during sampling;
can only be turned off for vectorized sampler.
Returns: Tensor of shape (num_samples, shape_of_single_theta).
"""
num_chains = initial_params.shape[0]
dim_samples = initial_params.shape[1]
if not vectorized: # Sample all chains sequentially
all_samples = []
for c in range(num_chains):
posterior_sampler = SliceSampler(
utils.tensor2numpy(initial_params[c, :]).reshape(-1),
lp_f=potential_function,
thin=thin,
verbose=show_progress_bars,
)
if warmup_steps > 0:
posterior_sampler.gen(int(warmup_steps))
all_samples.append(
posterior_sampler.gen(ceil(num_samples / num_chains))
)
all_samples = np.stack(all_samples).astype(np.float32)
samples = torch.from_numpy(all_samples) # chains x samples x dim
else: # Sample all chains at the same time
posterior_sampler = SliceSamplerVectorized(
init_params=utils.tensor2numpy(initial_params),
log_prob_fn=potential_function,
num_chains=num_chains,
verbose=show_progress_bars,
)
warmup_ = warmup_steps * thin
num_samples_ = ceil((num_samples * thin) / num_chains)
samples = posterior_sampler.run(warmup_ + num_samples_)
samples = samples[:, warmup_:, :] # discard warmup steps
samples = samples[:, ::thin, :] # thin chains
samples = torch.from_numpy(samples) # chains x samples x dim
# Save sample as potential next init (if init_strategy == 'latest_sample').
self._mcmc_init_params = samples[:, -1, :].reshape(num_chains, dim_samples)
samples = samples.reshape(-1, dim_samples)[:num_samples, :]
assert samples.shape[0] == num_samples
return samples.type(torch.float32)
def summarize(
summary_writer,
summary,
round_,
true_observation,
parameter_bank,
observation_bank,
simulator,
posterior_samples_acceptance_rate=None,
):
# get ground truth if available
try:
(
_,
prior,
ground_truth_parameters,
ground_truth_observation,
) = simulators.get_simulator_prior_and_groundtruth(simulator.name)
# Update summaries.
except:
pass
try:
mmd = utils.unbiased_mmd_squared(
parameter_bank[-1],
simulator.get_ground_truth_posterior_samples(num_samples=1000),
)
summary["mmds"].append(mmd.item())
except:
pass
try:
# Median |x - x0| for most recent round.
median_observation_distance = torch.median(
torch.sqrt(
torch.sum(
(observation_bank[-1] -
true_observation.reshape(1, -1))**2,
dim=-1,
)))
summary["median_observation_distances"].append(
median_observation_distance.item())
summary_writer.add_scalar(
tag="median_observation_distance",
scalar_value=summary["median_observation_distances"][-1],
global_step=round_ + 1,
)
except:
pass
try:
# KDE estimate of negative log prob true parameters using
# parameters from most recent round.
negative_log_prob_true_parameters = -utils.gaussian_kde_log_eval(
samples=parameter_bank[-1],
query=ground_truth_parameters.reshape(1, -1),
)
summary["negative_log_probs_true_parameters"].append(
negative_log_prob_true_parameters.item())
summary_writer.add_scalar(
tag="negative_log_prob_true_parameters",
scalar_value=summary["negative_log_probs_true_parameters"][-1],
global_step=round_ + 1,
)
except:
pass
try:
# Rejection sampling acceptance rate
summary["rejection_sampling_acceptance-rates"].append(
posterior_samples_acceptance_rate)
summary_writer.add_scalar(
tag="rejection_sampling_acceptance_rate",
scalar_value=summary["rejection_sampling_acceptance_rates"][-1],
global_step=round_ + 1,
)
except:
pass
try:
# Plot most recently sampled parameters.
parameters = utils.tensor2numpy(parameter_bank[-1])
figure = utils.plot_hist_marginals(
data=parameters,
ground_truth=utils.tensor2numpy(ground_truth_parameters).reshape(
-1),
lims=simulator.parameter_plotting_limits,
)
summary_writer.add_figure(tag="posterior_samples",
figure=figure,
global_step=round_ + 1)
except:
pass
# Write quantities using SummaryWriter.
summary_writer.add_scalar(
tag="epochs_trained",
scalar_value=summary["epochs"][-1],
global_step=round_ + 1,
)
summary_writer.add_scalar(
tag="best_validation_log_prob",
scalar_value=summary["best_validation_log_probs"][-1],
global_step=round_ + 1,
)
if summary["mmds"]:
summary_writer.add_scalar(
tag="mmd",
scalar_value=summary["mmds"][-1],
global_step=round_ + 1,
)
summary_writer.flush()
return summary_writer, summary