Rewrite abacus interface on devel branch; add unites for qe/pw/scf's crystal unit atomic positions read in.

dpdata/abacus/scf.py

@@ -0,0 +1,180 @@
+import os,sys
+import numpy as np
+
+bohr2ang = 0.5291770
+ry2ev = 13.605698
+kbar2evperang3 = 1e3 / 1.6021892e6
+# The consts are cited from $ABACUS_ROOT/source/src_global/constant.h
+
+
+def get_block (lines, keyword, skip = 0, nlines = None):
+    ret = []
+    found = False
+    if not nlines:
+        nlines = 1e6
+    for idx,ii in enumerate(lines) :
+        if keyword in ii :
+            found = True
+            blk_idx = idx + 1 + skip
+            line_idx = 0
+            while len(lines[blk_idx]) == 0:
+                blk_idx += 1
+            while len(lines[blk_idx]) != 0 and line_idx < nlines and blk_idx != len(lines):
+                ret.append(lines[blk_idx])
+                blk_idx += 1
+                line_idx += 1
+            break
+    if not found:
+        raise RuntimeError("The keyword %s is not found in the script." %keyword)
+    return ret
+
+def get_geometry_in(fname, inlines):
+    geometry_path_in = os.path.join(fname, "STRU")
+    for line in inlines:
+        if "atom_file" in line and "atom_file"==line.split()[0]:
+           atom_file = line.split()[1]
+           geometry_path_in = os.path.join(fname, atom_file)
+           break
+    return geometry_path_in
+
+def get_path_out(fname, inlines):
+    path_out = os.path.join(fname, "OUT.ABACUS/running_scf.log")
+    for line in inlines:
+        if "suffix" in line and "suffix"==line.split()[0]:
+           suffix = line.split()[1]
+           path_out = os.path.join(fname, "OUT.%s/running_scf.log" % suffix)
+           break
+    return path_out
+
+def get_cell(geometry_inlines):
+    cell_lines = get_block(geometry_inlines, "LATTICE_VECTORS", skip = 0, nlines = 3)
+    celldm_lines = get_block(geometry_inlines, "LATTICE_CONSTANT", skip=0, nlines=1)
+
+    celldm = float(celldm_lines[0].split()[0]) * bohr2ang # lattice const is in Bohr
+    cell = []
+    for ii in range(3):
+        cell.append([float(jj) for jj in cell_lines[ii].split()[0:3]])
+    cell = celldm*np.array(cell)
+    return celldm, cell
+
+def get_coords(celldm, cell, geometry_inlines, inlines):
+    coords_lines = get_block(geometry_inlines, "ATOMIC_POSITIONS", skip=0)
+    # assuming that ATOMIC_POSITIONS is at the bottom of the STRU file
+    coord_type = coords_lines[0].split()[0].lower() # cartisan or direct
+    atom_names = [] # element abbr in periodic table 
+    atom_types = [] # index of atom_names of each atom in the geometry
+    atom_numbs = [] # of atoms for each element
+    coords = [] # coordinations of atoms
+    ntype = 0
+    for line in inlines:
+        if "ntype" in line and "ntype"==line.split()[0]:
+            ntype = int(line.split()[1])
+            break
+    if ntype <= 0:
+        raise RuntimeError('ntype cannot be found in INPUT file.')
+    line_idx = 1 # starting line of first element
+    for it in range(ntype):
+        atom_names.append(coords_lines[line_idx].split()[0])
+        line_idx+=2
+        atom_numbs.append(int(coords_lines[line_idx].split()[0]))
+        line_idx+=1
+        for iline in range(atom_numbs[it]):
+            xyz = np.array([float(xx) for xx in coords_lines[line_idx].split()[0:3]])
+            if coord_type == "cartesian":
+                xyz = xyz*celldm
+            elif coord_type == "direct":
+                tmp = np.matmul(xyz, cell)
+                xyz = tmp
+            else:
+                print("coord_type = %s" % coord_type)
+                raise RuntimeError("Input coordination type is invalid.\n Only direct and cartesian are accepted.")
+            coords.append(xyz)
+            atom_types.append(it)
+            line_idx += 1
+    coords = np.array(coords) # need transformation!!!
+    atom_types = np.array(atom_types)
+    return atom_names, atom_numbs, atom_types, coords
+
+def get_energy(outlines):
+    Etot = None
+    for line in outlines:
+        if "!FINAL_ETOT_IS" in line:
+            Etot = float(line.split()[1]) # in eV
+            break
+    if not Etot:
+        not_converge = False
+        for line in outlines:
+            if "convergence has NOT been achieved!" in line:
+                not_converge = True
+                raise RuntimeError("convergence has NOT been achieved in scf!")
+                break
+        if not not_converge:
+            raise RuntimeError("Final total energy cannot be found in output. Unknown problem.")
+    return Etot
+
+def get_force (outlines):
+    force = []
+    force_inlines = get_block (outlines, "TOTAL-FORCE (eV/Angstrom)", skip = 4)
+    for line in force_inlines:
+        force.append([float(f) for f in line.split()[1:4]])
+    force = np.array(force)
+    return force
+
+def get_stress(outlines):
+    stress = []
+    stress_inlines = get_block(outlines, "TOTAL-STRESS (KBAR)", skip = 3)
+    for line in stress_inlines:
+        stress.append([float(f) for f in line.split()])
+    stress = np.array(stress) * kbar2evperang3
+    return stress
+
+
+
+def get_frame (fname):
+    if type(fname) == str:
+        # if the input parameter is only one string, it is assumed that it is the 
+        # base directory containing INPUT file;
+        path_in = os.path.join(fname, "INPUT")
+    else:
+        raise RuntimeError('invalid input')    
+    with open(path_in, 'r') as fp:
+        inlines = fp.read().split('\n')
+
+    geometry_path_in = get_geometry_in(fname, inlines) 
+    path_out = get_path_out(fname, inlines) 
+
+    with open(geometry_path_in, 'r') as fp:
+        geometry_inlines = fp.read().split('\n')
+    with open(path_out, 'r') as fp:
+        outlines = fp.read().split('\n')
+
+    celldm, cell = get_cell(geometry_inlines) 
+    atom_names, natoms, types, coords = get_coords(celldm, cell, geometry_inlines, inlines) 
+
+    energy = get_energy(outlines) 
+    force = get_force (outlines) 
+    stress = get_stress(outlines) * np.linalg.det(cell) 
+
+    data = {}
+    data['atom_names'] = atom_names
+    data['atom_numbs'] = natoms
+    data['atom_types'] = types
+    data['cells'] = cell[np.newaxis, :, :]
+    data['coords'] = coords[np.newaxis, :, :]
+    data['energies'] = np.array(energy)[np.newaxis]
+    data['forces'] = force[np.newaxis, :, :]
+    data['virials'] = stress[np.newaxis, :, :]
+    data['orig'] = np.zeros(3)
+    # print("atom_names = ", data['atom_names'])
+    # print("natoms = ", data['atom_numbs'])
+    # print("types = ", data['atom_types'])
+    # print("cells = ", data['cells'])
+    # print("coords = ", data['coords'])
+    # print("energy = ", data['energies'])
+    # print("force = ", data['forces'])
+    # print("virial = ", data['virials'])
+    return data
+
+if __name__ == "__main__":
+    path = "/home/lrx/work/12_ABACUS_dpgen_interface/dpdata/dpdata/tests/abacus.scf"
+    data = get_frame(path)

dpdata/qe/scf.py

@@ -28,25 +28,49 @@ def get_cell (lines) :
     blk = lines[idx:idx+2]
     ibrav = int(blk[0].replace(',','').split('=')[-1])
     if ibrav == 0:
+        for iline in lines:
+            if 'CELL_PARAMETERS' in iline and 'angstrom' not in iline.lower():
+                raise RuntimeError("CELL_PARAMETERS must be written in Angstrom. Other units are not supported yet.")
         blk = get_block(lines, 'CELL_PARAMETERS')
         for ii in blk:
             ret.append([float(jj) for jj in ii.split()[0:3]])
         ret = np.array(ret)
     elif ibrav == 1:
-        a = float(blk[1].split('=')[-1])
+        a = None
+        for iline in lines:
+            line = iline.replace("=", " ").replace(",", "").split()
+            if len(line) >= 2 and "a" == line[0]:
+                #print("line = ", line)
+                a = float(line[1])
+            if len(line) >= 2 and "celldm(1)" == line[0]:
+                a = float(line[1])*bohr2ang
+        #print("a = ", a)
+        if not a:
+            raise RuntimeError("parameter 'a' or 'celldm(1)' cannot be found.")
         ret = np.array([[a,0.,0.],[0.,a,0.],[0.,0.,a]])
     else:
         sys.exit('ibrav > 1 not supported yet.')
     return ret
 
-def get_coords (lines) :
+def get_coords (lines, cell) :
     coord = []
     atom_symbol_list = []
-    blk = get_block(lines, 'ATOMIC_POSITIONS')
-    for ii in blk:
-        coord.append([float(jj) for jj in ii.split()[1:4]])
-        atom_symbol_list.append(ii.split()[0])
-    coord = np.array(coord)
+    for iline in lines:
+        if 'ATOMIC_POSITIONS' in iline and ('angstrom' not in iline.lower() and 'crystal' not in iline.lower()):
+            raise RuntimeError("ATOMIC_POSITIONS must be written in Angstrom or crystal. Other units are not supported yet.")
+        if 'ATOMIC_POSITIONS' in iline and 'angstrom' in iline.lower():
+            blk = get_block(lines, 'ATOMIC_POSITIONS')
+            for ii in blk:
+                coord.append([float(jj) for jj in ii.split()[1:4]])
+                atom_symbol_list.append(ii.split()[0])
+            coord = np.array(coord)
+        elif 'ATOMIC_POSITIONS' in iline and 'crystal' in iline.lower():
+            blk = get_block(lines, 'ATOMIC_POSITIONS')
+            for ii in blk:
+                coord.append([float(jj) for jj in ii.split()[1:4]])
+                atom_symbol_list.append(ii.split()[0])
+            coord = np.array(coord)
+            coord = np.matmul(coord, cell.T)
     atom_symbol_list = np.array(atom_symbol_list)
     tmp_names, symbol_idx = np.unique(atom_symbol_list, return_index=True)
     atom_types = []
@@ -104,8 +128,8 @@ def get_frame (fname):
         outlines = fp.read().split('\n')
     with open(path_in, 'r') as fp:
         inlines = fp.read().split('\n')
-    atom_names, natoms, types, coords      = get_coords(inlines)
     cell        = get_cell  (inlines)
+    atom_names, natoms, types, coords      = get_coords(inlines, cell)
     energy      = get_energy(outlines)
     force       = get_force (outlines)
     stress      = get_stress(outlines) * np.linalg.det(cell)

dpdata/system.py

@@ -12,6 +12,7 @@
 import dpdata.deepmd.comp
 import dpdata.qe.traj
 import dpdata.qe.scf
+import dpdata.abacus.scf
 import dpdata.siesta.output
 import dpdata.siesta.aiMD_output
 import dpdata.md.pbc
@@ -96,6 +97,8 @@ def __init__ (self,
                 - ``deepmd/npy``: deepmd-kit compressed format (numpy binary)
                 - ``vasp/poscar``: vasp POSCAR
                 - ``qe/cp/traj``: Quantum Espresso CP trajectory files. should have: file_name+'.in' and file_name+'.pos'
+                - ``qe/pw/scf``: Quantum Espresso PW single point calculations. Both input and output files are required. If file_name is a string, it denotes the output file name. Input file name is obtained by replacing 'out' by 'in' from file_name. Or file_name is a list, with the first element being the input file name and the second element being the output filename.
+                - ``abacus/scf``: ABACUS plane wave scf. The directory containing INPUT file is required. 
                 - ``siesta/output``: siesta SCF output file
                 - ``siesta/aimd_output``: siesta aimd output file
                 - ``pwmat/atom.config``: pwmat atom.config
@@ -1366,6 +1369,11 @@ def from_qe_pw_scf(self, file_name) :
             self.data['virials'], \
             = dpdata.qe.scf.get_frame(file_name)
         self.rot_lower_triangular()
+
+    @register_from_funcs.register_funcs('abacus/scf')
+    def from_abacus_pw_scf(self, file_name) :
+        self.data = dpdata.abacus.scf.get_frame(file_name)
+        self.rot_lower_triangular()
 
     @register_from_funcs.register_funcs('siesta/output')
     def from_siesta_output(self, file_name) :

setup.py

@@ -38,6 +38,7 @@
               'dpdata/amber',
               'dpdata/fhi_aims',
               'dpdata/gromacs',
+              'dpdata/abacus',
               'dpdata/rdkit'
     ],
     package_data={'dpdata':['*.json']},

tests/abacus.scf/INPUT

@@ -0,0 +1,18 @@
+INPUT_PARAMETERS
+#Parameters	(General)
+suffix ch4
+atom_file           STRU.ch4 #the filename of file containing atom positions
+kpoint_file         KPT.ch4 #the name of file containing k points
+pseudo_dir ./	
+ntype		  	    2
+nbands		  		8
+#Parameters (Accuracy)
+ecutwfc				100
+symmetry			1
+niter				50
+smearing            gauss #type of smearing: gauss; fd; fixed; mp; mp2; mv
+sigma               0.01
+mixing_beta 0.5
+mixing_type plain
+force 1
+stress 1