diff --git a/dpti/lib/vasp.py b/dpti/lib/vasp.py index c435c3f..3753171 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 0000000..0f10288 --- /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()