Exemplo n.º 1
0
def make_enhc(iter_index, json_file, graph_files):
    graph_files.sort()
    fp = open(json_file, 'r')
    jdata = json.load(fp)
    bPosre = jdata.get("gmx_posre", False)
    numb_walkers = jdata["numb_walkers"]
    template_dir = jdata["template_dir"]
    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_threshhold = jdata["num_of_cluster_threshhold"]

    iter_name = make_iter_name(iter_index)
    work_path = iter_name + "/" + enhc_name + "/"
    mol_path = template_dir + "/" + mol_name + "/"
    enhc_path = template_dir + "/" + enhc_name + "/"
    conf_list = glob.glob(mol_path + "conf*gro")
    conf_list.sort()
    assert (len(conf_list) >=
            numb_walkers), "not enough conf files in mol dir %s" % mol_path

    create_path(work_path)
    kappa = np.linspace(2, 16, 8)
    dis_kappa = np.linspace(4.5, 1, 8)  # a weak distance restraint.
    for walker_idx in range(numb_walkers):
        kk = kappa[walker_idx]
        walker_path = work_path + make_walker_name(walker_idx) + "/"
        create_path(walker_path)
        # copy md ifles
        copy_file_list(mol_files, mol_path, walker_path)
        # copy conf file
        conf_file = conf_list[walker_idx]
        if os.path.exists(walker_path + "conf.gro"):
            os.remove(walker_path + "conf.gro")
        try:
            shutil.copy(conf_file, walker_path + "conf.gro")
        except:
            pass
        if os.path.exists(walker_path + "conf_init.gro"):
            os.remove(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):
            kk = kappa[(walker_idx + iter_index) % 8]
            prev_enhc_path = 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")
                os.remove(walker_path + "conf_init.gro")
                os.symlink(prev_enhc_path + "confout.gro",
                           walker_path + "conf.gro")
                os.symlink(prev_enhc_path + "conf_init.gro",
                           walker_path + "conf_init.gro")
                log_task("use conf of iter " + make_iter_name(iter_index - 1) +
                         " walker " + make_walker_name(walker_idx))
            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))
            ###########################################
            num_of_cluster = np.loadtxt(prev_enhc_path + 'num_of_cluster.dat')
            pre_trust_lvl1 = np.loadtxt(prev_enhc_path + 'trust_lvl1.dat')
            if num_of_cluster < num_of_cluster_threshhold:
                enhc_trust_lvl_1 = pre_trust_lvl1 * 1.5
                enhc_trust_lvl_2 = enhc_trust_lvl_1 + 1
            else:
                enhc_trust_lvl_1 = jdata["bias_trust_lvl_1"]
                enhc_trust_lvl_2 = enhc_trust_lvl_1 + 1
            if enhc_trust_lvl_1 > jdata["bias_trust_lvl_1"] * 8:
                enhc_trust_lvl_1 = jdata["bias_trust_lvl_1"]
                enhc_trust_lvl_2 = enhc_trust_lvl_1 + 1
        np.savetxt(walker_path + 'trust_lvl1.dat', [enhc_trust_lvl_1],
                   fmt='%.6f')
        # copy enhc file
        for ii in enhc_files:
            if os.path.exists(walker_path + ii):
                os.remove(walker_path + ii)
            try:
                shutil.copy(enhc_path + ii, walker_path)
            except:
                pass
        # copy graph files
        for ii in graph_files:
            file_name = os.path.basename(ii)
            abs_path = os.path.abspath(ii)
            if os.path.exists(walker_path + file_name):
                os.remove(walker_path + file_name)
            os.symlink(abs_path, walker_path + file_name)
        # config MD
        mol_conf_file = walker_path + "grompp.mdp"
        if bPosre:
            mol_conf_file = walker_path + "grompp_restraint.mdp"
        make_grompp_enhc(mol_conf_file, nsteps, frame_freq)
        # config plumed
        graph_list = ""
        counter = 0
        for ii in graph_files:
            file_name = os.path.basename(ii)
            if counter == 0:
                graph_list = "%s" % file_name
            else:
                graph_list = "%s,%s" % (graph_list, file_name)
            counter = counter + 1
        posre_files = glob.glob(walker_path + 'posre*.itp')
        for posre_file in posre_files:
            replace(posre_file, 'TEMP', '%d' % kk)
        plm_conf = walker_path + enhc_plm
        replace(plm_conf, "MODEL=[^ ]* ", ("MODEL=%s " % graph_list))
        replace(plm_conf, "TRUST_LVL_1=[^ ]* ",
                ("TRUST_LVL_1=%f " % enhc_trust_lvl_1))
        replace(plm_conf, "TRUST_LVL_2=[^ ]* ",
                ("TRUST_LVL_2=%f " % enhc_trust_lvl_2))
        replace(plm_conf, "STRIDE=[^ ]* ", ("STRIDE=%d " % frame_freq))
        replace(plm_conf, "FILE=[^ ]* ", ("FILE=%s " % enhc_out_plm))

        plm_bf_conf = walker_path + enhc_bf_plm
        replace(plm_bf_conf, "STRIDE=[^ ]* ", ("STRIDE=%d " % frame_freq))
        replace(plm_bf_conf, "FILE=[^ ]* ", ("FILE=%s " % enhc_out_plm))

        # molecule_name = os.getcwd().split('/')[-1].split('.')[0]
        # distance_list = get_distance('../{}/{}/{}.pdb'.format(molecule_name, molecule_name, molecule_name))
        # CA_list = get_CA_atom(walker_path + 'conf.gro')
        # ret = add_distance_restrain(CA_list, distance_list, dis_buttom=0.2, dis_kappa=dis_kappa[(walker_idx+iter_index)%8])
        # with open(plm_conf, 'a') as plm:
        #     plm.write(ret)

        # with open(plm_bf_conf, 'a') as plm:
        #     plm.write(ret)

        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))
Exemplo n.º 2
0
def make_temp(iter_index, json_file, graph_files):
    graph_files.sort()
    fp = open(json_file, 'r')
    jdata = json.load(fp)
    numb_walkers = jdata["numb_walkers"]
    template_dir = jdata["template_dir"]
    nsteps = jdata["temp_nsteps"]
    frame_freq = jdata["temp_frame_freq"]
    start_temp = jdata["start_temp"]

    iter_name = make_iter_name(iter_index)
    work_path = iter_name + "/" + temp_name + "/"
    mol_path = template_dir + "/" + mol_name + "/"
    temp_path = template_dir + "/" + temp_name + "/"
    conf_list = glob.glob(mol_path + "conf*gro")
    conf_list.sort()
    assert (len(conf_list) >=
            numb_walkers), "not enough conf files in mol dir %s" % mol_path

    create_path(work_path)

    for walker_idx in range(numb_walkers):
        walker_path = work_path + make_walker_name(walker_idx) + "/"
        create_path(walker_path)
        # copy md ifles
        for ii in mol_files:
            if os.path.exists(walker_path + ii):
                os.remove(walker_path + ii)
            shutil.copy(mol_path + ii, walker_path)
        # copy conf file
        conf_file = conf_list[walker_idx]
        if os.path.exists(walker_path + "conf.gro"):
            os.remove(walker_path + "conf.gro")
        shutil.copy(conf_file, walker_path + "conf.gro")
        # if have prev confout.gro, use as init conf
        if (iter_index > 0):
            prev_temp_path = make_iter_name(
                iter_index -
                1) + "/" + temp_name + "/" + make_walker_name(walker_idx) + "/"
            prev_temp_path = os.path.abspath(prev_temp_path) + "/"
            if os.path.isfile(prev_temp_path + "confout.gro"):
                os.remove(walker_path + "conf.gro")
                os.symlink(prev_temp_path + "confout.gro",
                           walker_path + "conf.gro")
            log_task("use conf of iter " + make_iter_name(iter_index - 1) +
                     " walker " + make_walker_name(walker_idx))
        # copy temp file
        for ii in temp_files:
            if os.path.exists(walker_path + ii):
                os.remove(walker_path + ii)
            shutil.copy(temp_path + ii, walker_path)
        # copy graph files
        for ii in graph_files:
            file_name = os.path.basename(ii)
            abs_path = os.path.abspath(ii)
            if os.path.exists(walker_path + file_name):
                os.remove(walker_path + file_name)
            os.symlink(abs_path, walker_path + file_name)
        # config MD
        mol_conf_file = walker_path + "grompp.mdp"
        make_grompp_enhc(mol_conf_file, nsteps, frame_freq)
        # config plumed
        if iter_index == 0:
            cur_temp = start_temp
        else:
            cur_temp = np.loadtxt(os.path.join(prev_temp_path, "next.temp"))
        log_task(("use temp of %f") % (cur_temp))
        log_task(("length of traj %d") % (nsteps))
        np.savetxt(os.path.join(walker_path, 'cur.temp'), [cur_temp])
        plm_conf = walker_path + temp_plm
        replace(plm_conf, "TEMP=[^ ]* ", ("TEMP=%s " % cur_temp))
        replace(plm_conf, "STRIDE=[^ ]* ", ("STRIDE=%d " % frame_freq))
        replace(plm_conf, "FILE=[^ ]* ", ("FILE=%s " % temp_out_plm))
Exemplo n.º 3
0
def make_enhc(iter_index, json_file, graph_files):
    graph_files.sort()
    fp = open(json_file, 'r')
    jdata = json.load(fp)
    numb_walkers = jdata["numb_walkers"]
    template_dir = jdata["template_dir"]
    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_threshhold = jdata["num_of_cluster_threshhold"]

    iter_name = make_iter_name(iter_index)
    work_path = iter_name + "/" + enhc_name + "/"
    mol_path = template_dir + "/" + mol_name + "/"
    enhc_path = template_dir + "/" + enhc_name + "/"
    conf_list = glob.glob(mol_path + "conf*gro")
    conf_list.sort()
    assert (len(conf_list) >=
            numb_walkers), "not enough conf files in mol dir %s" % mol_path

    create_path(work_path)

    for walker_idx in range(numb_walkers):
        walker_path = work_path + make_walker_name(walker_idx) + "/"
        create_path(walker_path)
        # copy md ifles
        for ii in mol_files:
            if os.path.exists(walker_path + ii):
                os.remove(walker_path + ii)
            shutil.copy(mol_path + ii, walker_path)
        # copy conf file
        conf_file = conf_list[walker_idx]
        if os.path.exists(walker_path + "conf.gro"):
            os.remove(walker_path + "conf.gro")
        shutil.copy(conf_file, walker_path + "conf.gro")
        # if have prev confout.gro, use as init conf
        if (iter_index > 0):
            prev_enhc_path = 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")
                os.symlink(prev_enhc_path + "confout.gro",
                           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))
            ###########################################
            num_of_cluster = np.loadtxt(prev_enhc_path + 'num_of_cluster.dat')
            pre_trust_lvl1 = np.loadtxt(prev_enhc_path + 'trust_lvl1.dat')
            if num_of_cluster < num_of_cluster_threshhold:
                enhc_trust_lvl_1 = pre_trust_lvl1 * 1.5
                enhc_trust_lvl_2 = enhc_trust_lvl_1 + 1
            else:
                enhc_trust_lvl_1 = jdata["bias_trust_lvl_1"]
                enhc_trust_lvl_2 = enhc_trust_lvl_1 + 1
            if enhc_trust_lvl_1 > jdata["bias_trust_lvl_1"] * 8:
                enhc_trust_lvl_1 = jdata["bias_trust_lvl_1"]
                enhc_trust_lvl_2 = enhc_trust_lvl_1 + 1
        np.savetxt(walker_path + 'trust_lvl1.dat', [enhc_trust_lvl_1],
                   fmt='%.6f')
        # copy enhc file
        for ii in enhc_files:
            if os.path.exists(walker_path + ii):
                os.remove(walker_path + ii)
            shutil.copy(enhc_path + ii, walker_path)
        # copy graph files
        for ii in graph_files:
            file_name = os.path.basename(ii)
            abs_path = os.path.abspath(ii)
            if os.path.exists(walker_path + file_name):
                os.remove(walker_path + file_name)
            os.symlink(abs_path, walker_path + file_name)
        # config MD
        mol_conf_file = walker_path + "grompp.mdp"
        make_grompp_enhc(mol_conf_file, nsteps, frame_freq)
        # config plumed
        graph_list = ""
        counter = 0
        for ii in graph_files:
            file_name = os.path.basename(ii)
            if counter == 0:
                graph_list = "%s" % file_name
            else:
                graph_list = "%s,%s" % (graph_list, file_name)
            counter = counter + 1
        plm_conf = walker_path + enhc_plm
        replace(plm_conf, "MODEL=[^ ]* ", ("MODEL=%s " % graph_list))
        replace(plm_conf, "TRUST_LVL_1=[^ ]* ",
                ("TRUST_LVL_1=%f " % enhc_trust_lvl_1))
        replace(plm_conf, "TRUST_LVL_2=[^ ]* ",
                ("TRUST_LVL_2=%f " % enhc_trust_lvl_2))
        replace(plm_conf, "STRIDE=[^ ]* ", ("STRIDE=%d " % frame_freq))
        replace(plm_conf, "FILE=[^ ]* ", ("FILE=%s " % enhc_out_plm))
        plm_bf_conf = walker_path + enhc_bf_plm
        replace(plm_bf_conf, "STRIDE=[^ ]* ", ("STRIDE=%d " % frame_freq))
        replace(plm_bf_conf, "FILE=[^ ]* ", ("FILE=%s " % 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))