Merge pull request #63 from deepmodeling/abacus

update develop branch so that it is in sync with abacus branch

.gitignore

@@ -52,4 +52,7 @@ env/
 venv/
 ENV/
 env.bak/
-venv.bak/
+venv.bak/
+
+# examples
+ABACUS/

.readthedocs.yaml

@@ -0,0 +1,18 @@
+# .readthedocs.yaml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+formats: all
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+   configuration: docs/source/conf.py
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+   version: 3.8
+   install:
+   - requirements: ./requirements.txt

README.rst

@@ -0,0 +1,9 @@
+Template for the Read the Docs tutorial
+=======================================
+
+This GitHub template includes fictional Python library
+with some basic Sphinx docs.
+
+Read the tutorial here:
+
+https://docs.readthedocs.io/en/stable/tutorial/

deepks/iterate/__init__.py

@@ -1,6 +1,9 @@
 __all__ = [
     "iterate",
     "template",
+    "template_abacus",
+    "generator_abacus"
 ]
 
-from .iterate import make_scf, make_train, make_iterate
+from .iterate import make_scf, make_train, make_iterate
+from .iterate import make_scf_abacus    #caoyu add 2021-07-25

deepks/iterate/generator_abacus.py

@@ -0,0 +1,144 @@
+#These 3 functions are copied from dpgen to generate ABACUS INPUT , KPT and STRU file.
+
+bohr2ang = 0.52917721067
+
+def make_abacus_scf_kpt(fp_params):
+    # Make KPT file for abacus pw scf calculation.
+    # KPT file is the file containing k points infomation in ABACUS scf calculation.
+    k_points = [1, 1, 1, 0, 0, 0]
+    if "k_points" in fp_params:
+        k_points = fp_params["k_points"]
+        if len(k_points) != 6:
+            raise RuntimeError("k_points has to be a list containig 6 integers specifying MP k points generation.")
+    ret = "K_POINTS\n0\nGamma\n"
+    for i in range(6):
+        ret += str(k_points[i]) + " "
+    return ret
+
+def make_abacus_scf_input(fp_params):
+    # Make INPUT file for abacus pw scf calculation.
+    ret = "INPUT_PARAMETERS\n"
+    ret += "calculation scf\n"
+    assert(fp_params['ntype'] >= 0 and type(fp_params["ntype"]) == int),  "'ntype' should be a positive integer."
+    ret += "ntype %d\n" % fp_params['ntype']
+    #ret += "pseudo_dir ./\n"
+    if "ecutwfc" in fp_params:
+        assert(fp_params["ecutwfc"] >= 0) ,  "'ntype' should be non-negative."
+        ret += "ecutwfc %f\n" % fp_params["ecutwfc"]
+    if "scf_thr" in fp_params:
+        ret += "scf_thr %e\n" % fp_params["scf_thr"]
+    if "scf_nmax" in fp_params:
+        assert(fp_params['scf_nmax'] >= 0 and type(fp_params["scf_nmax"])== int), "'scf_nmax' should be a positive integer."
+        ret += "scf_nmax %d\n" % fp_params["scf_nmax"]    
+    if "basis_type" in fp_params:
+        assert(fp_params["basis_type"] in ["pw", "lcao", "lcao_in_pw"]) , "'basis_type' must in 'pw', 'lcao' or 'lcao_in_pw'."
+        ret+= "basis_type %s\n" % fp_params["basis_type"]
+    if "dft_functional" in fp_params:
+        ret += "dft_functional %s\n" % fp_params["dft_functional"]
+    if "gamma_only" in fp_params:
+        assert(fp_params["gamma_only"] ==0 or fp_params["gamma_only"] ==1 ) , "'gamma_only' should be 0 or 1."
+        ret+= "gamma_only %d\n" % fp_params["gamma_only"]  
+    if "mixing_type" in fp_params:
+        assert(fp_params["mixing_type"] in ["plain", "kerker", "pulay", "pulay-kerker", "broyden"])
+        ret += "mixing_type %s\n" % fp_params["mixing_type"]
+    if "mixing_beta" in fp_params:
+        assert(fp_params["mixing_beta"] >= 0 and fp_params["mixing_beta"] < 1), "'mixing_beta' should between 0 and 1."
+        ret += "mixing_beta %f\n" % fp_params["mixing_beta"]
+    if "symmetry" in fp_params:
+        #assert(fp_params["symmetry"] == 0 or fp_params["symmetry"] == 1), "'symmetry' should be either 0 or 1."
+        ret += "symmetry %d\n" % fp_params["symmetry"]
+    if "nbands" in fp_params:
+        if(type(fp_params["nbands"]) == int and fp_params["nbands"] > 0):
+            ret += "nbands %d\n" % fp_params["nbands"]
+        else:
+            print("warnning: Parameter [nbands] given is not a positive integer, the default value of [nbands] in ABACUS will be used. ")
+    if "nspin" in fp_params:
+        assert(fp_params["nspin"] == 1 or fp_params["nspin"] == 2 or fp_params["nspin"] == 4), "'nspin' can anly take 1, 2 or 4"
+        ret += "nspin %d\n" % fp_params["nspin"]
+    if "ks_solver" in fp_params:
+        assert(fp_params["ks_solver"] in ["cg", "dav", "lapack", "genelpa", "hpseps", "scalapack_gvx"]), "'ks_sover' should in 'cgx', 'dav', 'lapack', 'genelpa', 'hpseps', 'scalapack_gvx'."
+        ret += "ks_solver %s\n" % fp_params["ks_solver"]
+    if "smearing_method" in fp_params:
+        assert(fp_params["smearing_method"] in ["gaussian", "fd", "fixed", "mp", "mp2", "mv"]), "'smearing' should in 'gaussian', 'fd', 'fixed', 'mp', 'mp2', 'mv'. "
+        ret += "smearing_method %s\n" % fp_params["smearing_method"]
+    if "smearing_sigma" in fp_params:
+        assert(fp_params["smearing_sigma"] >= 0), "'smearing_sigma' should be non-negative."
+        ret += "smearing_sigma %f\n" % fp_params["smearing_sigma"]
+    if (("kspacing" in fp_params) and (fp_params["k_points"] is None) and (fp_params["gamma_only"] == 0)):
+        assert(fp_params["kspacing"] > 0), "'kspacing' should be positive."
+        ret += "kspacing %f\n" % fp_params["kspacing"]
+    if "cal_force" in fp_params:
+        assert(fp_params["cal_force"] == 0  or fp_params["cal_force"] == 1), "'cal_force' should be either 0 or 1."
+        ret += "cal_force %d\n" % fp_params["cal_force"]
+    if "cal_stress" in fp_params:
+        assert(fp_params["cal_stress"] == 0  or fp_params["cal_stress"] == 1), "'cal_stress' should be either 0 or 1."
+        ret += "cal_stress %d\n" % fp_params["cal_stress"]    
+    #paras for deepks
+    if "deepks_out_labels" in fp_params:
+        assert(fp_params["deepks_out_labels"] == 0 or fp_params["deepks_out_labels"] == 1), "'deepks_out_labels' should be either 0 or 1."
+        ret += "deepks_out_labels %d\n" % fp_params["deepks_out_labels"]
+    if "deepks_scf" in fp_params:
+        assert(fp_params["deepks_scf"] == 0  or fp_params["deepks_scf"] == 1), "'deepks_scf' should be either 0 or 1."
+        ret += "deepks_scf %d\n" % fp_params["deepks_scf"]
+    if "deepks_bandgap" in fp_params:
+        assert(fp_params["deepks_bandgap"] == 0  or fp_params["deepks_bandgap"] == 1), "'deepks_scf' should be either 0 or 1."
+        ret += "deepks_bandgap %d\n" % fp_params["deepks_bandgap"]
+    if "model_file" in fp_params:
+        ret += "deepks_model %s\n" % fp_params["model_file"]
+    if fp_params["dft_functional"] == "hse":
+        ret += "exx_pca_threshold 1e-4\n"
+        ret += "exx_c_threshold 1e-4\n"
+        ret += "exx_dm_threshold 1e-4\n"
+        ret += "exx_schwarz_threshold 1e-5\n"
+        ret += "exx_cauchy_threshold 1e-7\n"
+        ret += "exx_ccp_rmesh_times 1\n"
+    return ret
+
+def make_abacus_scf_stru(sys_data, fp_pp_files, fp_params):
+    atom_names = sys_data['atom_names']
+    atom_numbs = sys_data['atom_numbs']
+    assert(len(atom_names) == len(fp_pp_files)), "the number of pp_files must be equal to the number of atom types. "
+    assert(len(atom_names) == len(atom_numbs)), "Please check the name of atoms. "
+    cell = sys_data["cells"][0].reshape([3, 3])
+    if "lattice_vector" in fp_params:
+        cell = fp_params["lattice_vector"]
+    coord = sys_data['coords'][0]
+    #volume = np.linalg.det(cell)
+    #lattice_const = np.power(volume, 1/3)
+    ret = "ATOMIC_SPECIES\n"
+    for iatom in range(len(atom_names)):
+        ret += atom_names[iatom] + " 1.00 " + fp_pp_files[iatom] + "\n"
+    ret += "\n"
+    if "lattice_constant" in fp_params:
+        ret += "\nLATTICE_CONSTANT\n"
+        ret += str(fp_params["lattice_constant"]) + "\n\n" # in Bohr, in this way coord and cell are in Angstrom 
+    else:
+        ret += "\nLATTICE_CONSTANT\n"
+        ret += str(1/bohr2ang) + "\n\n"
+    ret += "LATTICE_VECTORS\n"
+    for ix in range(3):
+        for iy in range(3):
+            ret += str(cell[ix][iy]) + " "
+        ret += "\n"
+    ret += "\n"
+    ret += "ATOMIC_POSITIONS\n"
+    ret += fp_params["coord_type"]
+    ret += "\n\n"
+    natom_tot = 0
+    for iele in range(len(atom_names)):
+        ret += atom_names[iele] + "\n"
+        ret += "0.0\n"
+        ret += str(atom_numbs[iele]) + "\n"
+        for iatom in range(atom_numbs[iele]):
+            ret += "%.12f %.12f %.12f %d %d %d\n" % (coord[natom_tot, 0], coord[natom_tot, 1], coord[natom_tot, 2], 0, 0, 0)
+            natom_tot += 1
+    assert(natom_tot == sum(atom_numbs))
+    if "basis_type" in fp_params and fp_params["basis_type"]=="lcao":
+        ret +="\nNUMERICAL_ORBITAL\n"
+        assert(len(fp_params["orb_files"])==len(atom_names))
+        for iatom in range(len(atom_names)):
+            ret += fp_params["orb_files"][iatom] +"\n"
+    if "deepks_scf" in fp_params and fp_params["deepks_out_labels"]==1:
+        ret +="\nNUMERICAL_DESCRIPTOR\n"
+        ret +=fp_params["proj_file"][0]+"\n"
+    return ret

deepks/iterate/iterate.py

@@ -11,6 +11,8 @@
 from deepks.utils import load_basis, save_basis
 from deepks.task.workflow import Sequence, Iteration
 from deepks.iterate.template import make_scf, make_train
+from deepks.iterate.template_abacus import make_scf_abacus  #caoyu add 2021-07-22 
+from deepks.iterate.template_abacus import DEFAULT_SCF_ARGS_ABACUS
 
 
 # args not specified here may cause error
@@ -21,17 +23,23 @@
     "ingroup_parallel": 1, #how many tasks can run at same time in one job
     "dispatcher": None, # use default lazy-local slurm defined in task.py
     "resources": None, # use default 10 core defined in templete.py
-    "python": "python" # use current python in path
+    "python": "python", # use current python in path
+    "dpdispatcher_machine": None, # (only need for dispatcher=="dpdispatcher")
+    "dpdispatcher_resources": None # (only need for dispatcher=="dpdispatcher")
 }
 
 # args not specified here may cause error
 DEFAULT_TRN_MACHINE = {
     "dispatcher": None, # use default lazy-local slurm defined in task.py
     "resources": None, # use default 10 core defined in templete.py
-    "python": "python" # use current python in path
+    "python": "python", # use current python in path
+    "dpdispatcher_machine": None, # (only need for dispatcher=="dpdispatcher")
+    "dpdispatcher_resources": None # (only need for dispatcher=="dpdispatcher")
 }
 
 SCF_ARGS_NAME = "scf_input.yaml"
+SCF_ARGS_NAME_ABACUS="scf_abacus.yaml"   #for abacus, caoyu add 2021-07-26
+INIT_SCF_NAME_ABACUS="init_scf_abacus.yaml"   #for abacus init, caoyu add 2021-12-17
 TRN_ARGS_NAME = "train_input.yaml"
 INIT_SCF_NAME = "init_scf.yaml"
 INIT_TRN_NAME = "init_train.yaml"
@@ -136,7 +144,8 @@ def make_iterate(systems_train=None, systems_test=None, n_iter=0,
                  train_input=True, train_machine=None,
                  init_model=False, init_scf=True, init_train=True,
                  init_scf_machine=None, init_train_machine=None,
-                 cleanup=False, strict=True):
+                 cleanup=False, strict=True, 
+                 use_abacus=False, scf_abacus=None, init_scf_abacus=None):#caoyu add 2021-07-22
     r"""
     Make a `Workflow` to do the iterative training procedure.
 
@@ -210,7 +219,10 @@ def make_iterate(systems_train=None, systems_test=None, n_iter=0,
     strict: bool, optional 
         Whether to allow additional arguments to be passed to task constructor,
         through `scf_machine` and `train_machine`. Defaults to True.
-
+    use_abacus: bool, optional
+        If set to`True`,  do SCF calculation with ABACUS.
+        If set to `False`,  do SCF calculation with PySCF.
+        Defaults to `False`.
     Returns
     -------
     iterate: Iteration (subclass of Workflow)
@@ -243,18 +255,29 @@ def make_iterate(systems_train=None, systems_test=None, n_iter=0,
     # check required machine parameters
     scf_machine = check_arg_dict(scf_machine, DEFAULT_SCF_MACHINE, strict)
     train_machine = check_arg_dict(train_machine, DEFAULT_TRN_MACHINE, strict)
-    # handle projection basis
-    if proj_basis is not None:
-        save_basis(os.path.join(share_folder, PROJ_BASIS), load_basis(proj_basis))
-        proj_basis = PROJ_BASIS
+
     # make tasks
-    scf_step = make_scf(
-        systems_train=systems_train, systems_test=systems_test,
-        train_dump=DATA_TRAIN, test_dump=DATA_TEST, no_model=False,
-        workdir=SCF_STEP_DIR, share_folder=share_folder,
-        source_arg=scf_args_name, source_model=MODEL_FILE,
-        source_pbasis=proj_basis, cleanup=cleanup, **scf_machine
-    )
+    if use_abacus:  #caoyu add 2021-07-22
+        scf_abacus_name = check_share_folder(scf_abacus, SCF_ARGS_NAME_ABACUS, share_folder)
+        scf_abacus = check_arg_dict(scf_abacus, DEFAULT_SCF_ARGS_ABACUS, strict)
+        scf_abacus = dict(scf_abacus, **scf_machine)
+        scf_step = make_scf_abacus(systems_train=systems_train, systems_test=systems_test,
+            train_dump=DATA_TRAIN, test_dump=DATA_TEST, no_model=False,
+            model_file=MODEL_FILE, workdir=SCF_STEP_DIR, share_folder=share_folder,
+            cleanup=cleanup, **scf_abacus)
+        proj_basis=None     # discussion needed
+    else:
+        # handle projection basis
+        if proj_basis is not None:
+            save_basis(os.path.join(share_folder, PROJ_BASIS), load_basis(proj_basis))
+            proj_basis = PROJ_BASIS
+        scf_step = make_scf(
+            systems_train=systems_train, systems_test=systems_test,
+            train_dump=DATA_TRAIN, test_dump=DATA_TEST, no_model=False,
+            workdir=SCF_STEP_DIR, share_folder=share_folder,
+            source_arg=scf_args_name, source_model=MODEL_FILE,
+            source_pbasis=proj_basis, cleanup=cleanup, **scf_machine
+        )
     train_step = make_train(
         source_train=DATA_TRAIN, source_test=DATA_TEST,
         restart=True, source_model=MODEL_FILE, save_model=MODEL_FILE, 
@@ -265,25 +288,37 @@ def make_iterate(systems_train=None, systems_test=None, n_iter=0,
     per_iter = Sequence([scf_step, train_step])
     iterate = Iteration(per_iter, n_iter, 
                         workdir=".", record_file=os.path.join(workdir, RECORD))
+
     # make init
     if init_model: # if set true or give str, check share/init/model.pth
         init_folder=os.path.join(share_folder, "init")
         check_share_folder(init_model, MODEL_FILE, init_folder)
         iterate.set_init_folder(init_folder)
     elif init_scf or init_train: # otherwise, make an init iteration to train the first model
-        init_scf_name = check_share_folder(init_scf, INIT_SCF_NAME, share_folder)
-        init_train_name = check_share_folder(init_train, INIT_TRN_NAME, share_folder)
         init_scf_machine = (check_arg_dict(init_scf_machine, DEFAULT_SCF_MACHINE, strict)
             if init_scf_machine is not None else scf_machine)
+        if use_abacus:  #caoyu add 2021-07-22
+            init_scf_abacus_name = check_share_folder(init_scf_abacus, INIT_SCF_NAME_ABACUS, share_folder)
+            init_scf_abacus = check_arg_dict(init_scf_abacus, DEFAULT_SCF_ARGS_ABACUS, strict)
+            init_scf_abacus = dict(init_scf_abacus, **scf_machine)
+            scf_init = make_scf_abacus(
+                systems_train=systems_train, systems_test=systems_test,
+                train_dump=DATA_TRAIN, test_dump=DATA_TEST, no_model=True,
+                workdir=SCF_STEP_DIR, share_folder=share_folder, model_file=None, 
+                cleanup=cleanup, **init_scf_abacus
+            )
+        else:
+            init_scf_name = check_share_folder(init_scf, INIT_SCF_NAME, share_folder)
+            scf_init = make_scf(
+                systems_train=systems_train, systems_test=systems_test,
+                train_dump=DATA_TRAIN, test_dump=DATA_TEST, no_model=True,
+                workdir=SCF_STEP_DIR, share_folder=share_folder,
+                source_arg=init_scf_name, source_model=None, source_pbasis=proj_basis,
+                cleanup=cleanup, **scf_machine
+            )
+        init_train_name = check_share_folder(init_train, INIT_TRN_NAME, share_folder)
         init_train_machine = (check_arg_dict(init_train_machine, DEFAULT_SCF_MACHINE, strict)
             if init_train_machine is not None else train_machine)
-        scf_init = make_scf(
-        systems_train=systems_train, systems_test=systems_test,
-        train_dump=DATA_TRAIN, test_dump=DATA_TEST, no_model=True,
-        workdir=SCF_STEP_DIR, share_folder=share_folder,
-        source_arg=init_scf_name, source_model=None, source_pbasis=proj_basis,
-        cleanup=cleanup, **scf_machine
-        )
         train_init = make_train(
             source_train=DATA_TRAIN, source_test=DATA_TEST,
             restart=False, source_model=MODEL_FILE, save_model=MODEL_FILE,