From ed2dca38304a2011837fcc11a27f421e20a7ce53 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 24 Aug 2026 04:45:37 +0800 Subject: [PATCH 1/2] Preserve unlabeled POSCAR coordinates Derive coordinate ownership from POSCAR header counts when element labels are absent, while retaining explicit labels when present. Preserve coordinate mode lines, validate requested ordering, and add duplicate-species regression tests. Coding-Agent: Codex Codex-Version: codex-cli 0.149.0 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- dpti/lib/vasp.py | 110 ++++++++++++++++++++++++--------------------- tests/test_vasp.py | 62 +++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 52 deletions(-) create mode 100644 tests/test_vasp.py diff --git a/dpti/lib/vasp.py b/dpti/lib/vasp.py index c435c3f1..a8de79b1 100644 --- a/dpti/lib/vasp.py +++ b/dpti/lib/vasp.py @@ -3,65 +3,71 @@ import numpy as np -def regulate_poscar(poscar_in, poscar_out): - with open(poscar_in) as fp: - lines = fp.read().split("\n") +def _poscar_coordinate_records(lines): + """Return header-derived element ownership for POSCAR coordinate lines.""" names = lines[5].split() counts = [int(ii) for ii in lines[6].split()] - assert len(names) == len(counts) - uniq_name = [] - for ii in names: - if ii not in uniq_name: - uniq_name.append(ii) - uniq_count = np.zeros(len(uniq_name), dtype=int) - for nn, cc in zip(names, counts): - uniq_count[uniq_name.index(nn)] += cc - natoms = np.sum(uniq_count) - posis = lines[8 : 8 + natoms] - all_lines = [] - for ele in uniq_name: - ele_lines = [] - for ii in posis: - ele_name = ii.split()[-1] - if ele_name == ele: - ele_lines.append(ii) - all_lines += ele_lines - all_lines.append("") - ret = lines[0:5] - ret.append(" ".join(uniq_name)) - ret.append(" ".join([str(ii) for ii in uniq_count])) - ret.append("Direct") - ret += all_lines + if len(names) != len(counts): + raise ValueError("POSCAR element names and counts must have equal lengths") + + coordinate_mode_index = 7 + if lines[coordinate_mode_index].strip().lower().startswith("s"): + coordinate_mode_index += 1 + coordinate_start = coordinate_mode_index + 1 + natoms = sum(counts) + positions = lines[coordinate_start : coordinate_start + natoms] + if len(positions) != natoms: + raise ValueError("POSCAR contains fewer coordinate lines than declared atoms") + + header_elements = [ + name for name, count in zip(names, counts) for _ in range(count) + ] + explicit_elements = [line.split()[-1] if line.split() else "" for line in positions] + if all(element in names for element in explicit_elements): + elements = explicit_elements + else: + # Standard POSCAR coordinates are unlabeled and follow header count order. + elements = header_elements + return names, counts, coordinate_start, list(zip(elements, positions)) + + +def _write_grouped_poscar(poscar_in, poscar_out, ordered_names): + with open(poscar_in) as fp: + lines = fp.read().splitlines() + names, counts, coordinate_start, records = _poscar_coordinate_records(lines) + unique_names = list(dict.fromkeys(names)) + if len(ordered_names) != len(set(ordered_names)) or set(ordered_names) != set( + unique_names + ): + raise ValueError("requested POSCAR order must contain each element exactly once") + + grouped = { + name: [line for element, line in records if element == name] + for name in ordered_names + } + new_counts = [len(grouped[name]) for name in ordered_names] + coordinate_lines = [line for name in ordered_names for line in grouped[name]] + ret = lines[:5] + ret.append(" ".join(ordered_names)) + ret.append(" ".join(str(count) for count in new_counts)) + ret.extend(lines[7:coordinate_start]) + ret.extend(coordinate_lines) + ret.extend(lines[coordinate_start + sum(counts) :]) with open(poscar_out, "w") as fp: - fp.write("\n".join(ret)) + fp.write("\n".join(ret) + "\n") -def sort_poscar(poscar_in, poscar_out, new_names): +def regulate_poscar(poscar_in, poscar_out): + """Merge duplicate POSCAR element groups while retaining all coordinates.""" with open(poscar_in) as fp: - lines = fp.read().split("\n") + lines = fp.read().splitlines() names = lines[5].split() - counts = [int(ii) for ii in lines[6].split()] - new_counts = np.zeros(len(counts), dtype=int) - for nn, cc in zip(names, counts): - new_counts[new_names.index(nn)] += cc - natoms = np.sum(new_counts) - posis = lines[8 : 8 + natoms] - all_lines = [] - for ele in new_names: - ele_lines = [] - for ii in posis: - ele_name = ii.split()[-1] - if ele_name == ele: - ele_lines.append(ii) - all_lines += ele_lines - all_lines.append("") - ret = lines[0:5] - ret.append(" ".join(new_names)) - ret.append(" ".join([str(ii) for ii in new_counts])) - ret.append("Direct") - ret += all_lines - with open(poscar_out, "w") as fp: - fp.write("\n".join(ret)) + _write_grouped_poscar(poscar_in, poscar_out, list(dict.fromkeys(names))) + + +def sort_poscar(poscar_in, poscar_out, new_names): + """Reorder POSCAR element groups using header counts for unlabeled coordinates.""" + _write_grouped_poscar(poscar_in, poscar_out, new_names) def perturb_xz(poscar_in, poscar_out, pert=0.01): diff --git a/tests/test_vasp.py b/tests/test_vasp.py new file mode 100644 index 00000000..0f102881 --- /dev/null +++ b/tests/test_vasp.py @@ -0,0 +1,62 @@ +import os +import tempfile +import unittest + +from dpti.lib.vasp import regulate_poscar, sort_poscar + +POSCAR = """water +1.0 +1 0 0 +0 1 0 +0 0 1 +O H O +1 2 1 +Direct +0.0 0.0 0.0 +0.1 0.1 0.1 +0.2 0.2 0.2 +0.3 0.3 0.3 +""" + + +class TestPoscarGrouping(unittest.TestCase): + def setUp(self): + self.tempdir = tempfile.TemporaryDirectory() + self.input_path = os.path.join(self.tempdir.name, "POSCAR") + with open(self.input_path, "w") as fp: + fp.write(POSCAR) + + def tearDown(self): + self.tempdir.cleanup() + + def test_regulate_retains_unlabeled_coordinates(self): + output_path = os.path.join(self.tempdir.name, "regulated.POSCAR") + + regulate_poscar(self.input_path, output_path) + + with open(output_path) as fp: + lines = fp.read().splitlines() + self.assertEqual(lines[5], "O H") + self.assertEqual(lines[6], "2 2") + self.assertEqual( + lines[8:12], + ["0.0 0.0 0.0", "0.3 0.3 0.3", "0.1 0.1 0.1", "0.2 0.2 0.2"], + ) + + def test_sort_uses_header_element_ownership(self): + output_path = os.path.join(self.tempdir.name, "sorted.POSCAR") + + sort_poscar(self.input_path, output_path, ["H", "O"]) + + with open(output_path) as fp: + lines = fp.read().splitlines() + self.assertEqual(lines[5], "H O") + self.assertEqual(lines[6], "2 2") + self.assertEqual( + lines[8:12], + ["0.1 0.1 0.1", "0.2 0.2 0.2", "0.0 0.0 0.0", "0.3 0.3 0.3"], + ) + + +if __name__ == "__main__": + unittest.main() From ecc55d4102ce2f8d0194df1a5263c6b703127334 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 20:46:12 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- dpti/lib/vasp.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dpti/lib/vasp.py b/dpti/lib/vasp.py index a8de79b1..37531713 100644 --- a/dpti/lib/vasp.py +++ b/dpti/lib/vasp.py @@ -19,9 +19,7 @@ def _poscar_coordinate_records(lines): if len(positions) != natoms: raise ValueError("POSCAR contains fewer coordinate lines than declared atoms") - header_elements = [ - name for name, count in zip(names, counts) for _ in range(count) - ] + header_elements = [name for name, count in zip(names, counts) for _ in range(count)] explicit_elements = [line.split()[-1] if line.split() else "" for line in positions] if all(element in names for element in explicit_elements): elements = explicit_elements @@ -39,7 +37,9 @@ def _write_grouped_poscar(poscar_in, poscar_out, ordered_names): if len(ordered_names) != len(set(ordered_names)) or set(ordered_names) != set( unique_names ): - raise ValueError("requested POSCAR order must contain each element exactly once") + raise ValueError( + "requested POSCAR order must contain each element exactly once" + ) grouped = { name: [line for element, line in records if element == name]