main_menu_5.py

#!/usr/bin/env python3

import numpy as np
import os, glob, shutil
import subprocess as sp
import json

import dpdata

from utils import prange, create_path, file_to_list, select_logs, process_bar, get_time_fp, str_to_type_map

###########################################################
def fp_task_functions():
    while True:
        print("\n")
        print("========================================================")
        print("          Various functions used in fp task             ")
        print("               Convergence Testing                      ")
        print(" (  1)  K-Point (baed on KSPACING) convergence test     ")
        print(" (  2)  ENCUT Convergence Test                          ")
        print("                                                        ")
        print(" ( 10)  Generate deepmd raw files from AIMD traj        ")
        print(" ( 11)  Generate confs for exploring from AIMD traj     ")
        print(" ( 12)  Calculate time cost for all VASP tasks in specified dir")
        print(" ( 13)  Generate deepmd raw file from Optimization OUTCAR")
        print("                                                        ")
        print("              Delete Redundant files                    ")
        print(" ( 90)  Delete all redundant files in current dirs (contains sub-dirs)")
        print(" (100)  Relabel all data generated by dpgen             ")
        print("\n Tips: Input -10 to return main menu                  ")
        print("========================================================")

        jsel = input("\n Please input the menu index: ")

        if jsel == str(1):
            from sklearn.metrics import mean_squared_error, r2_score
            str_in = input("\n Please input the KSPACING start, end and step, e.g., 0.10, 0.30, 0.05. ")
            str_in_in = input("\n If perform VASP calculations? y(Y) or n(N). If those calculations have been finished, you can choose n. ")
            start, end, step = list(map(float, str_in.split(",")))
            kspace = np.arange(start, end+step, step)
            if str_in_in == "y" or str_in_in == "Y":
                make_kspacing_convergence_test(kspace)

            str_in = input("\n If analyze results? y(Y) or n(N). ")
            if str_in == "Y" or str_in == "y":
                force = {}
                sys_min = dpdata.LabeledSystem(os.path.join("./", "%.6f/OUTCAR"%start), fmt="vasp/OUTCAR")
                force_min = np.array(sys_min["forces"]).flatten().tolist()

                sys_other = select_logs("./", "OUTCAR")
                for sys_tmp in sys_other:
                    test = dpdata.LabeledSystem(sys_tmp, fmt="vasp/OUTCAR")
                    force[sys_tmp] = np.array(test["forces"]).flatten().tolist()
                    print(" KSPACING: %s, Force RMSE: %e"%(sys_tmp, mean_squared_error(force_min, force[sys_tmp])))                

        elif jsel == str(2):
            str_in = input("\n Please input the ENCUT start, end and step, e.g., 400, 600, 50. ")
            start, end, step = list(map(float, str_in.split(",")))
            encut = np.arange(start, end+step, step)
            make_encut_convergence_test(encut)
        elif jsel == str(10):
            generate_deepmd_raw_files_from_AIMD_traj()
        elif jsel == str(11):
            generate_confs_for_exploring_from_AIMD_traj()
        elif jsel == str(12):
            str_in =input("\n Please input the root dir name: ") 
            calculate_time_cost_for_all_VASP_tasks_in_specified_dir(root_dir=str_in)
        elif jsel == str(13):
            str_in = input("\n Please input the root dir name: ")
            out_dir = input("\n Please inout the out dir name: ")
            generate_deepmd_raw_files_from_optimization_OUTCAR(root_dir=str_in, out_dir=out_dir)

        elif jsel ==str(90):
            str_in = input("\n Please input the root dir name: ")
            file_in = input("\n Please input the file name to be deleted, e.g., WAVECAR. ")
            delete_all_redundant_files(root_dir=str_in, file_name=file_in)
        elif jsel == str(100):
            str_in = input("\n Firstly you should use the main menu 1 (sub menu 95) to generate total data set, if finished? Y(y) or N(n). ")
            if str_in == "y" or str_in == "Y":
                root_dir = input("\n Please input the root dir name containing data set: ")
                relabel_all_data_generate_by_dpgen(root_dir)
        elif jsel == str(-10):
            break
        else:
            print(" !!! Unspported menu index, please input again. !!!")
###########################################################


###########################################################
def update_kspacing_incar(KSPACING: float) -> None:
    lines = file_to_list("INCAR_template")

    for line in lines:
        if "KSPACING" in line:
            raise ValueError(" The INCAR_template file contains KSPACING info, please delete them.")
    lines.append("KGAMMA   = .FALSE.")
    lines.append("KSPACING = %f"%KSPACING)

    os.remove("INCAR_template")
    with open("INCAR", "w") as f:
        f.writelines(['\n'.join(lines), '\n'])

def updata_encut_incar(encut: int) -> None:
    lines = file_to_list("INCAR_template")

    for line in lines:
        if "ENCUT" in line:
            raise ValueError(" The INCAR_template file contains ENCUT info, please delete them.")
    lines.append("ENCUT = %d"%encut)

    os.remove("INCAR_template")
    with open("INCAR", "w") as f:
        f.writelines(['\n'.join(lines), '\n'])

def make_kspacing_convergence_test(kspace: list) -> None:
    if not os.path.exists("POSCAR"):
        raise FileNotFoundError(" Cannot find POSCAR file, program exit...")
    if not os.path.exists("POTCAR"):
        raise FileNotFoundError(" Cannot find POTCAR file, program exit...")
    if not os.path.exists("INCAR_template"):
        raise FileNotFoundError(" Cannot find INCAR_template file, program exit...")

    cwd = os.getcwd()
    for kpt in kspace:
        kpt_dir = "dpgen_work/k_testing/%.6f"%kpt
        create_path(kpt_dir)
        shutil.copy2("POSCAR", kpt_dir+"/POSCAR")
        shutil.copy2("POTCAR", kpt_dir+"/POTCAR")
        shutil.copy2("INCAR_template", kpt_dir+"/INCAR_template")
        os.chdir(kpt_dir)
        update_kspacing_incar(kpt)
        os.chdir(cwd)

def make_encut_convergence_test(encut: list) -> None:
    if not os.path.exists("POSCAR"):
        raise FileNotFoundError(" Cannot find POSCAR file, program exit...")
    if not os.path.exists("POTCAR"):
        raise FileNotFoundError(" Cannot find POTCAR file, program exit...")
    if not os.path.exists("INCAR_template"):
        raise FileNotFoundError(" Cannot find INCAR_template file, program exit...")

    cwd = os.getcwd()
    for ect in encut:
        ect_dir = "dpgen_work/encut_testing/%d"%ect
        create_path(ect_dir)
        shutil.copy2("POSCAR", ect_dir+"/POSCAR")
        shutil.copy2("POTCAR", ect_dir+"/POTCAR")
        shutil.copy2("INCAR_template", ect_dir+"/INCAR_template")
        os.chdir(ect_dir)
        updata_encut_incar(ect)
        os.chdir(cwd)
###########################################################


###########################################################
def relabel_all_data_generate_by_dpgen(root_dir):
    import dpdata

    all_system_path = glob.glob(os.path.join(root_dir, "*"))
    all_system_path.sort()

    system = []
    for system_path in all_system_path:
        system.append(os.path.basename(system_path))
    system = set(system)

    for sys_tmp in system:      
        labeled_system = dpdata.LabeledSystem(os.path.join(root_dir, sys_tmp), fmt="deepmd/npy")
        for iframe in range(labeled_system.get_nframes()):
            path_name ="Relabeled_Data/%s/dpgen_work/%06d"%(sys_tmp, iframe) 
            create_path(path_name)
            labeled_system.sub_system([iframe,]).to_vasp_poscar(path_name+"/POSCAR")
            shutil.copy2("INCAR_%s"%sys_tmp, path_name+"/INCAR")
            shutil.copy2("POTCAR", path_name+"/POTCAR")
###########################################################


###########################################################
def generate_deepmd_raw_files_from_optimization_OUTCAR(root_dir, out_dir):
    """
    generate deepmd raw file from all optimization OUTCAR file in root_dir.

    Written at 20220116
    """
    import dpdata
    from dpdata  import MultiSystems
    all_outcar = select_logs(root_dir, "OUTCAR")
    if len(all_outcar) == 0:
        raise ValueError(" Cannot find a OUTCAR in %s, please check again..."%root_dir)
    else:
        print(" Total OUTCAR numbers: %06d"%len(all_outcar))
    #final_sys = MultiSystems()
    final_sys = None
    counter = 1
    for outcar in all_outcar:
        system = dpdata.LabeledSystem(outcar, fmt="vasp/outcar")
        if final_sys is None:
            final_sys = system[-1]
        else:
            final_sys.append(system[-1])
        process_bar(counter, len(all_outcar), " Reading OUTCAR...")
        counter += 1
    
    final_sys.to_deepmd_raw(out_dir)
    final_sys.to_deepmd_npy(out_dir)
        
def generate_deepmd_raw_files_from_AIMD_traj():
    import dpdata
    str_in = input("\n Please input the software that you use, vasp, qe, or cp2k? ")
    str_in_in = input("\n Please input the type map, e.g., Li, Si. ")
    if str_in.lower() == "vasp": # OUTCAR file
        labeled_system = dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar", type_map=str_to_type_map(str_in_in))
        #labeled_system = dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar")
    print(" System Info: ")
    print(" Total frames: %d"%labeled_system.get_nframes())
    str_in = input("\n Please input the frame start, end and step, e.g., 0, 100, 2. ")
    start, end, step = list(map(int, str_in.split(",")))
    frame = range(start, end, step)
    labeled_system.sub_system(frame).to("deepmd/raw", "deepmd")
    labeled_system.sub_system(frame).to("deepmd/npy", "deepmd")

def generate_confs_for_exploring_from_AIMD_traj():
    import dpdata
    str_in = input("\n Please input the software that you use, vasp, qe, or cp2k? ")
    if str_in.lower() == "vasp": # OUTCAR file
        labeled_system = dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar")
    print(" System Info: ")
    print(" Total frames: %d"%labeled_system.get_nframes())
    str_in = input("\n Please input the frame start, end and step, e.g., 0, 1000, 100. ")
    start, end, step = list(map(int, str_in.split(",")))
    frame = range(start, end, step)
    for i in range(len(frame)):
        idx = frame[i]
        n = str("%06d" % i)
        sp.call('mkdir -p configurations/' + n, shell=True)
        labeled_system.to('vasp/poscar', 'POSCAR', frame_idx=idx)
        sp.call('mv POSCAR configurations/' + n, shell=True)
###########################################################


###########################################################
def delete_all_redundant_files(root_dir, file_name):
    all_files = select_logs(root_dir, file_name)

    if len(all_files) == 0:
        raise ValueError("Cannot find the %s file in %s dir, please check again..."%(file_name, root_dir))
    else:
        print("\n Find file numbers: %d"%len(all_files))
    counter = 1
    for file_tmp in all_files:
        os.remove(file_tmp)
        process_bar(counter, len(all_files), " Deleting files...")
        counter += 1
###########################################################


###########################################################
def calculate_time_cost_for_all_VASP_tasks_in_specified_dir(root_dir):
    all_files = select_logs(root_dir, "OUTCAR")
    if len(all_files) == 0:
        print(" Cannot find the OUTCAR file in %s, please check again."%root_dir)
    else:
        for file_tmp in all_files:
            print(" %s, %.2f h."%(file_tmp, get_time_fp(file_tmp)))
###########################################################