Enable model training on Hamiltonian matrix, wavefunction coefficients and energy levels. Supporting for mixed training. Enable some output.

deepks/iterate/generator_abacus.py

@@ -81,10 +81,15 @@ def make_abacus_scf_input(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."
+        assert(fp_params["deepks_bandgap"] == 0  or fp_params["deepks_bandgap"] == 1), "'deepks_bandgap' should be either 0 or 1."
         ret += "deepks_bandgap %d\n" % fp_params["deepks_bandgap"]
+    if "deepks_v_delta" in fp_params:
+        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."
+        ret += "deepks_v_delta %d\n" % fp_params["deepks_v_delta"]
     if "model_file" in fp_params:
         ret += "deepks_model %s\n" % fp_params["model_file"]
+    if "out_wfc_lcao" in fp_params:
+        ret += "out_wfc_lcao %s\n" % fp_params["out_wfc_lcao"]
     if fp_params["dft_functional"] == "hse":
         ret += "exx_pca_threshold 1e-4\n"
         ret += "exx_c_threshold 1e-4\n"

deepks/iterate/template.py

@@ -449,6 +449,7 @@ def make_train(source_train="data_train", source_test="data_test", *,
     )
     # concat
     seq = [run_train, post_train]
+    # seq = [run_train]
     if cleanup:
         clean_train = make_cleanup(
             ["slurm-*.out", "err", "fin.record", "tag_*finished"],

deepks/iterate/template_abacus.py

@@ -64,6 +64,7 @@
     "cal_force": 0,
     "cal_stress": 0,
     "deepks_bandgap": 0,
+    "deepks_v_delta": 0,
     "deepks_out_labels":1,
     "deepks_scf":0,
     "lattice_constant": 1,
@@ -72,6 +73,7 @@
     "run_cmd": "mpirun",
     "sub_size": 1,
     "abacus_path": "/usr/local/bin/ABACUS.mpi",
+    "out_wfc_lcao": 0,
 }
 
 def coord_to_atom(path):
@@ -186,7 +188,7 @@ def convert_data(systems_train, systems_test=None, *,
             cell_data = np.load(f"{sys_paths[i]}/box.npy")
         nframes = atom_data.shape[0]
         natoms = atom_data.shape[1]
-        atoms = atom_data[0,:,0]
+        atoms = atom_data[0,:,0] # if use atom_data[1,:,0], will need at least two frames
         #atoms.sort() # type order
         types = np.unique(atoms) #index in type list
         ntype = types.size
@@ -392,7 +394,7 @@ def make_run_scf_abacus(systems_train, systems_test=None,
 
 
 def gather_stats_abacus(systems_train, systems_test, 
-                train_dump, test_dump, cal_force=0, cal_stress=0, deepks_bandgap=0, **stat_args):
+                train_dump, test_dump, cal_force=0, cal_stress=0, deepks_bandgap=0, deepks_v_delta=0, **stat_args):
     sys_train_paths = [os.path.abspath(s) for s in load_sys_paths(systems_train)]
     sys_test_paths = [os.path.abspath(s) for s in load_sys_paths(systems_test)]
     sys_train_paths = [get_sys_name(s) for s in sys_train_paths]
@@ -420,11 +422,16 @@ def gather_stats_abacus(systems_train, systems_test,
         f0_list=[]
         s0_list=[]
         o0_list=[]
+        h0_list=[]
         e_list=[]
         f_list=[]
         s_list=[]
         o_list=[]
+        h_list=[]
         op_list=[]
+        vdp_list=[] #v_delta_precalc
+        psialpha_list=[]
+        gevdm_list=[]
         gvx_list=[]
         gvepsl_list=[]
         for f in range(nframes):
@@ -462,6 +469,21 @@ def gather_stats_abacus(systems_train, systems_test,
                 if os.path.exists(f"{sys_train_paths[i]}/ABACUS/{f}/orbital_precalc.npy"):
                     orbital_precalc=np.load(f"{sys_train_paths[i]}/ABACUS/{f}/orbital_precalc.npy")
                     op_list.append(orbital_precalc)             
+            if(deepks_v_delta):
+                hcs=np.load(f"{sys_train_paths[i]}/ABACUS/{f}/h_base.npy")
+                h0_list.append(hcs/2)      
+                hcs=np.load(f"{sys_train_paths[i]}/ABACUS/{f}/h_tot.npy")
+                h_list.append(hcs/2)
+                if deepks_v_delta==1:
+                    if os.path.exists(f"{sys_train_paths[i]}/ABACUS/{f}/v_delta_precalc.npy"):
+                        v_delta_precalc=np.load(f"{sys_train_paths[i]}/ABACUS/{f}/v_delta_precalc.npy")
+                        vdp_list.append(v_delta_precalc)
+                elif deepks_v_delta==2:
+                    if os.path.exists(f"{sys_train_paths[i]}/ABACUS/{f}/psialpha.npy") and os.path.exists(f"{sys_train_paths[i]}/ABACUS/{f}/grad_evdm.npy"):
+                        psialpha=np.load(f"{sys_train_paths[i]}/ABACUS/{f}/psialpha.npy")
+                        psialpha_list.append(psialpha)
+                        gevdm=np.load(f"{sys_train_paths[i]}/ABACUS/{f}/grad_evdm.npy")
+                        gevdm_list.append(gevdm)          
         np.save(f"{train_dump}/{sys_train_names[i]}/conv.npy", c_list)
         dm_eig=np.array(d_list)   #concatenate
         np.save(f"{train_dump}/{sys_train_names[i]}/dm_eig.npy", dm_eig)
@@ -500,6 +522,21 @@ def gather_stats_abacus(systems_train, systems_test,
             np.save(f"{train_dump}/{sys_train_names[i]}/o_tot.npy", np.array(o_list))
             if len(op_list) > 0:
                 np.save(f"{train_dump}/{sys_train_names[i]}/orbital_precalc.npy", np.array(op_list))
+        if(deepks_v_delta): 
+            h_base=np.array(h0_list)
+            np.save(f"{train_dump}/{sys_train_names[i]}/h_base.npy", h_base)
+            h_ref=np.load(f"{sys_train_paths[i]}/hamiltonian.npy")
+            np.save(f"{train_dump}/{sys_train_names[i]}/hamiltonian.npy", h_ref)
+            np.save(f"{train_dump}/{sys_train_names[i]}/l_h_delta.npy", h_ref-h_base)
+            np.save(f"{train_dump}/{sys_train_names[i]}/h_tot.npy", np.array(h_list))
+            if len(vdp_list) > 0:
+                np.save(f"{train_dump}/{sys_train_names[i]}/v_delta_precalc.npy", np.array(vdp_list))
+            elif len(psialpha_list) > 0 and len(gevdm_list) > 0:
+                np.save(f"{train_dump}/{sys_train_names[i]}/psialpha.npy", np.array(psialpha_list))
+                np.save(f"{train_dump}/{sys_train_names[i]}/grad_evdm.npy", np.array(gevdm_list))
+            if os.path.exists(f"{sys_train_paths[i]}/overlap.npy"):
+                overlap=np.load(f"{sys_train_paths[i]}/overlap.npy")
+                np.save(f"{train_dump}/{sys_train_names[i]}/overlap.npy", overlap)
     #concatenate data (test)
     if not os.path.exists(test_dump):
             os.mkdir(test_dump)
@@ -517,11 +554,16 @@ def gather_stats_abacus(systems_train, systems_test,
         f0_list=[]
         s0_list=[]
         o0_list=[]
+        h0_list=[]
         e_list=[]
         f_list=[]
         s_list=[]
         o_list=[]
+        h_list=[]
         op_list=[]
+        vdp_list=[] #v_delta_precalc
+        psialpha_list=[]
+        gevdm_list=[]        
         gvx_list=[]
         gvepsl_list=[]
         for f in range(nframes):
@@ -559,6 +601,21 @@ def gather_stats_abacus(systems_train, systems_test,
                 if os.path.exists(f"{sys_test_paths[i]}/ABACUS/{f}/orbital_precalc.npy"):
                     orbital_precalc=np.load(f"{sys_test_paths[i]}/ABACUS/{f}/orbital_precalc.npy")
                     op_list.append(orbital_precalc)
+            if(deepks_v_delta):
+                hcs=np.load(f"{sys_test_paths[i]}/ABACUS/{f}/h_base.npy")
+                h0_list.append(hcs/2)      
+                hcs=np.load(f"{sys_test_paths[i]}/ABACUS/{f}/h_tot.npy")
+                h_list.append(hcs/2)
+                if deepks_v_delta==1:
+                    if os.path.exists(f"{sys_test_paths[i]}/ABACUS/{f}/v_delta_precalc.npy"):
+                        v_delta_precalc=np.load(f"{sys_test_paths[i]}/ABACUS/{f}/v_delta_precalc.npy")
+                        vdp_list.append(v_delta_precalc)
+                elif deepks_v_delta==2:
+                    if os.path.exists(f"{sys_test_paths[i]}/ABACUS/{f}/psialpha.npy") and os.path.exists(f"{sys_test_paths[i]}/ABACUS/{f}/grad_evdm.npy"):
+                        psialpha=np.load(f"{sys_test_paths[i]}/ABACUS/{f}/psialpha.npy")
+                        psialpha_list.append(psialpha)
+                        gevdm=np.load(f"{sys_test_paths[i]}/ABACUS/{f}/grad_evdm.npy")
+                        gevdm_list.append(gevdm)   
         dm_eig=np.array(d_list)   #concatenate
         np.save(f"{test_dump}/{sys_test_names[i]}/dm_eig.npy", dm_eig)
         e_base=np.array(e0_list)
@@ -596,6 +653,21 @@ def gather_stats_abacus(systems_train, systems_test,
             np.save(f"{test_dump}/{sys_test_names[i]}/o_tot.npy", np.array(o_list))
             if len(op_list) > 0:
                 np.save(f"{test_dump}/{sys_test_names[i]}/orbital_precalc.npy", np.array(op_list))
+        if(deepks_v_delta): 
+            h_base=np.array(h0_list)
+            np.save(f"{test_dump}/{sys_test_names[i]}/h_base.npy", h_base)
+            h_ref=np.load(f"{sys_test_paths[i]}/hamiltonian.npy")
+            np.save(f"{test_dump}/{sys_test_names[i]}/hamiltonian.npy", h_ref)
+            np.save(f"{test_dump}/{sys_test_names[i]}/l_h_delta.npy", h_ref-h_base)
+            np.save(f"{test_dump}/{sys_test_names[i]}/h_tot.npy", np.array(h_list))
+            if len(vdp_list) > 0:
+                np.save(f"{test_dump}/{sys_test_names[i]}/v_delta_precalc.npy", np.array(vdp_list))
+            elif len(psialpha_list) > 0 and len(gevdm_list) > 0:
+                np.save(f"{test_dump}/{sys_test_names[i]}/psialpha.npy", np.array(psialpha_list))
+                np.save(f"{test_dump}/{sys_test_names[i]}/grad_evdm.npy", np.array(gevdm_list))
+            if os.path.exists(f"{sys_test_paths[i]}/overlap.npy"):
+                overlap=np.load(f"{sys_test_paths[i]}/overlap.npy")
+                np.save(f"{test_dump}/{sys_test_names[i]}/overlap.npy", overlap)
         np.save(f"{test_dump}/{sys_test_names[i]}/conv.npy",c_list)
     #check convergence and print in log
     from deepks.scf.stats import print_stats
@@ -607,7 +679,7 @@ def gather_stats_abacus(systems_train, systems_test,
 
 
 def make_stat_scf_abacus(systems_train, systems_test=None, *, 
-                  train_dump="data_train", test_dump="data_test", cal_force=0, cal_stress=0, deepks_bandgap=0,
+                  train_dump="data_train", test_dump="data_test", cal_force=0, cal_stress=0, deepks_bandgap=0, deepks_v_delta=0,
                   workdir='.', outlog="log.data", **stat_args):
     # follow same convention for systems as run_scf
     systems_train = [os.path.abspath(s) for s in load_sys_paths(systems_train)]
@@ -624,7 +696,8 @@ def make_stat_scf_abacus(systems_train, systems_test=None, *,
         test_dump=test_dump,
         cal_force=cal_force,
         cal_stress=cal_stress,
-        deepks_bandgap=deepks_bandgap)
+        deepks_bandgap=deepks_bandgap,
+        deepks_v_delta=deepks_v_delta)
     # make task
     return PythonTask(
         gather_stats_abacus,

deepks/model/reader.py

@@ -1,7 +1,7 @@
 import os,time,sys
 import numpy as np
 import torch
-
+# import psutil
 
 def concat_batch(tdicts, dim=0):
     keys = tdicts[0].keys()
@@ -20,13 +20,55 @@ def split_batch(tdict, size, dim=0):
         for i in range(nsecs[0])
     ]
 
+def generalized_eigh(h,L_inv):
+    symm_h=L_inv @ h @ L_inv.mT
+    e,v=torch.linalg.eigh(symm_h)
+    psi=L_inv.mT @ v 
+    return e,psi
+
+#not used now
+def cal_vdp(psialpha,gevdm):
+    # process = psutil.Process(os.getpid())
+    # before_memory_usage = process.memory_info().rss
+    # start=time.time()
+
+    v_delta_precalc_temp=torch.einsum("...kyan,...avmn->...kyavm",psialpha,gevdm)
+
+    v_delta_precalc=torch.einsum("...kxam,...kyavm->...kxyav",psialpha,v_delta_precalc_temp)
+    del v_delta_precalc_temp
+
+    #reshape v_delta_precalc
+    mmax=v_delta_precalc.size(-1)
+    lmax=int((mmax-1)/2)
+    n=int(v_delta_precalc.size(1)/(lmax+1))
+
+    vdp_vector=[]
+    for l in range(lmax+1):
+        ll=v_delta_precalc[:,n*l:n*(l+1),...,:2*l+1]
+        ll=ll.permute(0,2,3,4,5,1,6)
+        ll=ll.flatten(start_dim=-2)
+        vdp_vector.append(ll)
+    vdp=torch.cat(vdp_vector,dim=-1)
+    del vdp_vector
+
+    # end=time.time()
+    # after_memory_usage = process.memory_info().rss
+    # memory_growth = after_memory_usage - before_memory_usage
+    # print(f"Memory growth during cal vdp: {memory_growth / 1024 / 1024} MB")
+
+    # print("all cal_vdp time:",end-start)
+    return vdp
 
 class Reader(object):
     def __init__(self, data_path, batch_size, 
                  e_name="l_e_delta", d_name="dm_eig", 
                  f_name="l_f_delta", gvx_name="grad_vx", 
                  s_name="l_s_delta", gvepsl_name="grad_vepsl", 
                  o_name="l_o_delta", op_name="orbital_precalc",
+                 h_name="l_h_delta", vdp_name="v_delta_precalc",
+                 psialpha_name="psialpha",gevdm_name="grad_evdm",
+                 h_base_name="h_base",h_ref_name="hamiltonian",
+                 read_overlap = False, overlap_name="overlap",
                  eg_name="eg_base", gveg_name="grad_veg", 
                  gldv_name="grad_ldv", conv_name="conv", 
                  atom_name="atom", **kwargs):
@@ -36,15 +78,23 @@ def __init__(self, data_path, batch_size,
         self.f_path = self.check_exist(f_name+".npy")
         self.s_path = self.check_exist(s_name+".npy")
         self.o_path = self.check_exist(o_name+".npy")
+        self.h_path = self.check_exist(h_name+".npy")
+        self.h_base_path=self.check_exist(h_base_name+".npy")
+        self.h_ref_path=self.check_exist(h_ref_name+".npy")
+        self.overlap_path=self.check_exist(overlap_name+".npy")
+        self.psialpha_path=self.check_exist(psialpha_name+".npy")
+        self.gevdm_path=self.check_exist(gevdm_name+".npy")
         self.d_path = self.check_exist(d_name+".npy")
         self.gvx_path = self.check_exist(gvx_name+".npy")
         self.gvepsl_path = self.check_exist(gvepsl_name+".npy")
         self.op_path = self.check_exist(op_name+".npy")
+        self.vdp_path = self.check_exist(vdp_name+".npy")
         self.eg_path = self.check_exist(eg_name+".npy")
         self.gveg_path = self.check_exist(gveg_name+".npy")
         self.gldv_path = self.check_exist(gldv_name+".npy")
         self.c_path = self.check_exist(conv_name+".npy")
         self.a_path = self.check_exist(atom_name+".npy")
+        self.read_overlap=read_overlap
         # load data
         self.load_meta()
         self.prepare()
@@ -117,6 +167,59 @@ def prepare(self):
                 np.load(self.o_path)[conv])
             self.t_data["op"] = torch.tensor(
                 np.load(self.op_path)[conv])
+        if self.h_path is not None and (self.vdp_path is not None or (self.psialpha_path is not None and self.gevdm_path is not None)):
+            h_shape = np.load(self.h_path).shape
+            assert h_shape[-1] == h_shape[-2], \
+                f"The last two dimension of H must have the same size , which is nlocal"
+            self.nlocal = h_shape[-1]
+            self.t_data["lb_vd"] = torch.tensor(
+                np.load(self.h_path)\
+                  .reshape(raw_nframes, -1, self.nlocal, self.nlocal)[conv]) #-1 for nks
+
+            #for v_delta_precalc
+            if self.vdp_path is not None and (self.psialpha_path is not None and self.gevdm_path is not None): #both file exist, choose newer ones
+                if os.path.getmtime(self.vdp_path) >= os.path.getmtime(self.psialpha_path):#psialpha and gevdm modified at the same time
+                    self.psialpha_path=None
+                    self.gevdm_path=None
+                else:
+                    self.vdp_path=None
+            if self.vdp_path is not None:
+                self.t_data["vdp"] = torch.tensor(
+                    np.load(self.vdp_path)\
+                        .reshape(raw_nframes, -1, self.nlocal, self.nlocal, self.natm, self.ndesc)[conv])
+            elif self.psialpha_path is not None and self.gevdm_path is not None:
+                psialpha=np.load(self.psialpha_path)
+                nl=psialpha.shape[2]
+                mmax=psialpha.shape[-1]
+                self.t_data["psialpha"] = torch.tensor(
+                    psialpha\
+                        .reshape(raw_nframes, self.natm, nl, -1, self.nlocal, mmax)[conv])#-1 for nks
+                self.t_data["gevdm"] = torch.tensor(
+                    np.load(self.gevdm_path)\
+                      .reshape(raw_nframes, self.natm, nl, mmax, mmax, mmax)[conv]) 
+                # vdp=cal_vdp(psialpha,gevdm)
+                # self.t_data["vdp"] = vdp\
+                #         .reshape(raw_nframes, -1, self.nlocal, self.nlocal, self.natm, self.ndesc)[conv].clone()
+
+            #for psi labels and band labels
+            self.t_data["h_base"]=torch.tensor(
+                np.load(self.h_base_path)\
+                  .reshape(raw_nframes, -1, self.nlocal, self.nlocal)[conv]) #-1 for nks
+            h_ref=torch.tensor(np.load(self.h_ref_path))
+            if self.read_overlap is True and self.overlap_path is not None:
+                #print("use generalized eigh")
+                overlap=torch.tensor(np.load(self.overlap_path))
+                L=torch.linalg.cholesky(overlap)
+                L_inv=torch.linalg.inv(L)
+                self.t_data["L_inv"]=L_inv\
+                        .reshape(raw_nframes, -1, self.nlocal, self.nlocal)[conv].clone()  
+                band_ref,psi_ref=generalized_eigh(h_ref,L_inv)    
+            else:
+                band_ref,psi_ref=torch.linalg.eigh(h_ref,UPLO='U')
+            self.t_data["lb_band"]=band_ref\
+                  .reshape(raw_nframes, -1, self.nlocal)[conv].clone()             
+            self.t_data["lb_psi"]=psi_ref\
+                  .reshape(raw_nframes, -1, self.nlocal, self.nlocal)[conv].clone()      
         if self.eg_path is not None and self.gveg_path is not None:
             self.t_data['eg0'] = torch.tensor(
                 np.load(self.eg_path)\