def make_task(iter_name, jdata, pres):
equi_conf = jdata['equi_conf']
equi_conf = os.path.abspath(equi_conf)
model = jdata['model']
model = os.path.abspath(model)
model_mass_map = jdata['model_mass_map']
if pres == None:
pres = jdata['pres']
elif 'pres' in jdata:
print('P = %f overrides the pres in json data' % pres)
jdata['pres'] = pres
create_path(iter_name)
cwd = os.getcwd()
os.chdir(iter_name)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
os.symlink(os.path.relpath(equi_conf), 'conf.lmp')
os.symlink(os.path.relpath(model), 'graph.pb')
lmp_str \
= _gen_lammps_relax('conf.lmp',
model_mass_map,
'graph.pb',
pres)
with open('in.lammps', 'w') as fp:
fp.write(lmp_str)
os.chdir(cwd)
def refine_tasks(from_task, to_task, err) :
jdata = json.load(open(os.path.join(from_task, 'in.json')))
equi_conf = get_task_file_abspath(from_task, jdata['equi_conf'])
model = get_task_file_abspath(from_task, jdata['model'])
create_path(to_task)
copied_conf = os.path.join(os.path.abspath(to_task), 'conf.lmp')
shutil.copyfile(equi_conf, copied_conf)
jdata['equi_conf'] = 'conf.lmp'
linked_model = os.path.join(os.path.abspath(to_task), 'graph.pb')
shutil.copyfile(model, linked_model)
jdata['model'] = 'graph.pb'
jdata['orig_task'] = from_task
jdata['refine_error'] = err
cwd = os.getcwd()
os.chdir(to_task)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
os.chdir(cwd)
from_name = os.path.join(from_task, '00.angle_on')
to_name = os.path.join(to_task, '00.angle_on')
_refine_tasks(from_name, to_name, err, 'angle_on')
from_name = os.path.join(from_task, '01.deep_on')
to_name = os.path.join(to_task, '01.deep_on')
_refine_tasks(from_name, to_name, err, 'deep_on')
from_name = os.path.join(from_task, '02.bond_angle_off')
to_name = os.path.join(to_task, '02.bond_angle_off')
_refine_tasks(from_name, to_name, err, 'bond_angle_off')
def make_task(iter_name, jdata, ens, temp, pres, avg_posi, npt_conf):
equi_conf = jdata['equi_conf']
equi_conf = os.path.abspath(equi_conf)
if npt_conf is not None:
npt_conf = os.path.abspath(npt_conf)
model = jdata['model']
model = os.path.abspath(model)
model_mass_map = jdata['model_mass_map']
nsteps = jdata['nsteps']
dt = jdata['dt']
stat_freq = jdata['stat_freq']
dump_freq = jdata['dump_freq']
tau_t = jdata['tau_t']
tau_p = jdata['tau_p']
if ens == None:
ens = jdata['ens']
elif 'ens' in jdata:
print('ens = %s overrides the ens in json data' % ens)
jdata['ens'] = ens
if temp == None:
temp = jdata['temp']
elif 'temp' in jdata:
print('T = %f overrides the temp in json data' % temp)
jdata['temp'] = temp
if 'npt' in ens:
if pres == None:
pres = jdata['pres']
elif 'pres' in jdata:
print('P = %f overrides the pres in json data' % pres)
jdata['pres'] = pres
create_path(iter_name)
cwd = os.getcwd()
os.chdir(iter_name)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
if npt_conf is None:
os.symlink(os.path.relpath(equi_conf), 'conf.lmp')
else:
open('conf.lmp', 'w').write(npt_equi_conf(npt_conf))
os.symlink(os.path.relpath(model), 'graph.pb')
lmp_str \
= _gen_lammps_input('conf.lmp',
model_mass_map,
'graph.pb',
nsteps,
dt,
ens,
temp,
pres,
tau_t = tau_t,
tau_p = tau_p,
prt_freq = stat_freq,
dump_freq = dump_freq,
dump_ave_posi = avg_posi)
with open('in.lammps', 'w') as fp:
fp.write(lmp_str)
os.chdir(cwd)
def _make_tasks(iter_name, jdata, step):
if step == 'soft_on':
all_lambda = parse_seq(jdata['lambda_soft_on'])
elif step == 'deep_on':
all_lambda = parse_seq(jdata['lambda_deep_on'])
elif step == 'soft_off':
all_lambda = parse_seq(jdata['lambda_soft_off'])
else:
raise RuntimeError('unknow step')
equi_conf = jdata['equi_conf']
equi_conf = os.path.abspath(equi_conf)
model_mass_map = jdata['model_mass_map']
model = jdata['model']
model = os.path.abspath(model)
soft_param = jdata['soft_param']
nsteps = jdata['nsteps']
dt = jdata['dt']
stat_freq = jdata['stat_freq']
copies = None
if 'copies' in jdata:
copies = jdata['copies']
temp = jdata['temp']
create_path(iter_name)
cwd = os.getcwd()
os.chdir(iter_name)
os.symlink(os.path.join('..', 'in.json'), 'in.json')
os.symlink(os.path.join('..', 'conf.lmp'), 'conf.lmp')
os.symlink(os.path.join('..', 'graph.pb'), 'graph.pb')
os.chdir(cwd)
for idx, ii in enumerate(all_lambda):
work_path = os.path.join(iter_name, 'task.%06d' % idx)
create_path(work_path)
os.chdir(work_path)
os.symlink(os.path.join('..', 'conf.lmp'), 'conf.lmp')
os.symlink(os.path.join('..', 'graph.pb'), 'graph.pb')
lmp_str \
= _gen_lammps_input_ideal(step,
'conf.lmp',
model_mass_map,
ii,
soft_param,
'graph.pb',
nsteps,
dt,
'nvt',
temp,
prt_freq = stat_freq,
copies = copies)
with open('in.lammps', 'w') as fp:
fp.write(lmp_str)
with open('lambda.out', 'w') as fp:
fp.write(str(ii))
os.chdir(cwd)
def train(self, train_ds, test_ds):
"""
train_ds: [enc_train_ds, dec_train_ds]
val_ds: [enc_val_ds, dec_val_ds]
test_ds: [enc_test_ds, dec_test_ds]
"""
root_model_dir = generate_model_dir(self.flags.root_model_dir)
create_path(root_model_dir)
checkpoint_dir = os.path.join(root_model_dir, 'training_checkpoints')
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(
model=self.model)
# save absl flags
flags_txt = (
self.flags.flags_into_string() +
'--enc_vocab_size={}\n'.format(self.model.enc_vocab_size) +
'--dec_vocab_size={}\n'.format(self.model.dec_vocab_size))
flags_file = open(os.path.join(root_model_dir, 'flags.txt'), 'w')
flags_file.write(flags_txt)
flags_file.close()
batch_count = 0
epoch_count = -1
batch_manager = Dataset(train_ds[0], train_ds[1], self.flags.num_epochs)
while batch_manager.epochs_done <= self.flags.num_epochs:
batch_count += 1
enc_input, dec_input, dec_output_true = batch_manager.batch(self.flags.batch_size)
batch_loss = self.train_step(enc_input, dec_input, dec_output_true)
if epoch_count != batch_manager.epochs_done:
test_enc_input = keras.preprocessing.sequence.pad_sequences(
test_ds[0], padding='post', value=0)
test_dec_input = keras.preprocessing.sequence.pad_sequences(
test_ds[1], padding='post', value=0)
test_dec_true = np.roll(test_dec_input, -1)
test_dec_true[:, -1] = 0
test_loss = self.eval(test_enc_input, test_dec_input, test_dec_true)
checkpoint.save(file_prefix=checkpoint_prefix)
epoch_count += 1
if batch_count % 100 == 0:
logging.info('Epoch {}, Batch {}, Loss {}, Last test loss: {}'.format(
batch_manager.epochs_done, batch_count, batch_loss, test_loss))
def _make_tasks_onephase(temp,
pres,
task_path,
jdata,
conf_file='conf.lmp',
graph_file='graph.pb'):
# assume that model and conf.lmp exist in the current dir
assert (os.path.isfile(conf_file))
assert (os.path.isfile(graph_file))
conf_file = os.path.abspath(conf_file)
graph_file = os.path.abspath(graph_file)
model_mass_map = jdata['model_mass_map']
# MD simulation protocol
nsteps = jdata['nsteps']
dt = jdata['dt']
stat_freq = jdata['stat_freq']
tau_t = jdata['tau_t']
tau_p = jdata['tau_p']
cwd = os.getcwd()
create_path(task_path)
os.chdir(task_path)
os.symlink(os.path.relpath(conf_file), 'conf.lmp')
os.symlink(os.path.relpath(graph_file), 'graph.pb')
# input for NPT MD
lmp_str \
= ti._gen_lammps_input('conf.lmp',
model_mass_map,
'graph.pb',
nsteps,
dt,
'npt',
temp,
pres,
tau_t = tau_t,
tau_p = tau_p,
prt_freq = stat_freq)
with open('thermo.out', 'w') as fp:
fp.write('%.16e %.16e' % (temp, pres))
with open('in.lammps', 'w') as fp:
fp.write(lmp_str)
os.chdir(cwd)
def make_tasks(iter_name, jdata) :
equi_conf = os.path.abspath(jdata['equi_conf'])
model = os.path.abspath(jdata['model'])
create_path(iter_name)
copied_conf = os.path.join(os.path.abspath(iter_name), 'conf.lmp')
shutil.copyfile(equi_conf, copied_conf)
jdata['equi_conf'] = 'conf.lmp'
linked_model = os.path.join(os.path.abspath(iter_name), 'graph.pb')
shutil.copyfile(model, linked_model)
jdata['model'] = 'graph.pb'
cwd = os.getcwd()
os.chdir(iter_name)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
os.chdir(cwd)
subtask_name = os.path.join(iter_name, '00.angle_on')
_make_tasks(subtask_name, jdata, 'angle_on')
subtask_name = os.path.join(iter_name, '01.deep_on')
_make_tasks(subtask_name, jdata, 'deep_on')
subtask_name = os.path.join(iter_name, '02.bond_angle_off')
_make_tasks(subtask_name, jdata, 'bond_angle_off')
def _setup_dpdt(task_path, jdata):
name_0 = jdata['phase_i']['name']
name_1 = jdata['phase_ii']['name']
conf_0 = jdata['phase_i']['equi_conf']
conf_1 = jdata['phase_ii']['equi_conf']
conf_0 = os.path.abspath(conf_0)
conf_1 = os.path.abspath(conf_1)
model = jdata['model']
model = os.path.abspath(model)
create_path(task_path)
conf_0_name = 'conf.%s.lmp' % '0'
conf_1_name = 'conf.%s.lmp' % '1'
# conf_0_name = 'conf.%s.lmp' % name_0
# conf_1_name = 'conf.%s.lmp' % name_1
copied_conf_0 = os.path.join(os.path.abspath(task_path), conf_0_name)
copied_conf_1 = os.path.join(os.path.abspath(task_path), conf_1_name)
shutil.copyfile(conf_0, copied_conf_0)
shutil.copyfile(conf_1, copied_conf_1)
linked_model = os.path.join(os.path.abspath(task_path), 'graph.pb')
shutil.copyfile(model, linked_model)
with open(os.path.join(os.path.abspath(task_path), 'in.json'), 'w') as fp:
json.dump(jdata, fp, indent=4)
def make_tasks(iter_name, jdata):
equi_conf = jdata['equi_conf']
equi_conf = os.path.abspath(equi_conf)
copies = None
if 'copies' in jdata:
copies = jdata['copies']
model = jdata['model']
model = os.path.abspath(model)
model_mass_map = jdata['model_mass_map']
nsteps = jdata['nsteps']
dt = jdata['dt']
stat_freq = jdata['stat_freq']
# thermos = jdata['thermos']
ens = jdata['ens']
path = jdata['path']
if 'nvt' in ens:
if path == 't':
temps = parse_seq(jdata['temps'])
tau_t = jdata['tau_t']
ntasks = len(temps)
else:
raise RuntimeError('supported path of nvt ens is \'t\'')
elif 'npt' in ens:
if path == 't':
temps = parse_seq(jdata['temps'])
press = jdata['press']
ntasks = len(temps)
elif path == 't-ginv':
temps = parse_seq_ginv(jdata['temps'])
press = jdata['press']
ntasks = len(temps)
elif path == 'p':
temps = jdata['temps']
press = parse_seq(jdata['press'])
ntasks = len(press)
else:
raise RuntimeError('supported path of npt ens are \'t\' or \'p\'')
tau_t = jdata['tau_t']
tau_p = jdata['tau_p']
else:
raise RuntimeError('invalid ens')
create_path(iter_name)
copied_conf = os.path.join(os.path.abspath(iter_name), 'conf.lmp')
shutil.copyfile(equi_conf, copied_conf)
jdata['equi_conf'] = 'conf.lmp'
linked_model = os.path.join(os.path.abspath(iter_name), 'graph.pb')
shutil.copyfile(model, linked_model)
jdata['model'] = 'graph.pb'
cwd = os.getcwd()
os.chdir(iter_name)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
os.chdir(cwd)
for ii in range(ntasks):
work_path = os.path.join(iter_name, 'task.%06d' % ii)
create_path(work_path)
os.chdir(work_path)
os.symlink(os.path.relpath(copied_conf), 'conf.lmp')
os.symlink(os.path.relpath(linked_model), 'graph.pb')
if 'nvt' in ens and path == 't':
lmp_str \
= _gen_lammps_input('conf.lmp',
model_mass_map,
'graph.pb',
nsteps,
dt,
ens,
temps[ii],
tau_t = tau_t,
prt_freq = stat_freq,
copies = copies)
with open('thermo.out', 'w') as fp:
fp.write('%f' % temps[ii])
elif 'npt' in ens and (path == 't' or path == 't-ginv'):
lmp_str \
= _gen_lammps_input('conf.lmp',
model_mass_map,
'graph.pb',
nsteps,
dt,
ens,
temps[ii],
press,
tau_t = tau_t,
tau_p = tau_p,
prt_freq = stat_freq,
copies = copies)
with open('thermo.out', 'w') as fp:
fp.write('%f' % (temps[ii]))
elif 'npt' in ens and path == 'p':
lmp_str \
= _gen_lammps_input('conf.lmp',
model_mass_map,
'graph.pb',
nsteps,
dt,
ens,
temps,
press[ii],
tau_t = tau_t,
tau_p = tau_p,
prt_freq = stat_freq,
copies = copies)
with open('thermo.out', 'w') as fp:
fp.write('%f' % (press[ii]))
else:
raise RuntimeError('invalid ens or path setting')
with open('in.lammps', 'w') as fp:
fp.write(lmp_str)
os.chdir(cwd)
def _make_tasks(iter_name,
jdata,
ref,
switch='one-step',
step='both',
link=False):
if 'crystal' not in jdata:
print('do not find crystal in jdata, assume vega')
jdata['crystal'] = 'vega'
crystal = jdata['crystal']
protect_eps = jdata['protect_eps']
if switch == 'one-step':
all_lambda = parse_seq(jdata['lambda'])
elif switch == 'two-step' or switch == 'three-step':
if step == 'deep_on':
all_lambda = parse_seq(jdata['lambda_deep_on'])
elif step == 'spring_off':
all_lambda = parse_seq(jdata['lambda_spring_off'])
elif step == 'lj_on':
all_lambda = parse_seq(jdata['lambda_lj_on'])
else:
raise RuntimeError('unknown step', step)
if all_lambda[0] == 0:
all_lambda[0] += protect_eps
if all_lambda[-1] == 1:
all_lambda[-1] -= protect_eps
equi_conf = jdata['equi_conf']
equi_conf = os.path.abspath(equi_conf)
model = jdata['model']
model = os.path.abspath(model)
model_mass_map = jdata['model_mass_map']
nsteps = jdata['nsteps']
dt = jdata['dt']
spring_k = jdata['spring_k']
sparam = jdata.get('soft_param', {})
if crystal == 'frenkel':
spring_k_1 = []
for ii in model_mass_map:
spring_k_1.append(spring_k * ii)
spring_k = spring_k_1
stat_freq = jdata['stat_freq']
copies = None
if 'copies' in jdata:
copies = jdata['copies']
temp = jdata['temp']
jdata['reference'] = ref
jdata['switch'] = switch
jdata['step'] = step
create_path(iter_name)
copied_conf = os.path.join(os.path.abspath(iter_name), 'conf.lmp')
if not link:
shutil.copyfile(equi_conf, copied_conf)
else:
cwd = os.getcwd()
os.chdir(iter_name)
os.symlink(os.path.relpath(equi_conf), 'conf.lmp')
os.chdir(cwd)
jdata['equi_conf'] = 'conf.lmp'
linked_model = os.path.join(os.path.abspath(iter_name), 'graph.pb')
if not link:
shutil.copyfile(model, linked_model)
else:
cwd = os.getcwd()
os.chdir(iter_name)
os.symlink(os.path.relpath(model), 'graph.pb')
os.chdir(cwd)
jdata['model'] = 'graph.pb'
langevin = jdata.get('langevin', True)
cwd = os.getcwd()
os.chdir(iter_name)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
os.chdir(cwd)
for idx, ii in enumerate(all_lambda):
work_path = os.path.join(iter_name, 'task.%06d' % idx)
create_path(work_path)
os.chdir(work_path)
os.symlink(os.path.relpath(copied_conf), 'conf.lmp')
os.symlink(os.path.relpath(linked_model), 'graph.pb')
if idx == 0:
ens = 'nvt-langevin'
else:
ens = 'nvt'
if langevin:
ens = 'nvt-langevin'
if jdata.get('ens', False):
ens = jdata.get('ens')
if ref == 'einstein':
lmp_str \
= _gen_lammps_input('conf.lmp',
model_mass_map,
ii,
'graph.pb',
spring_k,
nsteps,
dt,
ens,
temp,
prt_freq = stat_freq,
copies = copies,
switch = switch,
step = step,
sparam = sparam,
crystal = crystal)
elif ref == 'ideal':
lmp_str \
= _gen_lammps_input_ideal('conf.lmp',
model_mass_map,
ii,
'graph.pb',
nsteps,
dt,
ens,
temp,
prt_freq = stat_freq,
copies = copies)
else:
raise RuntimeError('unknow reference system type ' + ref)
with open('in.lammps', 'w') as fp:
fp.write(lmp_str)
with open('lambda.out', 'w') as fp:
fp.write(str(ii))
os.chdir(cwd)
def make_tasks(iter_name, jdata, ref, switch='one-step'):
equi_conf = os.path.abspath(jdata['equi_conf'])
model = os.path.abspath(jdata['model'])
if switch == 'one-step':
subtask_name = iter_name
_make_tasks(subtask_name, jdata, ref, step='both')
elif switch == 'two-step' or switch == 'three-step':
create_path(iter_name)
copied_conf = os.path.join(os.path.abspath(iter_name), 'conf.lmp')
shutil.copyfile(equi_conf, copied_conf)
jdata['equi_conf'] = 'conf.lmp'
linked_model = os.path.join(os.path.abspath(iter_name), 'graph.pb')
shutil.copyfile(model, linked_model)
jdata['model'] = 'graph.pb'
cwd = os.getcwd()
os.chdir(iter_name)
with open('in.json', 'w') as fp:
json.dump(jdata, fp, indent=4)
if switch == 'two-step':
subtask_name = '00.deep_on'
_make_tasks(subtask_name,
jdata,
ref,
switch=switch,
step='deep_on',
link=True)
subtask_name = '01.spring_off'
_make_tasks(subtask_name,
jdata,
ref,
switch=switch,
step='spring_off',
link=True)
elif switch == 'three-step':
subtask_name = '00.lj_on'
_make_tasks(subtask_name,
jdata,
ref,
switch=switch,
step='lj_on',
link=True)
subtask_name = '01.deep_on'
_make_tasks(subtask_name,
jdata,
ref,
switch=switch,
step='deep_on',
link=True)
subtask_name = '02.spring_off'
_make_tasks(subtask_name,
jdata,
ref,
switch=switch,
step='spring_off',
link=True)
else:
raise RuntimeError('unknow switch', switch)
os.chdir(cwd)
else:
raise RuntimeError('unknow switch', switch)
def save(self, save_to_path='./vocab'):
create_path(save_to_path)
with open(os.path.join(save_to_path, 'vocab_hub.pkl'), 'wb') as f:
pickle.dump(self, f)
def _make_tasks(iter_name, jdata, step) :
if step == 'angle_on' :
all_lambda = parse_seq(jdata['lambda_angle_on'])
protect_eps = jdata['protect_eps']
if all_lambda[0] == 0 :
all_lambda[0] += protect_eps
elif step == 'deep_on' :
all_lambda = parse_seq(jdata['lambda_deep_on'])
protect_eps = jdata['protect_eps']
if all_lambda[0] == 0 :
all_lambda[0] += protect_eps
elif step == 'bond_angle_off' :
all_lambda = parse_seq(jdata['lambda_bond_angle_off'])
protect_eps = jdata['protect_eps']
if all_lambda[-1] == 1 :
all_lambda[-1] -= protect_eps
equi_conf = jdata['equi_conf']
equi_conf = os.path.abspath(equi_conf)
model = jdata['model']
model = os.path.abspath(model)
model_mass_map = jdata['model_mass_map']
nsteps = jdata['nsteps']
dt = jdata['dt']
bparam = jdata['bond_param']
sparam = jdata['soft_param']
stat_freq = jdata['stat_freq']
ens = jdata['ens']
temp = jdata['temp']
pres = jdata['pres']
tau_t = jdata['tau_t']
tau_p = jdata['tau_p']
copies = None
if 'copies' in jdata :
copies = jdata['copies']
create_path(iter_name)
cwd = os.getcwd()
os.chdir(iter_name)
os.symlink(os.path.join('..', 'in.json'), 'in.json')
os.symlink(os.path.join('..', 'conf.lmp'), 'orig.lmp')
os.symlink(os.path.join('..', 'graph.pb'), 'graph.pb')
lines = water.add_bonds(open('orig.lmp').read().split('\n'))
open('conf.lmp', 'w').write('\n'.join(lines))
os.chdir(cwd)
for idx in range(len(all_lambda)) :
work_path = os.path.join(iter_name, 'task.%06d' % idx)
create_path(work_path)
os.chdir(work_path)
os.symlink(os.path.join('..', 'conf.lmp'), 'conf.lmp')
os.symlink(os.path.join('..', 'graph.pb'), 'graph.pb')
lmp_str \
= _gen_lammps_input(step,
'conf.lmp',
model_mass_map,
all_lambda[idx],
'graph.pb',
bparam,
sparam,
nsteps,
dt,
ens,
temp,
pres,
tau_t = tau_t,
tau_p = tau_p,
prt_freq = stat_freq,
copies = copies)
with open('in.lammps', 'w') as fp :
fp.write(lmp_str)
with open('lambda.out', 'w') as fp :
fp.write(str(all_lambda[idx]))
os.chdir(cwd)