def post_enhc(iter_index, json_file, machine_json, base_dir="./"):
base_dir = os.path.abspath(base_dir) + "/"
iter_name = make_iter_name(iter_index)
work_path = base_dir + iter_name + "/" + enhc_name + "/"
json_file = os.path.abspath(json_file)
json_file = os.path.abspath(json_file)
fp = open(json_file, 'r')
jdata = json.load(fp)
fp.close()
gmx_split = jdata["gmx_split_traj"]
gmx_split_log = "gmx_split.log"
gmx_split_cmd = cmd_append_log(gmx_split, gmx_split_log)
all_task = list(
filter(lambda x: os.path.isdir(x),
glob.glob(work_path + "/[0-9]*[0-9]")))
all_task.sort()
cwd = os.getcwd()
numb_walkers = jdata["numb_walkers"]
for ii in range(numb_walkers):
walker_path = work_path + make_walker_name(ii) + "/"
os.chdir(walker_path)
if os.path.isdir("confs"):
shutil.rmtree("confs")
os.makedirs("confs")
os.chdir(cwd)
print('rid.py:post_enhc:gmx_split_cmd', gmx_split_cmd)
print('rid.py:post_enhc:work path', work_path)
machine = set_machine(machine_json, target="post_enhc")
resources = set_resource(machine_json, target="post_enhc")
all_task_relpath = [os.path.relpath(ii, work_path) for ii in all_task]
gmx_split_task = [
Task(command=gmx_split_cmd,
task_work_path=ii,
outlog='gmx_split.log',
errlog='gmx_split.log') for ii in all_task_relpath
]
gmx_split_submission = Submission(work_base=work_path,
resources=resources,
machine=machine,
task_list=gmx_split_task)
gmx_split_submission.run_submission()
for ii in range(numb_walkers):
walker_path = work_path + make_walker_name(ii) + "/"
angles = np.loadtxt(walker_path + enhc_out_plm)
np.savetxt(walker_path + enhc_out_angle, angles[:, 1:], fmt="%.6f")
print("Post process of enhanced sampling finished.")
def run_enhc(iter_index,
json_file,
machine_json,
base_dir='./'):
json_file = os.path.abspath(json_file)
base_dir = os.path.abspath(base_dir) + "/"
iter_name = make_iter_name(iter_index)
work_path = base_dir + iter_name + "/" + enhc_name + "/"
fp = open(json_file, 'r')
jdata = json.load(fp)
fp.close()
gmx_prep = jdata["gmx_prep"]
gmx_run = jdata["gmx_run"]
enhc_thread = jdata["enhc_thread"]
gmx_run = gmx_run + (" -nt %d" % enhc_thread)
gmx_prep_log = "gmx_grompp.log"
gmx_run_log = "gmx_mdrun.log"
# assuming at least one walker
graph_files = glob.glob(work_path + (make_walker_name(0)) + "/*.pb")
if len(graph_files) != 0:
gmx_run = gmx_run + " -plumed " + enhc_plm
else:
gmx_run = gmx_run + " -plumed " + enhc_bf_plm
gmx_prep_cmd = cmd_append_log(gmx_prep, gmx_prep_log)
gmx_run_cmd = cmd_append_log(gmx_run, gmx_run_log)
numb_walkers = jdata["numb_walkers"]
all_task = list(filter(lambda x: os.path.isdir(
x), glob.glob(work_path + "/[0-9]*[0-9]")))
all_task.sort()
all_task_basedir = [os.path.relpath(ii, work_path) for ii in all_task]
print('run_enhc:work_path', work_path)
print('run_enhc:gmx_prep_cmd:', gmx_prep_cmd)
print('run_enhc:gmx_run_cmd:', gmx_run_cmd)
print('run_enhc:all_task:', all_task)
print('run_enhc:all_task_basedir:', all_task_basedir)
machine = set_machine(machine_json, target="enhcMD")
resources = set_resource(machine_json, target="enhcMD")
gmx_prep_task = [Task(command=gmx_prep_cmd, task_work_path=ii,
outlog='gmx_grompp.log', errlog='gmx_grompp.log') for ii in all_task_basedir]
gmx_prep_submission = Submission(
work_base=work_path, machine=machine, resources=resources, task_list=gmx_prep_task)
gmx_prep_submission.run_submission()
gmx_run_task = [Task(command=gmx_run_cmd, task_work_path=ii,
outlog='gmx_mdrun.log', errlog='gmx_mdrun.log') for ii in all_task_basedir]
gmx_run_submission = Submission(
work_base=work_path, machine=machine, resources=resources, task_list=gmx_run_task)
gmx_run_submission.run_submission()
def make_res(iter_index,
json_file,
cv_file,
mol_path,
base_dir="./",
custom_mdp=None):
json_file = os.path.abspath(json_file)
fp = open(json_file, 'r')
jdata = json.load(fp)
fp.close()
cv_file = os.path.abspath(cv_file)
numb_walkers = jdata["numb_walkers"]
bias_nsteps = jdata["bias_nsteps"]
bias_frame_freq = jdata["bias_frame_freq"]
nsteps = jdata["res_nsteps"]
frame_freq = jdata["res_frame_freq"]
sel_threshold = jdata["sel_threshold"]
max_sel = jdata["max_sel"]
cluster_threshold = jdata["cluster_threshold"]
init_numb_cluster_upper = int(jdata["init_numb_cluster_upper"])
init_numb_cluster_lower = int(jdata["init_numb_cluster_lower"])
init_numb_cluster = [init_numb_cluster_lower, init_numb_cluster_upper]
base_dir = os.path.abspath(base_dir) + "/"
iter_name = make_iter_name(iter_index)
enhc_path = base_dir + iter_name + "/" + enhc_name + "/"
res_path = base_dir + iter_name + "/" + res_name + "/"
create_path(res_path)
cwd = os.getcwd()
_conf_file = enhc_path + make_walker_name(0) + "/" + "conf.gro"
cv_dim_list = cal_cv_dim(_conf_file, cv_file)
cv_dim = sum(cv_dim_list)
cv_dih_dim = cv_dim_list[0]
ret_list = [True for ii in range(numb_walkers)]
weight = jdata["cv_weight_for_cluster"]
if type(weight) == list:
assert len(
weight
) == cv_dim, "Number of values in the weight list is not equal to the number of CVs."
elif type(weight) == float or type(weight) == int:
assert weight != 0
else:
raise TypeError(
"Invalid type of weight of CVs for clustering. Please use int or list instead."
)
# check if we have graph in enhc
for walker_idx in range(numb_walkers):
cls_sel = None
walker_path = enhc_path + walker_format % walker_idx + "/"
graph_files = glob.glob(walker_path + "*.pb")
if len(graph_files) != 0:
cluster_threshold = np.loadtxt(base_dir + "cluster_threshold.dat")
os.chdir(walker_path)
models = glob.glob("*.pb")
std_message = make_std(cv_dim,
dataset=enhc_out_angle,
models=models,
threshold=sel_threshold,
output="sel.out",
output_angle="sel.angle.out")
os.system('echo "{}" > sel.log'.format(std_message))
log_task("select with threshold %f" % sel_threshold)
os.chdir(cwd)
sel_idx = []
sel_angles = np.array([])
with open(walker_path + "sel.out") as fp:
for line in fp:
sel_idx += [int(x) for x in line.split()]
if len(sel_idx) != 0:
sel_angles = np.reshape(
np.loadtxt(walker_path + 'sel.angle.out'), [-1, cv_dim])
elif len(sel_idx) == 0:
np.savetxt(walker_path + 'num_of_cluster.dat', [0], fmt='%d')
np.savetxt(walker_path + 'cls.sel.out', [], fmt='%d')
continue
else:
cluster_threshold = jdata["cluster_threshold"]
sel_idx = range(
len(glob.glob(walker_path + enhc_out_conf + "conf*gro")))
sel_angles = np.loadtxt(walker_path + enhc_out_angle)
sel_angles = np.reshape(sel_angles, [-1, cv_dim])
np.savetxt(walker_path + 'sel.out', sel_idx, fmt='%d')
np.savetxt(walker_path + 'sel.angle.out', sel_angles, fmt='%.6f')
cls_sel, cluster_threshold = make_threshold(
walker_idx, walker_path, base_dir, sel_angles,
cluster_threshold, init_numb_cluster, cv_dih_dim, weight)
if cls_sel is None:
print(sel_angles, cluster_threshold, cv_dih_dim)
cls_sel = sel_from_cluster(sel_angles, cluster_threshold,
cv_dih_dim, weight)
conf_start = 0
conf_every = 1
sel_idx = np.array(sel_idx, dtype=np.int)
assert (
len(sel_idx) == sel_angles.shape[0]
), "{} selected indexes don't match {} selected angles.".format(
len(sel_idx), sel_angles.shape[0])
sel_idx = config_cls(sel_idx, cls_sel, max_sel, walker_path,
cluster_threshold, sel_angles)
res_angles = np.loadtxt(walker_path + enhc_out_angle)
res_angles = np.reshape(res_angles, [-1, cv_dim])
res_angles = res_angles[sel_idx]
np.savetxt(walker_path + 'cls.sel.out', sel_idx, fmt='%d')
np.savetxt(walker_path + 'cls.sel.angle.out', res_angles, fmt='%.6f')
res_confs = []
for ii in sel_idx:
res_confs.append(walker_path + enhc_out_conf + ("conf%d.gro" % ii))
assert (len(res_confs) == res_angles.shape[0]
), "number of enhc out conf does not match out angle"
assert (len(sel_idx) == res_angles.shape[0]
), "number of enhc out conf does not match number sel"
nconf = len(res_confs)
if nconf == 0:
ret_list[walker_idx] = False
continue
sel_list = make_sel_list(nconf, sel_idx)
log_task("selected %d confs, indexes: %s" % (nconf, sel_list))
make_conf(nconf,
res_path,
walker_idx,
walker_path,
sel_idx,
jdata,
mol_path,
conf_start=0,
conf_every=1,
custom_mdp=custom_mdp)
make_res_plumed(nconf,
jdata,
res_path,
walker_idx,
sel_idx,
res_angles,
_conf_file,
cv_file,
conf_start=0,
conf_every=1)
print("Restrained MD has been prepared.")
def make_enhc(iter_index,
json_file,
graph_files,
mol_dir,
cv_file,
base_dir='./',
custom_mdp=None):
base_dir = os.path.abspath(base_dir) + "/"
json_file = os.path.abspath(json_file)
cv_file = os.path.abspath(cv_file)
graph_files.sort()
fp = open(json_file, 'r')
jdata = json.load(fp)
fp.close()
numb_walkers = jdata["numb_walkers"]
enhc_trust_lvl_1 = jdata["bias_trust_lvl_1"]
enhc_trust_lvl_2 = jdata["bias_trust_lvl_2"]
nsteps = jdata["bias_nsteps"]
frame_freq = jdata["bias_frame_freq"]
num_of_cluster_threshold = jdata["num_of_cluster_threshold"]
dt = jdata["bias_dt"]
temperature = jdata["bias_temperature"]
iter_name = make_iter_name(iter_index)
work_path = base_dir + iter_name + "/" + enhc_name + "/"
mol_path = os.path.abspath(mol_dir) + "/"
conf_list = glob.glob(mol_path + "*gro")
conf_list.sort()
assert (len(conf_list) >=
numb_walkers), "not enough conf files in mol dir %s" % mol_path
create_path(work_path)
mol_files = ["topol.top"]
for walker_idx in range(numb_walkers):
walker_path = work_path + make_walker_name(walker_idx) + "/"
create_path(walker_path)
make_grompp(walker_path + "grompp.mdp", "bias", nsteps,
frame_freq, temperature=temperature, dt=dt, define='', custom_mdp=custom_mdp)
# make_grompp(walker_path + "grompp_restraint.mdp", "res", nsteps, frame_freq, temperature=temperature, dt=dt, define='-DPOSRE')
for ii in mol_files:
checkfile(walker_path + ii)
shutil.copy(mol_path + ii, walker_path)
# copy conf file
conf_file = conf_list[walker_idx]
checkfile(walker_path + "conf.gro")
shutil.copy(conf_file, walker_path + "conf.gro")
checkfile(walker_path + "conf_init.gro")
shutil.copy(conf_file, walker_path + "conf_init.gro")
# if have prev confout.gro, use as init conf
if iter_index > 0:
prev_enhc_path = base_dir + \
make_iter_name(iter_index-1) + "/" + enhc_name + \
"/" + make_walker_name(walker_idx) + "/"
prev_enhc_path = os.path.abspath(prev_enhc_path) + "/"
if os.path.isfile(prev_enhc_path + "confout.gro"):
os.remove(walker_path + "conf.gro")
rel_prev_enhc_path = os.path.relpath(
prev_enhc_path + "confout.gro", walker_path)
os.symlink(rel_prev_enhc_path, walker_path + "conf.gro")
else:
raise RuntimeError(
"cannot find prev output conf file " + prev_enhc_path + 'confout.gro')
log_task("use conf of iter " + make_iter_name(iter_index -
1) + " walker " + make_walker_name(walker_idx))
enhc_trust_lvl_1, enhc_trust_lvl_2 = adjust_lvl(
prev_enhc_path, num_of_cluster_threshold, jdata)
np.savetxt(walker_path+'trust_lvl1.dat',
[enhc_trust_lvl_1], fmt='%.6f')
make_plumed(walker_path, "dpbias", conf_file, cv_file)
make_plumed(walker_path, "bf", conf_file, cv_file)
prep_graph(graph_files, walker_path)
# config plumed
graph_list = get_graph_list(graph_files)
conf_enhc_plumed(walker_path + enhc_plm, "enhc", graph_list, enhc_trust_lvl_1=enhc_trust_lvl_1,
enhc_trust_lvl_2=enhc_trust_lvl_2, frame_freq=frame_freq, enhc_out_plm=enhc_out_plm)
conf_enhc_plumed(walker_path + enhc_bf_plm, "bf", graph_list,
frame_freq=frame_freq, enhc_out_plm=enhc_out_plm)
if len(graph_list) == 0:
log_task("brute force MD without NN acc")
else:
log_task("use NN model(s): " + graph_list)
log_task("set trust l1 and l2: %f %f" %
(enhc_trust_lvl_1, enhc_trust_lvl_2))
print("Enhanced sampling has prepared.")
def test_make_name(self):
self.assertTrue((utils.make_iter_name(5) == "iter.000005"))
self.assertTrue((utils.make_walker_name(5) == "005"))
pass