Refactor: remove association between ABACUS/STRU and the element number in param.json

dpgen/data/gen.py

@@ -142,7 +142,7 @@ def stru_ele(supercell_stru, stru_out, eles, natoms, jdata, path_work):
         dpks_descriptor_name = os.path.basename(jdata['dpks_descriptor'])
     supercell_stru["atom_masses"] = jdata["atom_masses"]
     supercell_stru["atom_names"] = eles
-    stru_text = make_abacus_scf_stru(supercell_stru, pp_file_names, orb_file_names, dpks_descriptor_name)
+    stru_text = make_abacus_scf_stru(supercell_stru, pp_file_names, orb_file_names, dpks_descriptor_name,type_map=jdata['elements'])
     with open(stru_out, "w") as f:
         f.write(stru_text)
     absolute_pp_file_path = [os.path.abspath(a) for a in jdata["potcars"]]
@@ -224,7 +224,7 @@ def poscar_scale (poscar_in, poscar_out, scale) :
         fout.write("".join(lines))
 
 def poscar_scale_abacus(poscar_in, poscar_out, scale, jdata):
-    stru = get_abacus_STRU(poscar_in, n_ele=len(jdata["elements"]))
+    stru = get_abacus_STRU(poscar_in)
     stru["cells"] *= scale
     stru["coords"] *= scale
     pp_files = [os.path.basename(a) for a in jdata['potcars']]
@@ -234,7 +234,7 @@ def poscar_scale_abacus(poscar_in, poscar_out, scale, jdata):
         orb_file_names = [os.path.basename(a) for a in jdata['orb_files']]
     if 'dpks_descriptor' in jdata:
         dpks_descriptor_name = os.path.basename(jdata['dpks_descriptor'])
-    ret = make_abacus_scf_stru(stru, pp_files, orb_file_names, dpks_descriptor_name)
+    ret = make_abacus_scf_stru(stru, pp_files, orb_file_names, dpks_descriptor_name,type_map=jdata['elements'])
     #ret = make_abacus_scf_stru(stru, pp_files)
     with open(poscar_out, "w") as fp:
         fp.write(ret)
@@ -366,7 +366,7 @@ def make_super_cell_STRU(jdata) :
     to_path = path_sc
     to_file = os.path.join(to_path, 'STRU')
 
-    from_struct=get_abacus_STRU(from_file, n_ele=len(jdata["elements"]))
+    from_struct=get_abacus_STRU(from_file)
     from_struct = make_supercell_abacus(from_struct, super_cell)
     pp_file_names = [os.path.basename(a) for a in jdata['potcars']]
     orb_file_names = None
@@ -375,7 +375,7 @@ def make_super_cell_STRU(jdata) :
         orb_file_names = [os.path.basename(a) for a in jdata['orb_files']]
     if 'dpks_descriptor' in jdata:
         dpks_descriptor_name = os.path.basename(jdata['dpks_descriptor'])
-    stru_text = make_abacus_scf_stru(from_struct, pp_file_names, orb_file_names, dpks_descriptor_name)
+    stru_text = make_abacus_scf_stru(from_struct, pp_file_names, orb_file_names, dpks_descriptor_name,type_map=jdata['elements'])
     with open(to_file, "w") as fp:
         fp.write(stru_text) 
     # make system dir (copy)
@@ -734,7 +734,7 @@ def pert_scaled(jdata) :
                         stru_in = get_abacus_STRU(pos_in)
                         stru_out = shuffle_stru_data(stru_in)
                         with open(pos_out, "w") as fp:
-                            fp.write(make_abacus_scf_stru(stru_out, pp_file, orb_file_names, dpks_descriptor_name))
+                            fp.write(make_abacus_scf_stru(stru_out, pp_file, orb_file_names, dpks_descriptor_name,type_map=jdata['elements']))
                 else :
                     shutil.copy2(pos_in, pos_out)
                 os.remove(pos_in)
@@ -756,7 +756,7 @@ def pert_scaled(jdata) :
                     stru_in = get_abacus_STRU(pos_in)
                     stru_out = shuffle_stru_data(stru_in)
                     with open(pos_out, "w") as fp:
-                        fp.write(make_abacus_scf_stru(stru_out, pp_file, orb_file_names, dpks_descriptor_name))
+                        fp.write(make_abacus_scf_stru(stru_out, pp_file, orb_file_names, dpks_descriptor_name,type_map=jdata['elements']))
             else :
                 shutil.copy2(pos_in, pos_out)
             os.chdir(cwd)

dpgen/generator/lib/abacus_scf.py

@@ -116,10 +116,12 @@ def make_abacus_scf_input(fp_params):
             ret += "%s %s\n" % (key, str(fp_params[key]))
     return ret
 
-def make_abacus_scf_stru(sys_data, fp_pp_files, fp_orb_files = None, fp_dpks_descriptor = None, fp_params = None):
+def make_abacus_scf_stru(sys_data, fp_pp_files, fp_orb_files = None, fp_dpks_descriptor = None, fp_params = None,type_map=None):
     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. "
+    if type_map == None:
+        type_map = atom_names
+
     assert(len(atom_names) == len(atom_numbs)), "Please check the name of atoms. "
     cell = sys_data["cells"].reshape([3, 3])
     coord = sys_data['coords'].reshape([sum(atom_numbs), 3])
@@ -128,10 +130,12 @@ def make_abacus_scf_stru(sys_data, fp_pp_files, fp_orb_files = None, fp_dpks_des
 
     ret = "ATOMIC_SPECIES\n"
     for iatom in range(len(atom_names)):
+        assert (atom_names[iatom] in type_map),"element %s is not defined in type_map" % atom_names[iatom]
+        idx = type_map.index(atom_names[iatom])
         if 'atom_masses' not in sys_data:
-            ret += atom_names[iatom] + " 1.00 " + fp_pp_files[iatom] + "\n"
+            ret += atom_names[iatom] + " 1.00 " + fp_pp_files[idx] + "\n"
         else:
-            ret += atom_names[iatom] + " %.3f "%sys_data['atom_masses'][iatom] + fp_pp_files[iatom] + "\n"
+            ret += atom_names[iatom] + " %.3f "%sys_data['atom_masses'][iatom] + fp_pp_files[idx] + "\n"
 
     if fp_params is not None and "lattice_constant" in fp_params:
         ret += "\nLATTICE_CONSTANT\n"
@@ -162,9 +166,10 @@ def make_abacus_scf_stru(sys_data, fp_pp_files, fp_orb_files = None, fp_dpks_des
 
     if fp_orb_files is not None:
         ret +="\nNUMERICAL_ORBITAL\n"
-        assert(len(fp_orb_files)==len(atom_names))
+        assert(len(fp_orb_files)==len(type_map))
         for iatom in range(len(atom_names)):
-            ret += fp_orb_files[iatom] +"\n"
+            idx = type_map.index(atom_names[iatom])
+            ret += fp_orb_files[idx] +"\n"
 
     if fp_dpks_descriptor is not None:
         ret +="\nNUMERICAL_DESCRIPTOR\n"

dpgen/generator/run.py

@@ -2755,7 +2755,7 @@ def make_fp_abacus_scf(iter_index,
         if 'kspacing' not in fp_params.keys():
             with open("KPT", "w") as fp:
                 fp.write(ret_kpt)
-        ret_stru = make_abacus_scf_stru(sys_data, fp_pp_files, fp_orb_files, fp_dpks_descriptor, fp_params)
+        ret_stru = make_abacus_scf_stru(sys_data, fp_pp_files, fp_orb_files, fp_dpks_descriptor, fp_params,type_map=jdata['type_map'])
         with open("STRU", "w") as fp:
             fp.write(ret_stru)
 

tests/data/Cu.STRU

@@ -0,0 +1,17 @@
+ATOMIC_SPECIES
+Cu 63.550 Cu_ONCV_PBE-1.0.upf
+
+LATTICE_CONSTANT
+1.8897261254578281
+
+LATTICE_VECTORS
+3.76 0.0 0.0 
+0.0 3.76 0.0 
+0.0 0.0 3.76 
+
+ATOMIC_POSITIONS
+Cartesian    # Cartesian(Unit is LATTICE_CONSTANT)
+Cu
+0.0
+1
+1.880000000000 0.000000000000 1.880000000000 1 1 1

tests/data/test_gen_bulk_abacus.py

@@ -18,18 +18,21 @@ def setUp(self):
         self.scale_numb=len(jdata["scale"])
         self.pert_numb=jdata["pert_numb"]
         self.root_dir= out_dir
+        self.jdata = jdata
         create_path(out_dir)
+
+
+    def tearDown(self):
+        shutil.rmtree(self.root_dir)
+
+    def test(self):
+        jdata = self.jdata
         stru_data = make_unit_cell_ABACUS(jdata)
         supercell_stru = make_super_cell_ABACUS(jdata, stru_data)            
         place_element_ABACUS(jdata, supercell_stru)
         make_abacus_relax(jdata, {"fp_resources":{}})
         make_scale_ABACUS(jdata)
         pert_scaled(jdata)
-
-    def tearDown(self):
-        shutil.rmtree(self.root_dir)
-
-    def test(self):
         path=self.out_dir+"/00.place_ele"
         #struct0=Structure.from_file(os.path.join(path,"STRU"))
         alloys=glob.glob(os.path.join(path,"sys-*"))
@@ -49,7 +52,33 @@ def test(self):
             for scale in scales:
                 perts=glob.glob(os.path.join(scale,"[0-9]*")) 
                 self.assertEqual(len(perts),self.pert_numb+1)
-
+
+    def testSTRU(self):
+        jdata = self.jdata
+        jdata['from_poscar_path'] = './Cu.STRU'
+        make_super_cell_STRU(jdata)
+        make_abacus_relax(jdata, {"fp_resources":{}})
+        make_scale_ABACUS(jdata)
+        pert_scaled(jdata)
+        path=self.out_dir+"/00.place_ele"
+        #struct0=Structure.from_file(os.path.join(path,"STRU"))
+        alloys=glob.glob(os.path.join(path,"sys-*"))
+        stru0 = get_abacus_STRU(os.path.join(alloys[0], "STRU"))
+        self.assertEqual(len(alloys),stru0['coords'].shape[0])
+        for ii in alloys:
+            elem_numb=[int(i) for i in ii.split('/')[-1].split('-')[1:]]
+            struct=get_abacus_STRU(os.path.join(ii,"STRU"))
+            self.assertEqual(struct["atom_numbs"], elem_numb)
+        path=self.out_dir+"/01.scale_pert"
+        alloys=glob.glob(os.path.join(path,"sys-*"))
+        self.assertEqual(len(alloys), stru0['coords'].shape[0])
+        for ii in alloys:
+            scales=glob.glob(os.path.join(ii,"scale-*"))
+            self.assertEqual(len(scales),self.scale_numb)
+            for scale in scales:
+                perts=glob.glob(os.path.join(scale,"[0-9]*")) 
+                self.assertEqual(len(perts),self.pert_numb+1)
+
 
 if __name__ == '__main__':
     unittest.main()