Add support for converting move flags between Abacus and VASP (#744)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

- **New Features**
- Enhanced atomic movement handling with the addition of a `move`
parameter across multiple functions, enabling more detailed tracking of
atom movements.
- Introduced selective dynamics support in VASP POSCAR data handling,
allowing for better control over atomic movement flags.
- Added structured representation for cubic boron nitride crystal in new
test files.

- **Bug Fixes**
- Updated atomic positions and magnetic properties in test structures to
reflect accurate simulation data.

- **Tests**
- Added new test methods to validate the handling of movement flags and
ensure the integrity of output structures across various scenarios.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: MabinogiX

dpdata/abacus/md.py

@@ -167,7 +167,7 @@ def get_frame(fname):
     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 = get_coords(
+    atom_names, natoms, types, coords, move = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
     # This coords is not to be used.
@@ -221,5 +221,7 @@ def get_frame(fname):
         data["spins"] = magmom
     if len(magforce) > 0:
         data["mag_forces"] = magforce
+    if len(move) > 0:
+        data["move"] = move
 
     return data

dpdata/abacus/relax.py

@@ -183,7 +183,7 @@ def get_frame(fname):
     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 = get_coords(
+    atom_names, natoms, types, coord_tmp, move = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
 
@@ -218,5 +218,7 @@ def get_frame(fname):
         data["spins"] = magmom
     if len(magforce) > 0:
         data["mag_forces"] = magforce
+    if len(move) > 0:
+        data["move"] = move
 
     return data

dpdata/abacus/scf.py

@@ -299,7 +299,8 @@ def get_coords(celldm, cell, geometry_inlines, inlines=None):
             coords.append(xyz)
             atom_types.append(it)
 
-            move.append(imove)
+            if imove is not None:
+                move.append(imove)
             velocity.append(ivelocity)
             mag.append(imagmom)
             angle1.append(iangle1)
@@ -310,7 +311,8 @@ def get_coords(celldm, cell, geometry_inlines, inlines=None):
             line_idx += 1
     coords = np.array(coords)  # need transformation!!!
     atom_types = np.array(atom_types)
-    return atom_names, atom_numbs, atom_types, coords
+    move = np.array(move, dtype=bool)
+    return atom_names, atom_numbs, atom_types, coords, move
 
 
 def get_energy(outlines):
@@ -477,7 +479,7 @@ def get_frame(fname):
         outlines = fp.read().split("\n")
 
     celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coords = get_coords(
+    atom_names, natoms, types, coords, move = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
     magmom, magforce = get_mag_force(outlines)
@@ -510,6 +512,8 @@ def get_frame(fname):
         data["spins"] = magmom
     if len(magforce) > 0:
         data["mag_forces"] = magforce
+    if len(move) > 0:
+        data["move"] = move[np.newaxis, :, :]
     # print("atom_names = ", data['atom_names'])
     # print("natoms = ", data['atom_numbs'])
     # print("types = ", data['atom_types'])
@@ -561,7 +565,7 @@ def get_frame_from_stru(fname):
     nele = get_nele_from_stru(geometry_inlines)
     inlines = [f"ntype {nele}"]
     celldm, cell = get_cell(geometry_inlines)
-    atom_names, natoms, types, coords = get_coords(
+    atom_names, natoms, types, coords, move = get_coords(
         celldm, cell, geometry_inlines, inlines
     )
     data = {}
@@ -571,6 +575,8 @@ def get_frame_from_stru(fname):
     data["cells"] = cell[np.newaxis, :, :]
     data["coords"] = coords[np.newaxis, :, :]
     data["orig"] = np.zeros(3)
+    if len(move) > 0:
+        data["move"] = move[np.newaxis, :, :]
 
     return data
 
@@ -609,8 +615,8 @@ def make_unlabeled_stru(
         numerical descriptor file
     mass : list of float, optional
         List of atomic masses
-    move : list of list of bool, optional
-        List of the move flag of each xyz direction of each atom
+    move : list of (list of list of bool), optional
+        List of the move flag of each xyz direction of each atom for each frame
     velocity : list of list of float, optional
         List of the velocity of each xyz direction of each atom
     mag : list of (list of float or float), optional
@@ -684,6 +690,9 @@ def process_file_input(file_input, atom_names, input_name):
     if mag is None and data.get("spins") is not None and len(data["spins"]) > 0:
         mag = data["spins"][frame_idx]
 
+    if move is None and data.get("move", None) is not None and len(data["move"]) > 0:
+        move = data["move"][frame_idx]
+
     atom_numbs = sum(data["atom_numbs"])
     for key in [move, velocity, mag, angle1, angle2, sc, lambda_]:
         if key is not None:

dpdata/plugins/abacus.py

@@ -67,6 +67,18 @@ def register_mag_data(data):
         dpdata.LabeledSystem.register_data_type(dt)
 
 
+def register_move_data(data):
+    if "move" in data:
+        dt = DataType(
+            "move",
+            np.ndarray,
+            (Axis.NFRAMES, Axis.NATOMS, 3),
+            required=False,
+            deepmd_name="move",
+        )
+        dpdata.System.register_data_type(dt)
+
+
 @Format.register("abacus/scf")
 @Format.register("abacus/pw/scf")
 @Format.register("abacus/lcao/scf")
@@ -75,6 +87,7 @@ class AbacusSCFFormat(Format):
     def from_labeled_system(self, file_name, **kwargs):
         data = dpdata.abacus.scf.get_frame(file_name)
         register_mag_data(data)
+        register_move_data(data)
         return data
 
 
@@ -86,6 +99,7 @@ class AbacusMDFormat(Format):
     def from_labeled_system(self, file_name, **kwargs):
         data = dpdata.abacus.md.get_frame(file_name)
         register_mag_data(data)
+        register_move_data(data)
         return data
 
 
@@ -97,4 +111,5 @@ class AbacusRelaxFormat(Format):
     def from_labeled_system(self, file_name, **kwargs):
         data = dpdata.abacus.relax.get_frame(file_name)
         register_mag_data(data)
+        register_move_data(data)
         return data

dpdata/plugins/vasp.py

@@ -7,13 +7,26 @@
 import dpdata.vasp.outcar
 import dpdata.vasp.poscar
 import dpdata.vasp.xml
+from dpdata.data_type import Axis, DataType
 from dpdata.format import Format
 from dpdata.utils import open_file, uniq_atom_names
 
 if TYPE_CHECKING:
     from dpdata.utils import FileType
 
 
+def register_move_data(data):
+    if "move" in data:
+        dt = DataType(
+            "move",
+            np.ndarray,
+            (Axis.NFRAMES, Axis.NATOMS, 3),
+            required=False,
+            deepmd_name="move",
+        )
+        dpdata.System.register_data_type(dt)
+
+
 @Format.register("poscar")
 @Format.register("contcar")
 @Format.register("vasp/poscar")
@@ -25,6 +38,7 @@ def from_system(self, file_name: FileType, **kwargs):
             lines = [line.rstrip("\n") for line in fp]
         data = dpdata.vasp.poscar.to_system_data(lines)
         data = uniq_atom_names(data)
+        register_move_data(data)
         return data
 
     def to_system(self, data, file_name: FileType, frame_idx=0, **kwargs):
@@ -99,6 +113,7 @@ def from_labeled_system(
                 vol = np.linalg.det(np.reshape(data["cells"][ii], [3, 3]))
                 data["virials"][ii] *= v_pref * vol
         data = uniq_atom_names(data)
+        register_move_data(data)
         return data
 
 
@@ -135,4 +150,5 @@ def from_labeled_system(self, file_name, begin=0, step=1, **kwargs):
                 vol = np.linalg.det(np.reshape(data["cells"][ii], [3, 3]))
                 data["virials"][ii] *= v_pref * vol
         data = uniq_atom_names(data)
+        register_move_data(data)
         return data

dpdata/vasp/poscar.py

@@ -4,13 +4,22 @@
 import numpy as np
 
 
-def _to_system_data_lower(lines, cartesian=True):
+def _to_system_data_lower(lines, cartesian=True, selective_dynamics=False):
+    def move_flag_mapper(flag):
+        if flag == "T":
+            return True
+        elif flag == "F":
+            return False
+        else:
+            raise RuntimeError(f"Invalid move flag: {flag}")
+
     """Treat as cartesian poscar."""
     system = {}
     system["atom_names"] = [str(ii) for ii in lines[5].split()]
     system["atom_numbs"] = [int(ii) for ii in lines[6].split()]
     scale = float(lines[1])
     cell = []
+    move_flags = []
     for ii in range(2, 5):
         boxv = [float(jj) for jj in lines[ii].split()]
         boxv = np.array(boxv) * scale
@@ -19,12 +28,21 @@ def _to_system_data_lower(lines, cartesian=True):
     natoms = sum(system["atom_numbs"])
     coord = []
     for ii in range(8, 8 + natoms):
-        tmpv = [float(jj) for jj in lines[ii].split()[:3]]
+        tmp = lines[ii].split()
+        tmpv = [float(jj) for jj in tmp[:3]]
         if cartesian:
             tmpv = np.array(tmpv) * scale
         else:
             tmpv = np.matmul(np.array(tmpv), system["cells"][0])
         coord.append(tmpv)
+        if selective_dynamics:
+            if len(tmp) == 6:
+                move_flags.append(list(map(move_flag_mapper, tmp[3:])))
+            else:
+                raise RuntimeError(
+                    f"Invalid move flags, should be 6 columns, got {tmp}"
+                )
+
     system["coords"] = [np.array(coord)]
     system["orig"] = np.zeros(3)
     atom_types = []
@@ -34,20 +52,26 @@ def _to_system_data_lower(lines, cartesian=True):
     system["atom_types"] = np.array(atom_types, dtype=int)
     system["cells"] = np.array(system["cells"])
     system["coords"] = np.array(system["coords"])
+    if move_flags:
+        move_flags = np.array(move_flags, dtype=bool)
+        move_flags = move_flags.reshape((1, natoms, 3))
+        system["move"] = np.array(move_flags, dtype=bool)
     return system
 
 
 def to_system_data(lines):
     # remove the line that has 'selective dynamics'
+    selective_dynamics = False
     if lines[7][0] == "S" or lines[7][0] == "s":
+        selective_dynamics = True
         lines.pop(7)
     is_cartesian = lines[7][0] in ["C", "c", "K", "k"]
     if not is_cartesian:
         if lines[7][0] not in ["d", "D"]:
             raise RuntimeError(
                 "seem not to be a valid POSCAR of vasp 5.x, may be a POSCAR of vasp 4.x?"
             )
-    return _to_system_data_lower(lines, is_cartesian)
+    return _to_system_data_lower(lines, is_cartesian, selective_dynamics)
 
 
 def from_system_data(system, f_idx=0, skip_zeros=True):
@@ -72,6 +96,10 @@ def from_system_data(system, f_idx=0, skip_zeros=True):
             continue
         ret += "%d " % ii
     ret += "\n"
+    move = system.get("move", None)
+    if move is not None and len(move) > 0:
+        ret += "Selective Dynamics\n"
+
     # should use Cartesian for VESTA software
     ret += "Cartesian\n"
     atype = system["atom_types"]
@@ -81,9 +109,26 @@ def from_system_data(system, f_idx=0, skip_zeros=True):
     sort_idx = np.lexsort((np.arange(len(atype)), atype))
     atype = atype[sort_idx]
     posis = posis[sort_idx]
+    if move is not None and len(move) > 0:
+        move = move[f_idx][sort_idx]
+
+    if isinstance(move, np.ndarray):
+        move = move.tolist()
+
     posi_list = []
-    for ii in posis:
-        posi_list.append(f"{ii[0]:15.10f} {ii[1]:15.10f} {ii[2]:15.10f}")
+    for idx in range(len(posis)):
+        ii_posi = posis[idx]
+        line = f"{ii_posi[0]:15.10f} {ii_posi[1]:15.10f} {ii_posi[2]:15.10f}"
+        if move is not None and len(move) > 0:
+            move_flags = move[idx]
+            if not isinstance(move_flags, list) or len(move_flags) != 3:
+                raise RuntimeError(
+                    f"Invalid move flags: {move_flags}, should be a list of 3 bools"
+                )
+            line += " " + " ".join("T" if flag else "F" for flag in move_flags)
+
+        posi_list.append(line)
+
     posi_list.append("")
     ret += "\n".join(posi_list)
     return ret

tests/abacus.scf/STRU.ch4

@@ -18,11 +18,11 @@ Cartesian 				#Cartesian(Unit is LATTICE_CONSTANT)
 C 						#Name of element	
 0.0						#Magnetic for this element.
 1				#Number of atoms
-0.981274803    0.861285385    0.838442496 0 0 0
+0.981274803    0.861285385    0.838442496 1 1 1
 H
 0.0
 4
 1.023557202    0.758025625    0.66351336 0 0 0
-0.78075702    0.889445935    0.837363468 0 0 0
-1.064091613    1.043438905    0.840995502 0 0 0
-1.039321214    0.756530859    1.009609207 0 0 0
+0.78075702    0.889445935    0.837363468 1 0 1
+1.064091613    1.043438905    0.840995502 1 0 1
+1.039321214    0.756530859    1.009609207 0 1 1

tests/abacus.scf/stru_test

@@ -26,7 +26,7 @@ C
 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
+5.416431453540 4.011298860305 3.511161492417 0 0 0
+4.131588222365 4.706745191323 4.431136645083 1 0 1
+5.630930319126 5.521640894956 4.450356541303 1 0 1
+5.499851012568 4.003388899277 5.342621842622 0 1 1

tests/poscars/POSCAR.oh.err1

@@ -0,0 +1,11 @@
+Cubic BN
+   3.57
+ 0.00 0.50 0.50
+ 0.45 0.00 0.50
+ 0.55 0.51 0.00
+ O H
+   1 1
+Selective dynamics
+Cartesian
+ 0.00 0.00 0.00 T T F
+ 0.25 0.25 0.25 F F

tests/poscars/POSCAR.oh.err2

@@ -0,0 +1,11 @@
+Cubic BN
+   3.57
+ 0.00 0.50 0.50
+ 0.45 0.00 0.50
+ 0.55 0.51 0.00
+ O H
+   1 1
+Selective dynamics
+Cartesian
+ 0.00 0.00 0.00 T T F
+ 0.25 0.25 0.25 a T F