Refactor the reading and writing of abacus/stru format

dpdata/abacus/md.py

@@ -9,12 +9,11 @@
 
 from .scf import (
     bohr2ang,
-    get_cell,
-    get_coords,
     get_geometry_in,
     get_mag_force,
     kbar2evperang3,
 )
+from .stru import get_frame_from_stru
 
 # Read in geometries from an ABACUS MD trajectory.
 # The atomic coordinates are read in from generated files in OUT.XXXX.
@@ -164,12 +163,12 @@ def get_frame(fname):
     geometry_path_in = get_geometry_in(fname, inlines)  # base dir of STRU
     path_out = get_path_out(fname, inlines)
 
-    with open_file(geometry_path_in) as fp:
-        geometry_inlines = fp.read().split("\n")
-    celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coords, move, magmom = get_coords(
-        celldm, cell, geometry_inlines, inlines
-    )
+    data = get_frame_from_stru(geometry_path_in)
+    natoms = data["atom_numbs"]
+    # should remove spins from STRU file
+    if "spins" in data:
+        data.pop("spins")
+
     # This coords is not to be used.
     dump_freq = get_coord_dump_freq(inlines=inlines)
     # ndump = int(os.popen("ls -l %s | grep 'md_pos_' | wc -l" %path_out).readlines()[0])
@@ -203,10 +202,6 @@ def get_frame(fname):
 
     magmom, magforce = get_mag_force(outlines)
 
-    data = {}
-    data["atom_names"] = atom_names
-    data["atom_numbs"] = natoms
-    data["atom_types"] = types
     data["cells"] = cells
     # for idx in range(ndump):
     #    data['cells'][:, :, :] = cell
@@ -221,7 +216,9 @@ def get_frame(fname):
         data["spins"] = magmom
     if len(magforce) > 0:
         data["force_mags"] = magforce
-    if len(move) > 0:
-        data["move"] = move
+
+    # need to expand the move.
+    if "move" in data:
+        data["move"] = [data["move"][0] for i in range(ndump)]
 
     return data

dpdata/abacus/relax.py

@@ -10,12 +10,11 @@
     bohr2ang,
     collect_force,
     collect_stress,
-    get_cell,
-    get_coords,
     get_geometry_in,
     get_mag_force,
     kbar2evperang3,
 )
+from .stru import get_frame_from_stru
 
 # Read in geometries from an ABACUS RELAX(CELL-RELAX) trajectory in OUT.XXXX/runnning_relax/cell-relax.log.
 
@@ -47,7 +46,7 @@ def get_coords_from_log(loglines, natoms):
             natoms_log += int(line.split()[-1])
 
     assert natoms_log > 0 and natoms_log == natoms, (
-        "ERROR: detected atom number in log file is %d" % natoms  # noqa: UP031
+        f"ERROR: detected atom number in log file is {natoms_log}, while the atom number in STRU file is {natoms}"
     )
 
     energy = []
@@ -180,31 +179,22 @@ def get_frame(fname):
     with open_file(path_in) as fp:
         inlines = fp.read().split("\n")
     geometry_path_in = get_geometry_in(fname, inlines)  # base dir of STRU
-    with open_file(geometry_path_in) as fp:
-        geometry_inlines = fp.read().split("\n")
-    celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coord_tmp, move, magmom = get_coords(
-        celldm, cell, geometry_inlines, inlines
-    )
+
+    data = get_frame_from_stru(geometry_path_in)
+    natoms = sum(data["atom_numbs"])
+    # should remove spins from STRU file
+    if "spins" in data:
+        data.pop("spins")
 
     logf = get_log_file(fname, inlines)
     assert os.path.isfile(logf), f"Error: can not find {logf}"
     with open_file(logf) as f1:
         lines = f1.readlines()
 
-    atomnumber = 0
-    for i in natoms:
-        atomnumber += i
-    energy, cells, coords, force, stress, virial = get_coords_from_log(
-        lines, atomnumber
-    )
+    energy, cells, coords, force, stress, virial = get_coords_from_log(lines, natoms)
 
     magmom, magforce = get_mag_force(lines)
 
-    data = {}
-    data["atom_names"] = atom_names
-    data["atom_numbs"] = natoms
-    data["atom_types"] = types
     data["cells"] = cells
     data["coords"] = coords
     data["energies"] = energy
@@ -218,7 +208,7 @@ def get_frame(fname):
         data["spins"] = magmom
     if len(magforce) > 0:
         data["force_mags"] = magforce
-    if len(move) > 0:
-        data["move"] = move
+    if "move" in data:
+        data["move"] = [data["move"][0] for i in range(len(data["energies"]))]
 
     return data