Skip to content

Commit 5cb6688

Browse files
ErjieWuErjieWu
authored
Merge pull request #10 from ErjieWu/refactor
Refactor: change places of default value and add annotations.
2 parents b6733d1 + e7d2d7b commit 5cb6688

23 files changed

Lines changed: 412 additions & 472 deletions

‎deepks/default.py‎

Lines changed: 82 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,93 @@
11
import os
22
import numpy as np
3+
import torch
34

4-
## from utils.py
5-
QCDIR = os.path.dirname(os.path.realpath(__file__))
5+
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
6+
BOHR2ANG = 0.52917721067
7+
NAME_TYPE = {'X': 0, # Ghost
8+
'H': 1, 'He': 2, 'Li': 3, 'Be': 4, 'B': 5, 'C': 6, 'N': 7,
9+
'O': 8, 'F': 9, 'Ne': 10, 'Na': 11, 'Mg': 12, 'Al': 13,
10+
'Si': 14, 'P': 15, 'S': 16, 'Cl': 17, 'Ar': 18, 'K': 19,
11+
'Ca': 20, 'Sc': 21, 'Ti': 22, 'V': 23, 'Cr': 24, 'Mn': 25,
12+
'Fe': 26, 'Co': 27, 'Ni': 28, 'Cu': 29, 'Zn': 30, 'Ga': 31,
13+
'Ge': 32, 'As': 33, 'Se': 34, 'Br': 35, 'Kr': 36, 'Rb': 37,
14+
'Sr': 38, 'Y': 39, 'Zr': 40, 'Nb': 41, 'Mo': 42, 'Tc': 43,
15+
'Ru': 44, 'Rh': 45, 'Pd': 46, 'Ag': 47, 'Cd': 48, 'In': 49,
16+
'Sn': 50, 'Sb': 51, 'Te': 52, 'I': 53, 'Xe': 54, 'Cs': 55,
17+
'Ba': 56, 'La': 57, 'Ce': 58, 'Pr': 59, 'Nd': 60, 'Pm': 61,
18+
'Sm': 62, 'Eu': 63, 'Gd': 64, 'Tb': 65, 'Dy': 66, 'Ho': 67,
19+
'Er': 68, 'Tm': 69, 'Yb': 70, 'Lu': 71, 'Hf': 72, 'Ta': 73,
20+
'W': 74, 'Re': 75, 'Os': 76, 'Ir': 77, 'Pt': 78, 'Au': 79,
21+
'Hg': 80, 'Tl': 81, 'Pb': 82, 'Bi': 83, 'Po': 84, 'At': 85,
22+
'Rn': 86, 'Fr': 87, 'Ra': 88, 'Ac': 89, 'Th': 90, 'Pa': 91,
23+
'U': 92, 'Np': 93, 'Pu': 94, 'Am': 95, 'Cm': 96, 'Bk': 97,
24+
'Cf': 98, 'Es': 99, 'Fm': 100, 'Md': 101, 'No': 102, 'Lr': 103,
25+
'Rf': 104, 'Db': 105, 'Sg': 106, 'Bh': 107, 'Hs': 108,
26+
'Mt': 109, 'Ds': 110, 'Rg': 111, 'Cn': 112, 'Uut': 113,
27+
'Fl': 114, 'Uup': 115, 'Lv': 116, 'Uus': 117, 'Uuo': 118
28+
} #dict
29+
TYPE_NAME = {v:k for k, v in NAME_TYPE.items()}
30+
ELEMENTS = list(NAME_TYPE.keys())
31+
32+
# Default settings for abacus
33+
DEFAULT_SCF_ARGS_ABACUS = {
34+
# for STRU
35+
"orb_files": ["orb"],
36+
"pp_files": ["upf"],
37+
"proj_file": ["orb"],
38+
"lattice_constant": 1,
39+
"lattice_vector": np.eye(3,dtype=int),
40+
"coord_type": "Cartesian",
41+
# for INPUT
42+
"nspin": 1,
43+
"symmetry": 0,
44+
"nbands": None,
45+
"ecutwfc": 50,
46+
"scf_thr": 1e-7,
47+
"scf_nmax": 50,
48+
"dft_functional": "pbe",
49+
"basis_type": "lcao",
50+
"gamma_only": 1,
51+
"k_points": None, # also used for KPT file if not None
52+
"kspacing": None,
53+
"smearing_method":"gaussian",
54+
"smearing_sigma": 0.02,
55+
"mixing_type": "pulay",
56+
"mixing_beta": 0.4,
57+
"cal_force": 0,
58+
"cal_stress": 0,
59+
"deepks_bandgap": 0,
60+
"deepks_v_delta": 0,
61+
"deepks_out_labels": 1,
62+
"deepks_scf": 0,
63+
"out_wfc_lcao": 0,
64+
# for running
65+
"run_cmd": "mpirun",
66+
"sub_size": 1,
67+
"abacus_path": "/usr/local/bin/ABACUS.mpi",
68+
}
69+
70+
## Default settings for pyscf
671
_zeta = 1.5**np.array([17,13,10,7,5,3,2,1,0,-1,-2,-3])
772
_coef = np.diag(np.ones(_zeta.size)) - np.diag(np.ones(_zeta.size-1), k=1)
873
_table = np.concatenate([_zeta.reshape(-1,1), _coef], axis=1)
974
DEFAULT_BASIS = [[0, *_table.tolist()], [1, *_table.tolist()], [2, *_table.tolist()]]
1075
DEFAULT_SYMB = "Ne"
76+
DEFAULT_UNIT = "Bohr"
77+
DEFAULT_FNAMES = {"e_tot", "e_base", "dm_eig", "conv"}
78+
DEFAULT_HF_ARGS = {
79+
"conv_tol": 1e-9
80+
}
81+
DEFAULT_SCF_ARGS = {
82+
"conv_tol": 1e-7,
83+
# "level_shift": 0.1,
84+
# "diis_space": 20
85+
}
86+
MOL_ATTRIBUTE = {"charge", "basis", "unit"} # other molecule properties
1187

12-
## from iterate/template.py
13-
SCF_CMD = " ".join([
14-
"{python} -u",
15-
"-m deepks.scf.run"
16-
# os.path.join(QCDIR, "scf/run.py") # this is the backup choice
17-
])
18-
19-
TRN_CMD = " ".join([
20-
"{python} -u",
21-
"-m deepks.model.train"
22-
# os.path.join(QCDIR, "train/train.py") # this is the backup choice
23-
])
24-
25-
88+
## default settings for make scf tasks
89+
SCF_CMD = "{python} -u -m deepks.scf.run"
90+
TRN_CMD = "{python} -u -m deepks.model.train"
2691
DEFAULT_SCF_RES = {
2792
"time_limit": "24:00:00",
2893
"cpus_per_task": 8,
@@ -31,21 +96,18 @@
3196
"PYSCF_MAX_MEMORY": 8000
3297
}
3398
}
34-
3599
DEFAULT_SCF_SUB_RES = {
36100
"numb_node": 1,
37101
"task_per_node": 1,
38102
"cpus_per_task": 8,
39103
"exclusive": True
40104
}
41-
42105
DEFAULT_TRN_RES = {
43106
"time_limit": "24:00:00",
44107
"cpus_per_task": 8,
45108
# "numb_gpu": 1, # do not use gpu by default
46109
"mem_limit": 8
47110
}
48-
49111
DEFAULT_DPDISPATCHER_RES = {
50112
"number_node": 1,
51113
"cpu_per_node": 8,
@@ -83,7 +145,6 @@
83145

84146
DATA_TRAIN = "data_train"
85147
DATA_TEST = "data_test"
86-
MODEL_FILE = "model.pth"
87148
PROJ_BASIS = "proj_basis.npz"
88149

89150
SCF_STEP_DIR = "00.scf"
@@ -96,104 +157,5 @@
96157
DEFAULT_TRAIN = "systems_train.raw"
97158
DEFAULT_TEST = "systems_test.raw"
98159

99-
100-
## from iterate/template_abacus.py
101160
MODEL_FILE = "model.pth"
102-
CMODEL_FILE = "model.ptg"
103-
104-
NAME_TYPE = { 'H': 1, 'He': 2, 'Li': 3, 'Be': 4, 'B': 5, 'C': 6, 'N': 7,
105-
'O': 8, 'F': 9, 'Ne': 10, 'Na': 11, 'Mg': 12, 'Al': 13,
106-
'Si': 14, 'P': 15, 'S': 16, 'Cl': 17, 'Ar': 18, 'K': 19,
107-
'Ca': 20, 'Sc': 21, 'Ti': 22, 'V': 23, 'Cr': 24, 'Mn': 25,
108-
'Fe': 26, 'Co': 27, 'Ni': 28, 'Cu': 29, 'Zn': 30, 'Ga': 31,
109-
'Ge': 32, 'As': 33, 'Se': 34, 'Br': 35, 'Kr': 36, 'Rb': 37,
110-
'Sr': 38, 'Y': 39, 'Zr': 40, 'Nb': 41, 'Mo': 42, 'Tc': 43,
111-
'Ru': 44, 'Rh': 45, 'Pd': 46, 'Ag': 47, 'Cd': 48, 'In': 49,
112-
'Sn': 50, 'Sb': 51, 'Te': 52, 'I': 53, 'Xe': 54, 'Cs': 55,
113-
'Ba': 56,
114-
## La-Lu
115-
## 'La': 57, 'Ce': 58, 'Pr': 59, 'Nd': 60, 'Pm': 61,
116-
## 'Sm': 62, 'Eu': 63, 'Gd': 64, 'Tb': 65, 'Dy': 66, 'Ho': 67,
117-
## 'Er': 68, 'Tm': 69, 'Yb': 70, 'Lu': 71,
118-
'Hf': 72, 'Ta': 73,
119-
'W': 74, 'Re': 75, 'Os': 76, 'Ir': 77, 'Pt': 78, 'Au': 79,
120-
'Hg': 80, 'Tl': 81, 'Pb': 82, 'Bi': 83,
121-
## Here after are radioactive elements
122-
## 'Po': 84, 'At': 85, 'Rn': 86, 'Fr': 87, 'Ra': 88,
123-
## Ac-Lr
124-
## 'Ac': 89, 'Th': 90, 'Pa': 91,
125-
## 'U': 92, 'Np': 93, 'Pu': 94, 'Am': 95, 'Cm': 96, 'Bk': 97,
126-
## 'Cf': 98, 'Es': 99, 'Fm': 100, 'Md': 101, 'No': 102, 'Lr': 103,
127-
## 'Rf': 104, 'Db': 105, 'Sg': 106, 'Bh': 107, 'Hs': 108,
128-
## 'Mt': 109, 'Ds': 110, 'Rg': 111, 'Cn': 112, 'Uut': 113,
129-
## 'Fl': 114, 'Uup': 115, 'Lv': 116, 'Uus': 117, 'Uuo': 118
130-
} #dict
131-
TYPE_NAME ={v:k for k, v in NAME_TYPE.items()}
132-
TYPE_INDEX = {k:v for k, v in NAME_TYPE.items()}
133-
134-
ABACUS_CMD="bash run_abacus.sh"
135-
136-
DEFAULT_SCF_ARGS_ABACUS={
137-
"orb_files": ["orb"], #atomic number order
138-
"pp_files": ["upf"], #atomic number order
139-
"proj_file": ["orb"],
140-
"ntype": 1,
141-
"nspin": 1,
142-
"symmetry": 0,
143-
"nbands": None,
144-
"ecutwfc": 50,
145-
"scf_thr": 1e-7,
146-
"scf_nmax": 50,
147-
"dft_functional": "pbe",
148-
"basis_type": "lcao",
149-
"gamma_only": 1,
150-
"k_points": None,
151-
"kspacing": None,
152-
"smearing_method":"gaussian",
153-
"smearing_sigma":0.02,
154-
"mixing_type": "pulay",
155-
"mixing_beta": 0.4,
156-
"cal_force": 0,
157-
"cal_stress": 0,
158-
"deepks_bandgap": 0,
159-
"deepks_v_delta": 0,
160-
"deepks_out_labels":1,
161-
"deepks_scf":0,
162-
"lattice_constant": 1,
163-
"lattice_vector": np.eye(3,dtype=int),
164-
"coord_type": "Cartesian",
165-
"run_cmd": "mpirun",
166-
"sub_size": 1,
167-
"abacus_path": "/usr/local/bin/ABACUS.mpi",
168-
"out_wfc_lcao": 0,
169-
}
170-
171-
172-
## from iterate/generator_abacus.py
173-
BOHR2ANG = 0.52917721067
174-
175-
## from model
176-
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
177-
178-
SCALE_EPS = 1e-8
179-
180-
## from tools
181-
BOHR = 0.52917721092
182-
183-
184-
## from scf
185-
DEFAULT_UNIT = "Bohr"
186-
187-
DEFAULT_FNAMES = {"e_tot", "e_base", "dm_eig", "conv"}
188-
189-
DEFAULT_HF_ARGS = {
190-
"conv_tol": 1e-9
191-
}
192-
193-
DEFAULT_SCF_ARGS = {
194-
"conv_tol": 1e-7,
195-
# "level_shift": 0.1,
196-
# "diis_space": 20
197-
}
198-
199-
MOL_ATTRIBUTE = {"charge", "basis", "unit"} # other molecule properties
161+
CMODEL_FILE = "model.ptg"

‎deepks/iterate/generator_abacus.py‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
1-
#These 3 functions are copied from dpgen to generate ABACUS INPUT , KPT and STRU file.
2-
3-
BOHR2ANG = 0.52917721067
1+
from deepks.default import BOHR2ANG
2+
# These 3 functions are used to generate ABACUS INPUT, KPT and STRU file.
43

54
def make_abacus_scf_kpt(fp_params):
6-
# Make KPT file for abacus pw scf calculation.
7-
# KPT file is the file containing k points infomation in ABACUS scf calculation.
8-
k_points = [1, 1, 1, 0, 0, 0]
5+
'''
6+
Make KPT file for abacus scf calculation.
7+
KPT file is the file containing k points infomation in ABACUS scf calculation.
8+
'''
9+
k_points = [1, 1, 1, 0, 0, 0] # Default k points
910
if "k_points" in fp_params:
1011
k_points = fp_params["k_points"]
1112
if len(k_points) != 6:
@@ -16,14 +17,14 @@ def make_abacus_scf_kpt(fp_params):
1617
return ret
1718

1819
def make_abacus_scf_input(fp_params):
19-
# Make INPUT file for abacus pw scf calculation.
20+
'''
21+
Make INPUT file for abacus scf calculation.
22+
'''
2023
ret = "INPUT_PARAMETERS\n"
2124
ret += "calculation scf\n"
22-
# assert(fp_params['ntype'] >= 0 and type(fp_params["ntype"]) == int), "'ntype' should be a positive integer."
23-
# ret += "ntype %d\n" % fp_params['ntype']
24-
#ret += "pseudo_dir ./\n"
25+
# ret += "pseudo_dir ./\n"
2526
if "ecutwfc" in fp_params:
26-
assert(fp_params["ecutwfc"] >= 0) , "'ntype' should be non-negative."
27+
assert(fp_params["ecutwfc"] >= 0) , "'ecutwfc' should be non-negative."
2728
ret += "ecutwfc %f\n" % fp_params["ecutwfc"]
2829
if "scf_thr" in fp_params:
2930
ret += "scf_thr %e\n" % fp_params["scf_thr"]
@@ -45,7 +46,7 @@ def make_abacus_scf_input(fp_params):
4546
assert(fp_params["mixing_beta"] >= 0 and fp_params["mixing_beta"] < 1), "'mixing_beta' should between 0 and 1."
4647
ret += "mixing_beta %f\n" % fp_params["mixing_beta"]
4748
if "symmetry" in fp_params:
48-
#assert(fp_params["symmetry"] == 0 or fp_params["symmetry"] == 1), "'symmetry' should be either 0 or 1."
49+
assert(fp_params["symmetry"] == -1 or fp_params["symmetry"] == 0 or fp_params["symmetry"] == 1), "'symmetry' should be either -1, 0 or 1."
4950
ret += "symmetry %d\n" % fp_params["symmetry"]
5051
if "nbands" in fp_params:
5152
if(type(fp_params["nbands"]) == int and fp_params["nbands"] > 0):
@@ -76,7 +77,7 @@ def make_abacus_scf_input(fp_params):
7677
if "out_dos" in fp_params:
7778
assert(type(fp_params["out_dos"]) == int), "'out_dos' should be integer."
7879
ret += "out_dos %d\n" % fp_params["out_dos"]
79-
#paras for deepks
80+
# Parameters for deepks
8081
if "deepks_out_labels" in fp_params:
8182
assert(fp_params["deepks_out_labels"] == 0 or fp_params["deepks_out_labels"] == 1), "'deepks_out_labels' should be either 0 or 1."
8283
ret += "deepks_out_labels %d\n" % fp_params["deepks_out_labels"]
@@ -90,12 +91,13 @@ def make_abacus_scf_input(fp_params):
9091
assert(len(fp_params["deepks_band_range"]) == 2), "length of 'deepks_band_range' should be 2."
9192
ret += "deepks_band_range %d %d\n" % (fp_params["deepks_band_range"][0], fp_params["deepks_band_range"][1])
9293
if "deepks_v_delta" in fp_params:
93-
assert(fp_params["deepks_v_delta"] == 0 or fp_params["deepks_v_delta"] == 1 or fp_params["deepks_v_delta"] == 2), "'deepks_v_delta' should be either 0/1/2."
94+
assert(fp_params["deepks_v_delta"] == -1 or fp_params["deepks_v_delta"] == 0 or fp_params["deepks_v_delta"] == 1 or fp_params["deepks_v_delta"] == 2), "'deepks_v_delta' should be either -1/0/1/2."
9495
ret += "deepks_v_delta %d\n" % fp_params["deepks_v_delta"]
9596
if "model_file" in fp_params:
9697
ret += "deepks_model %s\n" % fp_params["model_file"]
9798
if "out_wfc_lcao" in fp_params:
9899
ret += "out_wfc_lcao %s\n" % fp_params["out_wfc_lcao"]
100+
# Set the parameters for HSE calculation
99101
if fp_params["dft_functional"] == "hse":
100102
ret += "exx_pca_threshold 1e-4\n"
101103
ret += "exx_c_threshold 1e-4\n"
@@ -106,6 +108,9 @@ def make_abacus_scf_input(fp_params):
106108
return ret
107109

108110
def make_abacus_scf_stru(sys_data, fp_pp_files, fp_params):
111+
'''
112+
Make STRU file for abacus scf calculation.
113+
'''
109114
atom_names = sys_data['atom_names'] # Get the list of atom names, e.g., ['Cs', 'Pb', 'I']
110115
atom_numbs = sys_data['atom_numbs'] # Get the number of each atom type, e.g., [4, 4, 12]
111116

@@ -143,7 +148,7 @@ def make_abacus_scf_stru(sys_data, fp_pp_files, fp_params):
143148
ret += f"{fp_params['lattice_constant']}\n\n" # in Bohr
144149
else:
145150
ret += "\nLATTICE_CONSTANT\n"
146-
ret += f"{1 / bohr2ang}\n\n" # Default value is 1/bohr2ang
151+
ret += f"{1 / BOHR2ANG}\n\n" # Default value is 1/BOHR2ANG
147152

148153
ret += "LATTICE_VECTORS\n"
149154
cell = sys_data["cells"][0].reshape([3, 3])

‎deepks/iterate/iterate.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,14 @@
55
import deepks
66
except ImportError as e:
77
sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../../")
8+
from deepks.default import DEFAULT_SCF_ARGS_ABACUS
89
from deepks.utils import copy_file, link_file
910
from deepks.utils import load_yaml, save_yaml
1011
from deepks.utils import load_sys_paths
1112
from deepks.utils import load_basis, save_basis
1213
from deepks.task.workflow import Sequence, Iteration
1314
from deepks.iterate.template import make_scf, make_train
1415
from deepks.iterate.template_abacus import make_scf_abacus #caoyu add 2021-07-22
15-
from deepks.iterate.template_abacus import DEFAULT_SCF_ARGS_ABACUS
1616

1717

1818
# args not specified here may cause error
@@ -100,7 +100,7 @@ def check_arg_dict(data, default, strict=True):
100100
allowed = {k:v for k,v in data.items() if k in default}
101101
outside = {k:v for k,v in data.items() if k not in default}
102102
if outside:
103-
print(f"following ars are not in the default list: {list(outside.keys())}"
103+
print(f"following args are not in the default list: {list(outside.keys())}"
104104
+"and would be discarded" if strict else "but kept", file=sys.stderr)
105105
if strict:
106106
return {**default, **allowed}

0 commit comments

Comments
 (0)