def validate(model: Model, val_dataset: PairDataset, criterion: LossFunction):
model.eval()
total_loss = 0
count = 0
for batch_src, batch_tgt in tqdm(
val_dataset,
total=math.ceil(len(val_dataset) / val_dataset.batch_size)):
batch_src, lengths_src = pad_batch(batch_src, val_dataset.pad_index)
batch_tgt, lengths_tgt = pad_batch(batch_tgt, val_dataset.pad_index)
batch_src, lengths_src, batch_tgt, lengths_tgt = batch_src.cuda(
), lengths_src.cuda(), batch_tgt.cuda(), lengths_tgt.cuda()
with torch.no_grad():
output, mu, logvar = model(batch_src,
lengths_src,
batch_tgt,
lengths_tgt,
use_vae=False)
prediction_padding = torch.LongTensor([
val_dataset.pad_index for _ in range(lengths_tgt.size(0))
]).unsqueeze(0).cuda() # 1 x bs
prediction_tgt = torch.cat([batch_tgt[1:, :], prediction_padding],
dim=0) # still seqlen x bs
loss = criterion.forward(output, mu, logvar, prediction_tgt)
total_loss += loss
count += 1
average_val_loss = total_loss / count
print('average val loss: {}'.format(average_val_loss))
return average_val_loss
def train_mdn_with_proposal(save=True):
"""Use the prior proposal learnt by bootstrapping to train a brand new mdn."""
# load prior proposal and observations
_, obs_stats = helper.load(datadir + 'observed_data.pkl')
net, _, prior_proposal, _ = helper.load(netsdir + 'mdn_svi_proposal_prior_{0}.pkl'.format(n_bootstrap_iter-1))
n_inputs = n_percentiles
n_outputs = 3
n_samples = 5000
# generate data
ps = np.empty([n_samples, n_outputs])
stats = np.empty([n_samples, n_inputs])
for i in xrange(n_samples):
prior = 0.0
while prior < 0.5:
ps[i] = prior_proposal.gen()[0]
prior = eval_prior(*ps[i])
_, _, _, idts, _ = sim_likelihood(*ps[i])
stats[i] = calc_summary_stats(idts)
# train an mdn to give the posterior
minibatch = 100
maxiter = int(10000 * n_samples / minibatch)
monitor_every = 1000
net = mdn.replicate_gaussian_mdn(net, 8)
regularizer = lf.regularizerSvi(net.mps, net.sps, 0.1)
trainer = Trainer.Trainer(
model=net,
trn_data=[stats, ps],
trn_loss=net.mlprob + regularizer / n_samples,
trn_target=net.y
)
trainer.train(
maxiter=maxiter,
minibatch=minibatch,
show_progress=True,
monitor_every=monitor_every
)
# calculate the approximate posterior
mdn_mog = net.get_mog(obs_stats)
mdn_mog.prune_negligible_components(1.0e-6)
approx_posterior = mdn_mog / prior_proposal
# save the net
if save:
filename = netsdir + 'mdn_svi_proposal_hiddens_50_tanh_comps_8_sims_5k.pkl'
helper.save((net, approx_posterior), filename)
def train_mdn_with_proposal(save=True):
"""Use the prior proposal learnt by bootstrapping to train an mdn."""
# load prior proposal and observations
_, x, obs_data = helper.load(datadir + 'observed_data.pkl')
net, _, prior_proposal, _ = helper.load(
netsdir +
'mdn_svi_proposal_prior_{0}.pkl'.format(n_bootstrap_iter - 1))
n_inputs = n_data
n_outputs = n_dim
n_samples = 2000
# generate data
ws = np.empty([n_samples, n_outputs])
data = np.empty([n_samples, n_inputs])
for i in xrange(n_samples):
ws[i] = prior_proposal.gen()[0]
data[i] = gen_y_data(ws[i], x)
# train an mdn to give the posterior
minibatch = 100
maxiter = int(5000 * n_samples / minibatch)
monitor_every = 1000
regularizer = lf.regularizerSvi(net.mps, net.sps, 0.01)
trainer = Trainer.Trainer(model=net,
trn_data=[data, ws],
trn_loss=net.mlprob + regularizer / n_samples,
trn_target=net.y)
trainer.train(maxiter=maxiter,
minibatch=minibatch,
show_progress=True,
monitor_every=monitor_every)
# calculate the approximate posterior
mdn_mog = net.get_mog(obs_data)
approx_posterior = (mdn_mog * get_prior()) / prior_proposal
# save the net
if save:
filename = netsdir + 'mdn_svi_proposal_hiddens_50_tanh.pkl'
helper.save((net, approx_posterior), filename)
def train_mdn_proposal_prior(save=True):
"""Trains an svi mdn to return the proposal prior with boostrapping."""
n_iterations = n_bootstrap_iter
n_samples = 200
true_w, x, y = helper.load(datadir + 'observed_data.pkl')
obs_data = y
# create an mdn
n_inputs = obs_data.size
net = mdn.MDN_SVI(n_inputs=n_inputs,
n_hiddens=[50],
act_fun='tanh',
n_outputs=n_dim,
n_components=1)
regularizer = lf.regularizerSvi(net.mps, net.sps, 0.01)
prior = get_prior()
prior_proposal = prior
for iter in xrange(n_iterations):
# generate new data
ws = np.empty([n_samples, n_dim])
data = np.empty([n_samples, n_inputs])
dist = np.empty(n_samples)
for i in xrange(n_samples):
w = prior_proposal.gen()[0]
y = gen_y_data(w, x)
this_data = y
ws[i] = w
data[i] = this_data
dist[i] = calc_dist(this_data, obs_data)
print 'simulation {0}, distance = {1}'.format(i, dist[i])
# plot distance histogram
fig = plt.figure()
ax = fig.add_subplot(111)
ax.hist(dist, bins=int(np.sqrt(n_samples)))
ax.set_title('iteration = {0}'.format(iter + 1))
ax.set_xlim([0.0, 20.0])
plt.show(block=False)
# train an mdn to give the posterior
minibatch = 50
maxiter = int(1000 * n_samples / minibatch)
monitor_every = 10
trainer = Trainer.Trainer(model=net,
trn_data=[data, ws],
trn_loss=net.mlprob +
regularizer / n_samples,
trn_target=net.y)
trainer.train(maxiter=maxiter,
minibatch=minibatch,
show_progress=True,
monitor_every=monitor_every)
# calculate the approximate posterior
mdn_mog = net.get_mog(obs_data, n_samples=None)
approx_posterior = (mdn_mog * prior) / prior_proposal
prior_proposal = approx_posterior.project_to_gaussian()
# save the net and the approximate posterior
if save:
helper.save(
(net, approx_posterior, prior_proposal, dist),
netsdir + 'mdn_svi_proposal_prior_{0}.pkl'.format(iter))
def train_prior_proposal_with_bootstrapping(save=True):
"""Trains an svi mdn to return the posterior with boostrapping."""
n_samples = 400
true_ps, obs_stats = helper.load(datadir + 'observed_data.pkl')
# create an mdn
n_inputs = len(obs_stats)
n_outputs = len(true_ps)
net = mdn.MDN_SVI(n_inputs=n_inputs, n_hiddens=[50], act_fun='tanh', n_outputs=n_outputs, n_components=1)
regularizer = lf.regularizerSvi(net.mps, net.sps, 0.01)
prior_proposal = None
for iter in xrange(n_bootstrap_iter):
# generate new data
ps = np.empty([n_samples, n_outputs])
stats = np.empty([n_samples, n_inputs])
dist = np.empty(n_samples)
for i in xrange(n_samples):
prior = 0.0
while prior < 0.5:
ps[i] = sim_prior() if iter == 0 else prior_proposal.gen()[0]
prior = eval_prior(*ps[i])
_, _, _, idts, _ = sim_likelihood(*ps[i])
stats[i] = calc_summary_stats(idts)
dist[i] = calc_dist(stats[i], obs_stats)
print 'simulation {0}, distance = {1}'.format(i, dist[i])
# plot distance histogram
fig = plt.figure()
ax = fig.add_subplot(111)
ax.hist(dist, bins=int(np.sqrt(n_samples)))
ax.set_title('iteration = {0}'.format(iter + 1))
ax.set_xlim([0.0, 1.0])
plt.show(block=False)
# train an mdn to give the posterior
minibatch = 50
maxiter = int(1500 * n_samples / minibatch)
monitor_every = 10
trainer = Trainer.Trainer(
model=net,
trn_data=[stats, ps],
trn_loss=net.mlprob + regularizer / n_samples,
trn_target=net.y
)
trainer.train(
maxiter=maxiter,
minibatch=minibatch,
show_progress=True,
monitor_every=monitor_every
)
# calculate the approximate posterior
mdn_mog = net.get_mog(obs_stats, n_samples=None)
approx_posterior = mdn_mog if iter == 0 else mdn_mog / prior_proposal
prior_proposal = approx_posterior.project_to_gaussian()
# save the net and the approximate posterior
if save:
helper.save((net, approx_posterior, prior_proposal, dist), netsdir + 'mdn_svi_proposal_prior_{0}.pkl'.format(iter))
def train(model: Model,
train_dataset: PairDataset,
criterion: LossFunction,
optimizer: optim.Optimizer,
max_grad_norm=None,
original_parameter_vector=None,
parameter_crit=None,
parameter_crit_weight=None):
assert (original_parameter_vector is None) == (parameter_crit is None) == (
parameter_crit_weight is None)
model.train()
print_every = 200
count = 0
total_loss = 0
total_param_loss = 0
warned = False
for batch_src, batch_tgt in tqdm(
train_dataset,
total=math.ceil(len(train_dataset) / train_dataset.batch_size)):
optimizer.zero_grad()
batch_src, lengths_src = pad_batch(batch_src, train_dataset.pad_index)
batch_tgt, lengths_tgt = pad_batch(batch_tgt, train_dataset.pad_index)
batch_src, lengths_src, batch_tgt, lengths_tgt = batch_src.cuda(
), lengths_src.cuda(), batch_tgt.cuda(), lengths_tgt.cuda()
output, mu, logvar = model(batch_src, lengths_src, batch_tgt,
lengths_tgt)
prediction_padding = torch.LongTensor([
train_dataset.pad_index for _ in range(lengths_tgt.size(0))
]).unsqueeze(0).cuda() # 1 x bs
prediction_tgt = torch.cat([batch_tgt[1:, :], prediction_padding],
dim=0) # still seqlen x bs
loss = criterion.forward(output, mu, logvar, prediction_tgt)
total_loss += loss
if original_parameter_vector is not None:
parameter_diff_loss = parameter_crit(
parameters_to_vector(model.parameters()),
original_parameter_vector)
total_param_loss += parameter_diff_loss
loss = loss + parameter_diff_loss * parameter_crit_weight
loss.backward()
if max_grad_norm is not None:
total_norm = compute_grad_norm(model)
if not warned and total_norm > max_grad_norm:
print('clipping gradient norm')
warned = True
nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()
count += 1
if count % print_every == 0:
print('average train loss: {}'.format(total_loss / print_every))
total_loss = 0
if original_parameter_vector is not None:
print('average train param diff loss: {}'.format(
total_param_loss / print_every))
total_param_loss = 0
def main(args: Namespace):
if args.unconditional:
assert args.morgan_similarity_threshold == 0 # shouldn't care about inputs in this case
i2s = None
if args.checkpoint_dir is not None:
assert args.checkpoint_path is None
for _, _, files in os.walk(args.checkpoint_dir):
for fname in files:
if fname.endswith('.pt'):
args.checkpoint_path = os.path.join(
args.checkpoint_dir, fname)
if args.checkpoint_path is not None:
print('loading model from checkpoint')
model, i2s = load_model(args)
full_train_dataset = PairDataset(
path=args.train_path,
i2s=i2s,
batch_size=args.batch_size,
extra_vocab_path=args.extra_precursors_path
if args.extra_precursors_path is not None else None,
max_data=args.train_max_data
if args.train_max_data is not None else None)
pair_datasets = full_train_dataset.split([0.9, 0.1], seed=0)
train_dataset, val_dataset = pair_datasets[0], pair_datasets[1]
predict_dataset = SourceDataset(path=args.val_path,
i2s=train_dataset.i2s,
s2i=train_dataset.s2i,
pad_index=train_dataset.pad_index,
start_index=train_dataset.start_index,
end_index=train_dataset.end_index,
batch_size=args.batch_size)
if args.checkpoint_path is None:
print('building model from scratch')
model = Model(args=args,
vocab_size=len(train_dataset.i2s),
pad_index=train_dataset.pad_index,
start_index=train_dataset.start_index,
end_index=train_dataset.end_index)
for param in model.parameters():
if param.dim() == 1:
nn.init.constant_(param, 0)
else:
nn.init.xavier_normal_(param)
print(model)
print('num params: {:,}'.format(
sum(p.numel() for p in model.parameters() if p.requires_grad)))
model = model.cuda()
chemprop_predictor = ChempropPredictor(args)
criterion = LossFunction(train_dataset.pad_index, args.kl_weight)
optimizer = optim.Adam(model.parameters(), lr=args.init_lr)
scheduler = lr_scheduler.ExponentialLR(optimizer, 0.9)
for epoch in range(args.epochs):
print('epoch {}'.format(epoch))
train_dataset.reshuffle(seed=epoch)
train(model=model,
train_dataset=train_dataset,
criterion=criterion,
optimizer=optimizer,
max_grad_norm=args.max_grad_norm)
val_loss = validate(model=model,
val_dataset=val_dataset,
criterion=criterion)
os.makedirs(os.path.join(args.save_dir, 'epoch' + str(epoch)),
exist_ok=True)
train_dataset.save(
os.path.join(args.save_dir, 'epoch' + str(epoch),
'train_pairs.csv'))
save_model(model=model,
i2s=train_dataset.i2s,
path=os.path.join(args.save_dir, 'epoch' + str(epoch),
'val_loss_{}.pt'.format(val_loss)))
predict(model=model,
predict_dataset=predict_dataset,
save_dir=os.path.join(args.save_dir, 'epoch' + str(epoch)),
args=args,
chemprop_predictor=chemprop_predictor
if not args.no_predictor_at_val else None,
sample=not args.greedy_prediction,
num_predictions=args.val_num_predictions,
print_filter_frac=args.print_filter_frac)
if epoch % args.evaluate_every == 0:
evaluate(pred_smiles_dir=os.path.join(args.save_dir,
'epoch' + str(epoch)),
train_path=args.train_path,
val_path=args.val_path,
checkpoint_dir=args.chemprop_dir,
computed_prop=args.computed_prop,
prop_min=args.prop_min,
sim_thresholds=[0.2, 0.4, 0.6, 0.8, 0.9, 1.0],
chemprop_predictor=chemprop_predictor,
prop_max=args.prop_max,
unconditional=args.unconditional)
scheduler.step()
if args.self_train_epochs > 0:
# store parameters of current model for a loss to constrain it not to stray too far
original_parameter_vector = parameters_to_vector(
model.parameters()).data
parameter_crit = nn.MSELoss()
args.epoch_length = len(train_dataset.src) // 2
# Get properties of target molecules in train set
train_dataset.tgt_props = np.array(
chemprop_predictor(train_dataset.tgt_smiles))
augmented_train_dataset = deepcopy(train_dataset)
epochs_to_dataset_creation = 0
for epoch in range(args.epochs, args.epochs + args.self_train_epochs):
print('self train epoch {}'.format(epoch))
if epochs_to_dataset_creation == 0:
train_dataset.reshuffle(seed=epoch)
if args.self_train_max_data is not None:
self_train_dataset = deepcopy(train_dataset)
self_train_dataset.src, self_train_dataset.tgt = \
self_train_dataset.src[:args.self_train_max_data], self_train_dataset.tgt[:args.self_train_max_data]
self_train_dataset.src_smiles, self_train_dataset.tgt_smiles = \
self_train_dataset.src_smiles[:args.self_train_max_data], self_train_dataset.tgt_smiles[:args.self_train_max_data]
if hasattr(self_train_dataset, 'src_props'):
self_train_dataset.src_props = self_train_dataset.src_props[:
args
.
self_train_max_data]
if hasattr(self_train_dataset, 'tgt_props'):
self_train_dataset.tgt_props = self_train_dataset.tgt_props[:
args
.
self_train_max_data]
else:
self_train_dataset = deepcopy(train_dataset)
if args.extra_precursors_path is not None:
self_train_dataset.add_dummy_pairs(
args.extra_precursors_path)
translations, props = generate_self_train_translations(
train_dataset=self_train_dataset,
model=model,
chemprop_predictor=chemprop_predictor,
args=args,
k=args.k)
if not args.keep_translations: # drop old translations and restart
augmented_train_dataset = deepcopy(self_train_dataset)
if args.unconditional:
new_train_dataset = deepcopy(self_train_dataset)
new_train_dataset.tgt_smiles = translations
new_train_dataset.tgt = [
list(self_train_dataset.smiles2indices(smiles))
for smiles in new_train_dataset.tgt_smiles
]
new_train_dataset.tgt = np.array(new_train_dataset.tgt)
new_train_dataset.src_smiles = translations # any dummy is fine
new_train_dataset.src = [
list(self_train_dataset.smiles2indices(smiles))
for smiles in new_train_dataset.src_smiles
]
new_train_dataset.src = np.array(new_train_dataset.src)
else:
new_train_dataset = deepcopy(self_train_dataset)
new_train_dataset.src = np.concatenate(
[self_train_dataset.src for _ in range(args.k)])
new_train_dataset.src_smiles = []
for _ in range(args.k):
new_train_dataset.src_smiles += self_train_dataset.src_smiles
new_train_dataset.tgt = []
for i in range(args.k):
new_train_dataset.tgt += [
translations[j][i]
for j in range(len(translations))
]
new_train_dataset.tgt_smiles = [
self_train_dataset.indices2smiles(indices)
for indices in new_train_dataset.tgt
]
new_train_dataset.tgt = np.array(new_train_dataset.tgt)
if args.replace_old_dataset:
augmented_train_dataset = new_train_dataset
else:
augmented_train_dataset.add(new_train_dataset)
if not args.unconditional:
augmented_train_dataset.filter_dummy_pairs(
need_props=False) # filters src == tgt pairs
epochs_to_dataset_creation = args.epochs_per_dataset
augmented_train_dataset.reshuffle(seed=epoch, need_props=False)
epochs_to_dataset_creation -= 1
train(model=model,
train_dataset=augmented_train_dataset,
criterion=criterion,
optimizer=optimizer,
max_grad_norm=args.max_grad_norm,
original_parameter_vector=original_parameter_vector,
parameter_crit=parameter_crit,
parameter_crit_weight=args.l2_diff_weight)
val_loss = validate(model=model,
val_dataset=val_dataset,
criterion=criterion)
os.makedirs(os.path.join(args.save_dir, 'epoch' + str(epoch)),
exist_ok=True)
augmented_train_dataset.save(
os.path.join(args.save_dir, 'epoch' + str(epoch),
'train_pairs.csv'))
save_model(model=model,
i2s=train_dataset.i2s,
path=os.path.join(args.save_dir, 'epoch' + str(epoch),
'val_loss_{}.pt'.format(val_loss)))
predict(model=model,
predict_dataset=predict_dataset,
save_dir=os.path.join(args.save_dir, 'epoch' + str(epoch)),
args=args,
chemprop_predictor=chemprop_predictor
if not args.no_predictor_at_val else None,
sample=not args.greedy_prediction,
num_predictions=args.val_num_predictions,
print_filter_frac=args.print_filter_frac)
evaluate(pred_smiles_dir=os.path.join(args.save_dir,
'epoch' + str(epoch)),
train_path=args.train_path,
val_path=args.val_path,
checkpoint_dir=args.chemprop_dir,
computed_prop=args.computed_prop,
prop_min=args.prop_min,
sim_thresholds=[0.2, 0.4, 0.6, 0.8, 0.9, 1.0],
chemprop_predictor=chemprop_predictor,
prop_max=args.prop_max,
unconditional=args.unconditional)
scheduler.step()
# for convenient evaluation
os.makedirs(os.path.join(args.save_dir, 'final_eval'), exist_ok=True)
test_dataset = SourceDataset(path=args.test_path,
i2s=train_dataset.i2s,
s2i=train_dataset.s2i,
pad_index=train_dataset.pad_index,
start_index=train_dataset.start_index,
end_index=train_dataset.end_index,
batch_size=args.batch_size)
predict(model=model,
predict_dataset=test_dataset,
save_dir=os.path.join(args.save_dir, 'final_eval'),
args=args,
chemprop_predictor=chemprop_predictor
if not args.no_predictor_at_val else None,
sample=not args.greedy_prediction,
num_predictions=args.val_num_predictions,
print_filter_frac=args.print_filter_frac)
if args.final_eval_chemprop_dir is not None:
args.computed_prop = None
args.chemprop_dir = args.final_eval_chemprop_dir
chemprop_predictor = ChempropPredictor(args)
if args.final_eval_computed_prop is not None:
args.chemprop_dir = None
args.computed_prop = args.final_eval_computed_prop
chemprop_predictor = ChempropPredictor(args)
evaluate(pred_smiles_dir=os.path.join(args.save_dir, 'final_eval'),
train_path=args.train_path,
val_path=args.test_path,
checkpoint_dir=args.chemprop_dir,
computed_prop=args.computed_prop,
prop_min=args.prop_min,
sim_thresholds=[0.2, 0.4, 0.6, 0.8, 0.9, 1.0],
chemprop_predictor=chemprop_predictor,
prop_max=args.prop_max,
unconditional=args.unconditional)
def train_mdn_proposal_prior(save=True):
"""
Train a proposal prior using bootstrapping.
"""
n_iterations = n_bootstrap_iter
n_data = 500
# read data
pilot_means, pilot_stds = helper.load(datadir + 'pilot_run_results.pkl')
obs_stats = helper.load(datadir + 'obs_stats.pkl')
obs_stats -= pilot_means
obs_stats /= pilot_stds
# create an mdn
net = mdn.MDN_SVI(n_inputs=9, n_hiddens=[50], act_fun='tanh', n_outputs=4, n_components=1)
regularizer = lf.regularizerSvi(net.mps, net.sps, 0.01)
prior_proposal = None
for iter in xrange(n_iterations):
# generate new data
params = []
stats = []
dist = []
i = 0
while i < n_data:
prop_params = sim_prior_params() if iter == 0 else np.exp(prior_proposal.gen())[0]
if np.any(np.log(prop_params) < log_prior_min) or np.any(np.log(prop_params) > log_prior_max):
continue
try:
lv = mjp.LotkaVolterra(init, prop_params)
states = lv.sim_time(dt, duration, max_n_steps=max_n_steps)
except mjp.SimTooLongException:
continue
sum_stats = calc_summary_stats(states)
sum_stats -= pilot_means
sum_stats /= pilot_stds
params.append(prop_params)
stats.append(sum_stats)
dist.append(calc_dist(sum_stats, obs_stats))
i += 1
print 'simulation {0}, distance = {1}'.format(i, dist[-1])
params = np.array(params)
stats = np.array(stats)
dist = np.array(dist)
# plot distance histogram
fig = plt.figure()
ax = fig.add_subplot(111)
ax.hist(dist, bins=int(np.sqrt(n_data)))
ax.set_title('iteration = {0}'.format(iter + 1))
ax.set_xlim([0.0, 12.0])
plt.show(block=False)
# train an mdn to give the posterior
minibatch = 100
maxiter = int(2000 * n_data / minibatch)
monitor_every = 100
trainer = Trainer.Trainer(
model=net,
trn_data=[stats, np.log(params)],
trn_loss=net.mlprob + regularizer / n_data,
trn_target=net.y
)
trainer.train(
maxiter=maxiter,
minibatch=minibatch,
show_progress=True,
monitor_every=monitor_every
)
# calculate the approximate posterior
mdn_mog = net.get_mog(obs_stats)
approx_posterior = mdn_mog if iter == 0 else mdn_mog / prior_proposal
prior_proposal = approx_posterior.project_to_gaussian()
# save the net and the approximate posterior
if save:
helper.save((net, approx_posterior, prior_proposal, dist), netsdir + 'mdn_svi_proposal_prior_{0}.pkl'.format(iter))
def train_mdn_with_proposal(save=True):
"""Use the prior proposal learnt by bootstrapping to train an mdn."""
# load prior proposal and observations
pilot_means, pilot_stds = helper.load(datadir + 'pilot_run_results.pkl')
obs_stats = helper.load(datadir + 'obs_stats.pkl')
obs_stats -= pilot_means
obs_stats /= pilot_stds
net, _, prior_proposal, _ = helper.load(netsdir + 'mdn_svi_proposal_prior_{0}.pkl'.format(n_bootstrap_iter-1))
n_samples = 2000
# generate data
params = []
stats = []
i = 0
while i < n_samples:
prop_params = np.exp(prior_proposal.gen())[0]
if np.any(np.log(prop_params) < log_prior_min) or np.any(np.log(prop_params) > log_prior_max):
continue
try:
lv = mjp.LotkaVolterra(init, prop_params)
states = lv.sim_time(dt, duration, max_n_steps=max_n_steps)
except mjp.SimTooLongException:
continue
sum_stats = calc_summary_stats(states)
sum_stats -= pilot_means
sum_stats /= pilot_stds
params.append(prop_params)
stats.append(sum_stats)
i += 1
params = np.array(params)
stats = np.array(stats)
# train an mdn to give the posterior
minibatch = 100
maxiter = int(5000 * n_samples / minibatch)
monitor_every = 1000
regularizer = lf.regularizerSvi(net.mps, net.sps, 0.01)
trainer = Trainer.Trainer(
model=net,
trn_data=[stats, np.log(params)],
trn_loss=net.mlprob + regularizer / n_samples,
trn_target=net.y
)
trainer.train(
maxiter=maxiter,
minibatch=minibatch,
show_progress=True,
monitor_every=monitor_every
)
# calculate the approximate posterior
mdn_mog = net.get_mog(obs_stats)
mdn_mog.prune_negligible_components(1.0e-3)
approx_posterior = mdn_mog / prior_proposal
# save the net
if save:
filename = netsdir + 'mdn_svi_proposal_hiddens_50_tanh_comps_1_sims_2k.pkl'
helper.save((net, approx_posterior), filename)