Add unlabeled abacus STRU read/dump interface (#303)

* adapt abacus/md interface to MD output without stress calculation.

* add unit test for abacus md without stress output.

* Add files via upload

* add unlabeled read and dump interface for ABACUS STRU format.

* add assertion failed reason for abacus/md interface

* fix open file method

* fix typo; introduce np.testing in abacus tests.

* remove useless if clause

Co-authored-by: LiuRenxi <[email protected]>

Author: Liu-RX

dpdata/abacus/scf.py

@@ -1,7 +1,7 @@
 import os,sys
 import numpy as np
 from ..unit import EnergyConversion, PressureConversion, LengthConversion
-
+import re
 bohr2ang = LengthConversion("bohr", "angstrom").value()
 ry2ev = EnergyConversion("rydberg", "eV").value()
 kbar2evperang3 = PressureConversion("kbar", "eV/angstrom^3").value()
@@ -16,10 +16,10 @@ def get_block (lines, keyword, skip = 0, nlines = None):
             found = True
             blk_idx = idx + 1 + skip
             line_idx = 0
-            while len(lines[blk_idx].split("\s+")) == 0:
+            while len(re.split("\s+", lines[blk_idx])) == 0:
                 blk_idx += 1
             while line_idx < nlines and blk_idx != len(lines):
-                if len(lines[blk_idx].split("\s+")) == 0 or lines[blk_idx] == "":
+                if len(re.split("\s+", lines[blk_idx])) == 0 or lines[blk_idx] == "":
                     blk_idx+=1
                     continue
                 ret.append(lines[blk_idx])
@@ -184,6 +184,100 @@ def get_frame (fname):
     # print("virial = ", data['virials'])
     return data
 
+def get_nele_from_stru(geometry_inlines):
+    key_words_list = ["ATOMIC_SPECIES", "NUMERICAL_ORBITAL", "LATTICE_CONSTANT", "LATTICE_VECTORS", "ATOMIC_POSITIONS", "NUMERICAL_DESCRIPTOR"]
+    keyword_sequence = []
+    keyword_line_index = []
+    atom_names = []
+    atom_numbs = []
+    for iline, line in enumerate(geometry_inlines):
+        if line.split() == []:
+            continue
+        have_key_word = False
+        for keyword in key_words_list:
+            if keyword in line and keyword == line.split()[0]:
+                keyword_sequence.append(keyword)
+                keyword_line_index.append(iline)
+    assert(len(keyword_line_index) == len(keyword_sequence))
+    assert(len(keyword_sequence) > 0)
+    keyword_line_index.append(len(geometry_inlines))
+
+    nele = 0
+    for idx, keyword in enumerate(keyword_sequence):
+        if keyword == "ATOMIC_SPECIES":
+            for iline in range(keyword_line_index[idx]+1, keyword_line_index[idx+1]):
+                if len(re.split("\s+", geometry_inlines[iline])) >= 3:
+                    nele += 1
+    return nele
+
+def get_frame_from_stru(fname):
+    assert(type(fname) == str)
+    with open(fname, 'r') as fp:
+        geometry_inlines = fp.read().split('\n')
+    nele = get_nele_from_stru(geometry_inlines)
+    inlines = ["ntype %d" %nele]
+    celldm, cell = get_cell(geometry_inlines)
+    atom_names, natoms, types, coords = get_coords(celldm, cell, geometry_inlines, inlines) 
+    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['orig'] = np.zeros(3)
+
+    return data
+
+def make_unlabeled_stru(data, frame_idx, pp_file=None, numerical_orbital=None, numerical_descriptor=None, mass=None):
+    out = "ATOMIC_SPECIES\n"
+    for iele in range(len(data['atom_names'])):
+        out += data['atom_names'][iele] + " "
+        if mass is not None:
+            out += "%d "%mass[iele]
+        else:
+            out += "1 "
+        if pp_file is not None:
+            out += "%s\n"%pp_file[iele]
+        else:
+            out += "\n"
+    out += "\n"
+
+    if numerical_orbital is not None:
+        assert(len(numerical_orbital) == len(data['atom_names']))
+        out += "NUMERICAL_ORBITAL\n"
+        for iele in range(len(numerical_orbital)):
+            out += "%s\n"%numerical_orbital[iele]
+        out += "\n"
+
+    if numerical_descriptor is not None:
+        assert(type(numerical_descriptor) == str)
+        out += "NUMERICAL_DESCRIPTOR\n%s\n"%numerical_descriptor
+        out += "\n"
+    
+    out += "LATTICE_CONSTANT\n"
+    out += str(1/bohr2ang) + "\n\n"
+
+    out += "LATTICE_VECTORS\n"
+    for ix in range(3):
+        for iy in range(3):
+            out += str(data['cells'][frame_idx][ix][iy]) + " "
+        out += "\n"
+    out += "\n"
+
+    out += "ATOMIC_POSITIONS\n"
+    out += "Cartesian    # Cartesian(Unit is LATTICE_CONSTANT)\n"
+    #ret += "\n"
+    natom_tot = 0
+    for iele in range(len(data['atom_names'])):
+        out += data['atom_names'][iele] + "\n"
+        out += "0.0\n"
+        out += str(data['atom_numbs'][iele]) + "\n"
+        for iatom in range(data['atom_numbs'][iele]):
+            out += "%.12f %.12f %.12f %d %d %d\n" % (data['coords'][frame_idx][natom_tot, 0], data['coords'][frame_idx][natom_tot, 1], data['coords'][frame_idx][natom_tot, 2], 1, 1, 1)
+            natom_tot += 1
+    assert(natom_tot == sum(data['atom_numbs']))
+    return out
+
 #if __name__ == "__main__":
 #    path = "/home/lrx/work/12_ABACUS_dpgen_interface/dpdata/dpdata/tests/abacus.scf"
 #    data = get_frame(path)

dpdata/plugins/abacus.py

@@ -2,6 +2,39 @@
 import dpdata.abacus.md
 from dpdata.format import Format
 
+@Format.register("abacus/stru")
+@Format.register("stru")
+class AbacusSTRUFormat(Format):
+    def from_system(self, file_name, **kwargs):
+        return dpdata.abacus.scf.get_frame_from_stru(file_name)
+    
+    def to_system(self, data, file_name, frame_idx=0, **kwargs):
+        """
+        Dump the system into ABACUS STRU format file.
+
+        Parameters
+        ----------
+        file_name : str
+            The output file name
+        frame_idx : int
+            The index of the frame to dump
+        pp_file: list of string, optional
+            List of pseudo potential files
+        numerical_orbital: list of string, optional
+            List of orbital files
+        mass: list of float, optional
+            List of atomic masses
+        numerical_descriptor: str, optional
+            numerical descriptor file
+        """
+        
+        pp_file = kwargs.get('pp_file')
+        numerical_orbital = kwargs.get('numerical_orbital')
+        mass = kwargs.get('mass')
+        numerical_descriptor = kwargs.get('numerical_descriptor')
+        stru_string = dpdata.abacus.scf.make_unlabeled_stru(data=data, frame_idx=frame_idx, pp_file=pp_file, numerical_orbital=numerical_orbital, numerical_descriptor=numerical_descriptor, mass=mass)
+        with open(file_name, "w") as fp:
+            fp.write(stru_string)
 
 @Format.register("abacus/scf")
 @Format.register("abacus/pw/scf")

tests/abacus.scf/stru_test

@@ -0,0 +1,32 @@
+ATOMIC_SPECIES
+C 12 C.upf
+H 1 H.upf
+
+NUMERICAL_ORBITAL
+C.orb
+H.orb
+
+NUMERICAL_DESCRIPTOR
+jle.orb
+
+LATTICE_CONSTANT
+1.8897261246257702
+
+LATTICE_VECTORS
+5.291772109029999 0.0 0.0 
+0.0 5.291772109029999 0.0 
+0.0 0.0 5.291772109029999 
+
+ATOMIC_POSITIONS
+Cartesian    # Cartesian(Unit is LATTICE_CONSTANT)
+C
+0.0
+1
+5.192682633809 4.557725978258 4.436846615358 1 1 1
+H
+0.0
+4
+5.416431453540 4.011298860305 3.511161492417 1 1 1
+4.131588222365 4.706745191323 4.431136645083 1 1 1
+5.630930319126 5.521640894956 4.450356541303 1 1 1
+5.499851012568 4.003388899277 5.342621842622 1 1 1

tests/test_abacus_md.py

@@ -29,13 +29,9 @@ def test_cell(self) :
         cell = bohr2ang * 28 * np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
         cell2 = bohr2ang * 5.1 * np.array([[1, 1, 0], [1, 0, 1], [0, 1, 1]])
         for idx in range(np.shape(self.system_water.data['cells'])[0]):
-            for ii in range(cell.shape[0]) :
-                for jj in range(cell.shape[1]) :
-                    self.assertAlmostEqual(self.system_water.data['cells'][idx][ii][jj], cell[ii][jj])
+            np.testing.assert_almost_equal(cell, self.system_water.data['cells'][idx], decimal = 5)
         for idx in range(np.shape(self.system_Si.data['cells'])[0]):
-            for ii in range(cell2.shape[0]) :
-                for jj in range(cell2.shape[1]) :
-                    self.assertAlmostEqual(self.system_Si.data['cells'][idx][ii][jj], cell2[ii][jj])
+            np.testing.assert_almost_equal(self.system_Si.data['cells'][idx], cell2, decimal = 5)
 
     def test_coord(self) :
         with open('abacus.md/water_coord') as fp:
@@ -44,21 +40,15 @@ def test_coord(self) :
                 coord.append([float(jj) for jj in ii.split()])
             coord = np.array(coord)
             coord = coord.reshape([5, 3, 3])
-            for ii in range(coord.shape[0]) :
-                for jj in range(coord.shape[1]) :
-                    for kk in range(coord.shape[2]):
-                        self.assertAlmostEqual(self.system_water.data['coords'][ii][jj][kk], coord[ii][jj][kk])
+            np.testing.assert_almost_equal(self.system_water.data['coords'], coord, decimal = 5)
 
         with open('abacus.md.nostress/Si_coord') as fp2:
             coord = []
             for ii in fp2 :
                 coord.append([float(jj) for jj in ii.split()])
             coord = np.array(coord)
             coord = coord.reshape([4, 2, 3])
-            for ii in range(coord.shape[0]) :
-                for jj in range(coord.shape[1]) :
-                    for kk in range(coord.shape[2]):
-                        self.assertAlmostEqual(self.system_Si.data['coords'][ii][jj][kk], coord[ii][jj][kk])
+            np.testing.assert_almost_equal(self.system_Si.data['coords'], coord, decimal = 5)
 
     def test_force(self) :
         with open('abacus.md/water_force') as fp:
@@ -67,10 +57,7 @@ def test_force(self) :
                 force.append([float(jj) for jj in ii.split()])
             force = np.array(force)
             force = force.reshape([5, 3, 3])
-            for ii in range(force.shape[0]) :
-                for jj in range(force.shape[1]) :
-                    for kk in range(force.shape[2]):
-                        self.assertAlmostEqual(self.system_water.data['forces'][ii][jj][kk], force[ii][jj][kk])
+            np.testing.assert_almost_equal(self.system_water.data['forces'], force, decimal=5)
 
 
         with open('abacus.md.nostress/Si_force') as fp2:
@@ -79,10 +66,7 @@ def test_force(self) :
                 force.append([float(jj) for jj in ii.split()])
             force = np.array(force)
             force = force.reshape([4, 2, 3])
-            for ii in range(force.shape[0]) :
-                for jj in range(force.shape[1]) :
-                    for kk in range(force.shape[2]):
-                        self.assertAlmostEqual(self.system_Si.data['forces'][ii][jj][kk], force[ii][jj][kk])
+            np.testing.assert_almost_equal(self.system_Si.data['forces'], force, decimal=5)
 
 
     def test_virial(self) :
@@ -92,19 +76,14 @@ def test_virial(self) :
                 virial.append([float(jj) for jj in ii.split()])
             virial = np.array(virial)
             virial = virial.reshape([5, 3, 3])
-            for ii in range(virial.shape[0]) :
-                for jj in range(virial.shape[1]) :
-                    for kk in range(virial.shape[2]) :
-                        self.assertAlmostEqual(self.system_water.data['virials'][ii][jj][kk], virial[ii][jj][kk]) 
+            np.testing.assert_almost_equal(self.system_water.data['virials'], virial, decimal=5) 
 
     def test_energy(self) :
         ref_energy = np.array([-466.69285117, -466.69929051, -466.69829826, -466.70364664,
        -466.6976083])
         ref_energy2 = np.array([-211.77184603, -211.78111966, -211.79681663, -211.79875524])
-        for ii in range(5):
-            self.assertAlmostEqual(self.system_water.data['energies'][ii], ref_energy[ii])
-        for ii in range(4):
-            self.assertAlmostEqual(self.system_Si.data['energies'][ii], ref_energy2[ii])
+        np.testing.assert_almost_equal(self.system_water.data['energies'], ref_energy)
+        np.testing.assert_almost_equal(self.system_Si.data['energies'], ref_energy2)
 
 
 class TestABACUSMDLabeledOutput(unittest.TestCase, TestABACUSMD):

tests/test_abacus_pw_scf.py

@@ -11,17 +11,17 @@ class TestABACUSSinglePointEnergy:
     def test_atom_names(self) :
         self.assertEqual(self.system_ch4.data['atom_names'], ['C', 'H'])
         #self.assertEqual(self.system_h2o.data['atom_names'], ['O','H'])
-
+        self.assertEqual(self.system_ch4_unlabeled.data['atom_names'], ['C', 'H'])
     def test_atom_numbs(self) :
         self.assertEqual(self.system_ch4.data['atom_numbs'], [1, 4])
         #self.assertEqual(self.system_h2o.data['atom_numbs'], [64,128])
-
+        self.assertEqual(self.system_ch4_unlabeled.data['atom_numbs'], [1, 4])
     def test_atom_types(self) :
         ref_type = [0,1,1,1,1]
         ref_type =  np.array(ref_type)
         for ii in range(ref_type.shape[0]) :
             self.assertEqual(self.system_ch4.data['atom_types'][ii], ref_type[ii])
-
+            self.assertEqual(self.system_ch4_unlabeled['atom_types'][ii], ref_type[ii])
         # ref_type = [0]*64 + [1]*128
         # ref_type =  np.array(ref_type)
         # for ii in range(ref_type.shape[0]) :
@@ -30,10 +30,8 @@ def test_atom_types(self) :
     def test_cell(self) :
         # cell = 5.29177 * np.eye(3)
         cell = bohr2ang * 10 * np.eye(3)
-        for ii in range(cell.shape[0]) :
-            for jj in range(cell.shape[1]) :
-                self.assertAlmostEqual(self.system_ch4.data['cells'][0][ii][jj], cell[ii][jj])
-
+        np.testing.assert_almost_equal(self.system_ch4.data['cells'][0], cell)
+        np.testing.assert_almost_equal(self.system_ch4_unlabeled.data['cells'][0], cell)
         # fp = open('qe.scf/h2o_cell')
         # cell = []
         # for ii in fp :
@@ -46,15 +44,14 @@ def test_cell(self) :
 
 
     def test_coord(self) :
-        fp = open('abacus.scf/ch4_coord')
-        coord = []
-        for ii in fp :
-            coord.append([float(jj) for jj in ii.split()])
-        coord = np.array(coord)
-        for ii in range(coord.shape[0]) :
-            for jj in range(coord.shape[1]) :
-                self.assertAlmostEqual(self.system_ch4.data['coords'][0][ii][jj], coord[ii][jj], places=5)
-        fp.close()
+        with open('abacus.scf/ch4_coord') as fp:
+            coord = []
+            for ii in fp :
+                coord.append([float(jj) for jj in ii.split()])
+            coord = np.array(coord)
+            np.testing.assert_almost_equal(self.system_ch4.data['coords'][0], coord, decimal=5)
+            np.testing.assert_almost_equal(self.system_ch4_unlabeled.data['coords'][0], coord, decimal=5)
+
 
         # fp = open('qe.scf/h2o_coord')
         # coord = []
@@ -67,15 +64,13 @@ def test_coord(self) :
         # fp.close()
 
     def test_force(self) :
-        fp = open('abacus.scf/ch4_force')
-        force = []
-        for ii in fp :
-            force.append([float(jj) for jj in ii.split()])
-        force = np.array(force)
-        for ii in range(force.shape[0]) :
-            for jj in range(force.shape[1]) :
-                self.assertAlmostEqual(self.system_ch4.data['forces'][0][ii][jj], force[ii][jj])
-        fp.close()
+        with open('abacus.scf/ch4_force') as fp:
+            force = []
+            for ii in fp :
+                force.append([float(jj) for jj in ii.split()])
+            force = np.array(force)
+            np.testing.assert_almost_equal(self.system_ch4.data['forces'][0], force)
+
 
         # fp = open('qe.scf/h2o_force')
         # force = []
@@ -88,15 +83,13 @@ def test_force(self) :
         # fp.close()
 
     def test_virial(self) :
-        fp = open('abacus.scf/ch4_virial')
-        virial = []
-        for ii in fp :
-            virial.append([float(jj) for jj in ii.split()])
-        virial = np.array(virial)
-        for ii in range(virial.shape[0]) :
-            for jj in range(virial.shape[1]) :
-                self.assertAlmostEqual(self.system_ch4.data['virials'][0][ii][jj], virial[ii][jj], places = 3)
-        fp.close()
+        with open('abacus.scf/ch4_virial') as fp:
+            virial = []
+            for ii in fp :
+                virial.append([float(jj) for jj in ii.split()])
+            virial = np.array(virial)
+            np.testing.assert_almost_equal(self.system_ch4.data['virials'][0], virial, decimal = 3)
+
 
         # fp = open('qe.scf/h2o_virial')
         # virial = []
@@ -121,7 +114,7 @@ class TestABACUSLabeledOutput(unittest.TestCase, TestABACUSSinglePointEnergy):
     def setUp(self):
         self.system_ch4 = dpdata.LabeledSystem('abacus.scf',fmt='abacus/scf')
         # self.system_h2o = dpdata.LabeledSystem('qe.scf/02.out',fmt='qe/pw/scf')
-
+        self.system_ch4_unlabeled = dpdata.System('abacus.scf/STRU.ch4', fmt='abacus/stru')
 
 if __name__ == '__main__':
     unittest.main()