From d0a786835a00b5b35cbf6ba72cd2c8373f69443c Mon Sep 17 00:00:00 2001 From: AngelJia Date: Tue, 29 Oct 2024 00:53:34 -0500 Subject: [PATCH] Add support for converting move flags between Abacus and VASP (#744) ## 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. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- dpdata/abacus/md.py | 4 ++- dpdata/abacus/relax.py | 4 ++- dpdata/abacus/scf.py | 21 +++++++---- dpdata/plugins/abacus.py | 15 ++++++++ dpdata/plugins/vasp.py | 16 +++++++++ dpdata/vasp/poscar.py | 55 ++++++++++++++++++++++++++--- tests/abacus.scf/STRU.ch4 | 8 ++--- tests/abacus.scf/stru_test | 8 ++--- tests/poscars/POSCAR.oh.err1 | 11 ++++++ tests/poscars/POSCAR.oh.err2 | 11 ++++++ tests/test_abacus_stru_dump.py | 48 +++++++++++++++++++++++-- tests/test_vasp_poscar_dump.py | 15 ++++++++ tests/test_vasp_poscar_to_system.py | 31 ++++++++++++++++ 13 files changed, 224 insertions(+), 23 deletions(-) create mode 100644 tests/poscars/POSCAR.oh.err1 create mode 100644 tests/poscars/POSCAR.oh.err2 diff --git a/dpdata/abacus/md.py b/dpdata/abacus/md.py index bbb023439..aa4215d84 100644 --- a/dpdata/abacus/md.py +++ b/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 diff --git a/dpdata/abacus/relax.py b/dpdata/abacus/relax.py index be80bc90a..63f13678f 100644 --- a/dpdata/abacus/relax.py +++ b/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 diff --git a/dpdata/abacus/scf.py b/dpdata/abacus/scf.py index b1b2cfed9..43e65f7ea 100644 --- a/dpdata/abacus/scf.py +++ b/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: diff --git a/dpdata/plugins/abacus.py b/dpdata/plugins/abacus.py index cbd1475f2..b3e7c98a4 100644 --- a/dpdata/plugins/abacus.py +++ b/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 diff --git a/dpdata/plugins/vasp.py b/dpdata/plugins/vasp.py index 0160bde29..d25d0c25e 100644 --- a/dpdata/plugins/vasp.py +++ b/dpdata/plugins/vasp.py @@ -7,6 +7,7 @@ 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 @@ -14,6 +15,18 @@ 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 diff --git a/dpdata/vasp/poscar.py b/dpdata/vasp/poscar.py index 102e79041..30073e2b3 100644 --- a/dpdata/vasp/poscar.py +++ b/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,12 +52,18 @@ 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: @@ -47,7 +71,7 @@ def to_system_data(lines): 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 diff --git a/tests/abacus.scf/STRU.ch4 b/tests/abacus.scf/STRU.ch4 index cf97747aa..bc33cbe94 100644 --- a/tests/abacus.scf/STRU.ch4 +++ b/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 diff --git a/tests/abacus.scf/stru_test b/tests/abacus.scf/stru_test index 22d619c93..e40364093 100644 --- a/tests/abacus.scf/stru_test +++ b/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 diff --git a/tests/poscars/POSCAR.oh.err1 b/tests/poscars/POSCAR.oh.err1 new file mode 100644 index 000000000..1b7521e6e --- /dev/null +++ b/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 diff --git a/tests/poscars/POSCAR.oh.err2 b/tests/poscars/POSCAR.oh.err2 new file mode 100644 index 000000000..ed52d4d08 --- /dev/null +++ b/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 diff --git a/tests/test_abacus_stru_dump.py b/tests/test_abacus_stru_dump.py index 4549b6d16..8891dee3f 100644 --- a/tests/test_abacus_stru_dump.py +++ b/tests/test_abacus_stru_dump.py @@ -141,8 +141,52 @@ def test_dump_spin(self): 4 5.416431453540 4.011298860305 3.511161492417 1 1 1 mag 4.000000000000 5.000000000000 6.000000000000 4.131588222365 4.706745191323 4.431136645083 1 1 1 mag 1.000000000000 1.000000000000 1.000000000000 -5.630930319126 5.521640894956 4.450356541303 1 1 1 mag 2.000000000000 2.000000000000 2.000000000000 -5.499851012568 4.003388899277 5.342621842622 1 1 1 mag 3.000000000000 3.000000000000 3.000000000000 +5.630930319126 5.521640894956 4.450356541303 0 0 0 mag 2.000000000000 2.000000000000 2.000000000000 +5.499851012568 4.003388899277 5.342621842622 0 0 0 mag 3.000000000000 3.000000000000 3.000000000000 +""" + self.assertTrue(stru_ref in c) + + def test_dump_move_from_vasp(self): + self.system = dpdata.System() + self.system.from_vasp_poscar(os.path.join("poscars", "POSCAR.oh.c")) + self.system.to( + "abacus/stru", + "STRU_tmp", + pp_file={"O": "O.upf", "H": "H.upf"}, + ) + assert os.path.isfile("STRU_tmp") + with open("STRU_tmp") as f: + c = f.read() + + stru_ref = """O +0.0 +1 +0.000000000000 0.000000000000 0.000000000000 1 1 0 +H +0.0 +1 +1.262185604418 0.701802783513 0.551388341420 0 0 0 +""" + self.assertTrue(stru_ref in c) + + self.system.to( + "abacus/stru", + "STRU_tmp", + pp_file={"O": "O.upf", "H": "H.upf"}, + move=[[True, False, True], [False, True, False]], + ) + assert os.path.isfile("STRU_tmp") + with open("STRU_tmp") as f: + c = f.read() + + stru_ref = """O +0.0 +1 +0.000000000000 0.000000000000 0.000000000000 1 0 1 +H +0.0 +1 +1.262185604418 0.701802783513 0.551388341420 0 1 0 """ self.assertTrue(stru_ref in c) diff --git a/tests/test_vasp_poscar_dump.py b/tests/test_vasp_poscar_dump.py index 62f215986..d55228117 100644 --- a/tests/test_vasp_poscar_dump.py +++ b/tests/test_vasp_poscar_dump.py @@ -34,6 +34,21 @@ def setUp(self): self.system = dpdata.System() self.system.from_vasp_poscar("tmp.POSCAR") + def test_dump_move_flags(self): + tmp_system = dpdata.System() + tmp_system.from_vasp_poscar(os.path.join("poscars", "POSCAR.oh.c")) + tmp_system.to_vasp_poscar("tmp.POSCAR") + self.system = dpdata.System() + self.system.from_vasp_poscar("tmp.POSCAR") + with open("tmp.POSCAR") as f: + content = f.read() + + stru_ref = """Cartesian + 0.0000000000 0.0000000000 0.0000000000 T T F + 1.2621856044 0.7018027835 0.5513883414 F F F +""" + self.assertTrue(stru_ref in content) + class TestPOSCARSkipZeroAtomNumb(unittest.TestCase): def tearDown(self): diff --git a/tests/test_vasp_poscar_to_system.py b/tests/test_vasp_poscar_to_system.py index 7457d33d2..9e5d7f37f 100644 --- a/tests/test_vasp_poscar_to_system.py +++ b/tests/test_vasp_poscar_to_system.py @@ -14,6 +14,37 @@ def setUp(self): self.system = dpdata.System() self.system.from_vasp_poscar(os.path.join("poscars", "POSCAR.oh.c")) + def test_move_flags(self): + expected = np.array([[[True, True, False], [False, False, False]]]) + self.assertTrue(np.array_equal(self.system["move"], expected)) + + +class TestPOSCARMoveFlags(unittest.TestCase): + def setUp(self): + self.tmp_file = "POSCAR.tmp.1" + + def tearDown(self): + if os.path.exists(self.tmp_file): + os.remove(self.tmp_file) + + def test_move_flags_error1(self): + with self.assertRaisesRegex(RuntimeError, "Invalid move flags.*?"): + dpdata.System().from_vasp_poscar(os.path.join("poscars", "POSCAR.oh.err1")) + + def test_move_flags_error2(self): + with self.assertRaisesRegex(RuntimeError, "Invalid move flag: a"): + dpdata.System().from_vasp_poscar(os.path.join("poscars", "POSCAR.oh.err2")) + + def test_move_flags_error3(self): + system = dpdata.System().from_vasp_poscar( + os.path.join("poscars", "POSCAR.oh.c") + ) + system.data["move"] = np.array([[[True, True], [False, False]]]) + with self.assertRaisesRegex( + RuntimeError, "Invalid move flags:.*?should be a list of 3 bools" + ): + system.to_vasp_poscar(self.tmp_file) + class TestPOSCARDirect(unittest.TestCase, TestPOSCARoh): def setUp(self):