def fisher_pretraining(self, n_batch=5000, plot=True, batch_size=100, validation_split=0.1, epochs=1000, patience=20):
# Train on master only
if self.rank == 0:
# Generate fisher pre-training data
# Broader proposal
proposal = priors.TruncatedGaussian(self.theta_fiducial, 9*self.Finv, self.lower, self.upper)
# Anticipated covariance of the re-scaled data
Cdd = np.zeros((self.npar, self.npar))
for i in range(self.npar):
for j in range(self.npar):
Cdd[i,j] = self.Finv[i,j]/(self.fisher_errors[i]*self.fisher_errors[j])
Ldd = np.linalg.cholesky(Cdd)
Cddinv = np.linalg.inv(Cdd)
ln2pidetCdd = np.log(2*np.pi*np.linalg.det(Cdd))
# Sample parameters from some broad proposal
ps = np.zeros((3*n_batch, self.npar))
for i in range(0, n_batch):
# Draws from prior
ps[i,:] = (self.prior.draw() - self.theta_fiducial)/self.fisher_errors
# Draws from asymptotic posterior
ps[n_batch + i,:] = (self.asymptotic_posterior.draw() - self.theta_fiducial)/self.fisher_errors
# Drawn from Gaussian with 3x anticipated covariance matrix
ps[2*n_batch + i,:] = (proposal.draw() - self.theta_fiducial)/self.fisher_errors
# Sample data assuming a Gaussian likelihood
xs = np.array([pss + np.dot(Ldd, np.random.normal(0, 1, self.npar)) for pss in ps])
# Evaluate the logpdf at those values
fisher_logpdf_train = np.array([-0.5*np.dot(xs[i,:]-ps[i,:], np.dot(Cddinv, xs[i,:]-ps[i,:])) - 0.5*ln2pidetCdd for i in range(len(xs))])
# Construct the initial training-set
fisher_x_train = ps.astype(np.float32).reshape((3*n_batch, self.npar))
fisher_y_train = xs.astype(np.float32).reshape((3*n_batch, self.npar))
# Train the networks on these initial simulations
self.train_ndes(training_data=[fisher_x_train, fisher_y_train, np.atleast_2d(fisher_logpdf_train).reshape(-1,1)], validation_split = validation_split, epochs=epochs, batch_size=batch_size, patience=patience, mode='regression')
# Generate posterior samples
if plot==True:
print('Sampling approximate posterior...')
self.posterior_samples = self.emcee_sample(log_likelihood = lambda x: self.log_posterior_stacked(x, self.data)[0],
x0=[self.posterior_samples[-i,:] for i in range(self.nwalkers)], \
main_chain=self.posterior_chain_length)
print('Done.')
# Plot the posterior
self.triangle_plot([self.posterior_samples], \
savefig=True, \
filename='{}fisher_train_post.pdf'.format(self.results_dir))
def sequential_training(self, simulator, compressor, n_initial, n_batch, n_populations, proposal = None, \
simulator_args = None, compressor_args = None, safety = 5, plot = True, batch_size = 100, \
validation_split = 0.1, epochs = 300, patience = 20, seed_generator = None, \
save_intermediate_posteriors = True, sub_batch = 1):
# Set up the initial parameter proposal density
if proposal is None:
if self.input_normalization is 'fisher':
proposal = priors.TruncatedGaussian(self.theta_fiducial,
9 * self.Finv, self.lower,
self.upper)
elif self.Finv is not None:
proposal = priors.TruncatedGaussian(self.theta_fiducial,
9 * self.Finv, self.lower,
self.upper)
else:
proposal = self.prior
# Generate initial theta values from some broad proposal on
# master process and share with other processes. Overpropose
# by a factor of safety to (hopefully) cope gracefully with
# the possibility of some bad proposals. Assign indices into
# proposal array (self.inds_prop) and accepted arrays
# (self.inds_acpt) to allow for easy MPI communication.
if self.rank == 0:
ps = np.array([proposal.draw() for i in range(safety * n_initial)])
else:
ps = np.zeros((safety * n_initial, self.npar))
if self.use_mpi:
self.comm.Bcast(ps, root=0)
self.inds_prop = self.allocate_jobs(safety * n_initial)
self.inds_acpt = self.allocate_jobs(n_initial)
# Run simulations at those theta values
xs_batch, ps_batch = self.run_simulation_batch(
n_initial,
ps,
simulator,
compressor,
simulator_args,
compressor_args,
seed_generator=seed_generator,
sub_batch=sub_batch)
# Train on master only
if self.rank == 0:
# Construct the initial training-set
self.load_simulations(xs_batch, ps_batch)
# Train the network on these initial simulations
self.train_ndes(training_data=[self.x_train, self.y_train],
batch_size=max(self.n_sims // 8, batch_size),
validation_split=validation_split,
epochs=epochs,
patience=patience)
self.stacked_sequential_training_loss.append(
np.sum(
np.array([
self.training_loss[n][-1] * self.stacking_weights[n]
for n in range(self.n_ndes)
])))
self.stacked_sequential_validation_loss.append(
np.sum(
np.array([
self.validation_loss[n][-1] * self.stacking_weights[n]
for n in range(self.n_ndes)
])))
self.sequential_nsims.append(self.n_sims)
# Generate posterior samples
if save_intermediate_posteriors:
print('Sampling approximate posterior...')
self.posterior_samples = self.emcee_sample(log_likelihood=self.log_posterior_stacked, \
x0=[self.posterior_samples[-i,:] for i in range(self.nwalkers)], \
main_chain=self.posterior_chain_length)
# Save posterior samples to file
f = open('{}posterior_samples_0.dat'.format(self.results_dir),
'w')
np.savetxt(f, self.posterior_samples)
f.close()
print('Done.')
# If plot == True, plot the current posterior estimate
if plot == True:
self.triangle_plot([self.posterior_samples], \
savefig=True, \
filename='{}seq_train_post_0.pdf'.format(self.results_dir))
# Save attributes if save == True
if self.save == True:
self.saver()
# Loop through a number of populations
for i in range(n_populations):
# Propose theta values on master process and share with
# other processes. Again, ensure we propose more sets of
# parameters than needed to cope with bad params.
if self.rank == 0:
# Current population
print('Population {}/{}'.format(i + 1, n_populations))
# Sample the current posterior approximation
print('Sampling proposal density...')
self.proposal_samples = \
self.emcee_sample(log_likelihood=self.log_geometric_mean_proposal_stacked, \
x0=[self.proposal_samples[-j,:] for j in range(self.nwalkers)], \
main_chain=self.proposal_chain_length)
ps_batch = self.proposal_samples[-safety * n_batch:, :]
print('Done.')
else:
ps_batch = np.zeros((safety * n_batch, self.npar))
if self.use_mpi:
self.comm.Bcast(ps_batch, root=0)
# Run simulations
self.inds_prop = self.allocate_jobs(safety * n_batch)
self.inds_acpt = self.allocate_jobs(n_batch)
xs_batch, ps_batch = self.run_simulation_batch(
n_batch,
ps_batch,
simulator,
compressor,
simulator_args,
compressor_args,
seed_generator=seed_generator,
sub_batch=sub_batch)
# Train on master only
if self.rank == 0:
# Augment the training data
self.add_simulations(xs_batch, ps_batch)
# Train the network on these initial simulations
self.train_ndes(training_data=[self.x_train, self.y_train],
batch_size=max(self.n_sims // 8, batch_size),
validation_split=0.1,
epochs=epochs,
patience=patience)
self.stacked_sequential_training_loss.append(
np.sum(
np.array([
self.training_loss[n][-1] *
self.stacking_weights[n]
for n in range(self.n_ndes)
])))
self.stacked_sequential_validation_loss.append(
np.sum(
np.array([
self.validation_loss[n][-1] *
self.stacking_weights[n]
for n in range(self.n_ndes)
])))
self.sequential_nsims.append(self.n_sims)
# Generate posterior samples
if save_intermediate_posteriors:
print('Sampling approximate posterior...')
self.posterior_samples = self.emcee_sample(log_likelihood=self.log_posterior_stacked, \
x0=[self.posterior_samples[j] for j in range(self.nwalkers)], \
main_chain=self.posterior_chain_length)
# Save posterior samples to file
f = open(
'{}posterior_samples_{:d}.dat'.format(
self.results_dir, i + 1), 'w')
np.savetxt(f, self.posterior_samples)
f.close()
print('Done.')
# If plot == True
if plot == True:
# Plot the posterior
self.triangle_plot([self.posterior_samples], \
savefig=True, \
filename='{}seq_train_post_{:d}.pdf'.format(self.results_dir, i + 1))
# Plot training convergence
if plot == True:
# Plot the training loss convergence
self.sequential_training_plot(
savefig=True,
filename='{}seq_train_loss.pdf'.format(
self.results_dir))
def __init__(self, data, prior, nde, \
Finv = None, theta_fiducial = None, param_limits = None, param_names = None, nwalkers = 100, \
posterior_chain_length = 1000, proposal_chain_length = 100, \
rank = 0, n_procs = 1, comm = None, red_op = None, \
show_plot = True, results_dir = "", progress_bar = True, input_normalization = None,
graph_restore_filename = "graph_checkpoint", restore_filename = "restore.pkl", restore = False, save = True):
# Input validation
for i in range(len(nde)):
# Check all NDEs expect same number of parameters
if nde[0].n_parameters != nde[i].n_parameters:
err_msg = 'NDEs have inconsistent parameter counts. ' + \
'NDE 0: {:d} pars; NDE {:d}: {:d} pars.'
raise ValueError(err_msg.format(nde[0].n_parameters, \
i, nde[i].n_parameters))
# Check all NDEs expect same data length
if nde[0].n_data != nde[i].n_data:
err_msg = 'NDEs have inconsistent data counts. ' + \
'NDE 0: {:d} data; NDE {:d}: {:d} data.'
raise ValueError(err_msg.format(nde[0].n_data, \
i, nde[i].n_data))
# Check length of data provided is consistent with NDE
# expectations
if nde[i].n_data != len(data):
err_msg = 'inconsistent compressed data lengths. ' + \
'Compressed data have shape' + \
str(data.shape) + \
'; NDE {:d} expects length {:d}.'
raise ValueError(err_msg.format(i, nde[i].n_data))
# Data
self.data = data
self.D = len(data)
# Prior
self.prior = prior
# Number of parameters
self.npar = nde[0].n_parameters
# Initialize the NDEs, trainers, and stacking weights (for stacked density estimators)
self.n_ndes = len(nde)
self.nde = nde
self.trainer = [
pydelfi.train.ConditionalTrainer(nde[i])
for i in range(self.n_ndes)
]
self.stacking_weights = np.zeros(self.n_ndes)
# Tensorflow session for the NDE training
self.sess = tf.Session(config=tf.ConfigProto())
self.sess.run(tf.global_variables_initializer())
# Parameter limits
if param_limits is not None:
# Set to provided prior limits if provided
self.lower = param_limits[0]
self.upper = param_limits[1]
else:
# Else set to max and min float32
self.lower = np.ones(self.npar) * np.finfo(np.float32).min
self.upper = np.ones(self.npar) * np.finfo(np.float32).max
# Fisher matrix and fiducial parameters
if Finv is not None:
self.Finv = Finv
self.fisher_errors = np.sqrt(np.diag(self.Finv))
self.theta_fiducial = theta_fiducial
self.asymptotic_posterior = priors.TruncatedGaussian(
self.theta_fiducial, self.Finv, self.lower, self.upper)
else:
self.Finv = None
self.fisher_errors = None
self.theta_fiducial = None
self.asymptotic_posterior = None
# Re-scaling for inputs to NDE
self.input_normalization = input_normalization
if input_normalization is None:
self.x_mean = np.zeros(self.D)
self.x_std = np.ones(self.D)
self.p_mean = np.zeros(self.npar)
self.p_std = np.ones(self.npar)
elif input_normalization is 'fisher':
self.x_mean = self.theta_fiducial
self.x_std = self.fisher_errors
self.p_mean = self.theta_fiducial
self.p_std = self.fisher_errors
else:
self.x_mean, self.x_std, self.p_mean, self.p_std = input_normalization
# Training data [initialize empty]
self.ps = np.array([]).reshape(0, self.npar)
self.xs = np.array([]).reshape(0, self.D)
self.x_train = tf.placeholder(tf.float32, shape=(None, self.npar))
self.y_train = tf.placeholder(tf.float32, shape=(None, self.D))
self.n_sims = 0
# MCMC chain parameters
self.nwalkers = nwalkers
self.posterior_chain_length = posterior_chain_length
self.proposal_chain_length = proposal_chain_length
# MCMC samples of learned posterior
if self.asymptotic_posterior is not None:
self.posterior_samples = np.array([
self.asymptotic_posterior.draw()
for i in range(self.nwalkers * self.posterior_chain_length)
])
self.proposal_samples = np.array([
self.asymptotic_posterior.draw()
for i in range(self.nwalkers * self.proposal_chain_length)
])
else:
self.posterior_samples = np.array([
self.prior.draw()
for i in range(self.nwalkers * self.posterior_chain_length)
])
self.proposal_samples = np.array([
self.prior.draw()
for i in range(self.nwalkers * self.proposal_chain_length)
])
# Parameter names and ranges for plotting with GetDist
self.names = param_names
self.labels = param_names
self.ranges = dict(
zip(param_names,
[[self.lower[i], self.upper[i]] for i in range(self.npar)]))
self.show_plot = show_plot
# Results directory
self.results_dir = results_dir
# Training loss, validation loss
self.training_loss = [np.array([]) for i in range(self.n_ndes)]
self.validation_loss = [np.array([]) for i in range(self.n_ndes)]
self.stacked_sequential_training_loss = []
self.stacked_sequential_validation_loss = []
self.sequential_nsims = []
# MPI-specific setup
self.rank = rank
self.n_procs = n_procs
if n_procs > 1:
self.use_mpi = True
self.comm = comm
self.red_op = red_op
else:
self.use_mpi = False
# Show progress bars?
self.progress_bar = progress_bar
# Filenames for saving/restoring graph and attributes
self.graph_restore_filename = results_dir + graph_restore_filename
self.restore_filename = results_dir + restore_filename
# Save attributes of the ojbect as you go?
self.save = save
# Restore the graph and dynamic object attributes if restore = True
if restore == True:
# Restore the graph
saver = tf.train.Saver()
saver.restore(self.sess, self.graph_restore_filename)
# Restore the dynamic object attributes
self.stacking_weights, self.posterior_samples, self.proposal_samples, self.training_loss, self.validation_loss, self.stacked_sequential_training_loss, self.stacked_sequential_validation_loss, self.sequential_nsims, self.ps, self.xs, self.x_mean, self.x_std, self.p_mean, self.p_std = pickle.load(
open(self.restore_filename, 'rb'))