diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index e9269322..6ff3cba4 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -34,7 +34,7 @@ jobs: pip install build python -m build --sdist - - uses: actions/upload-artifact@v3 + - uses: actions/upload-artifact@v4 with: path: dist/*.tar.gz @@ -57,7 +57,7 @@ jobs: CIBW_ARCHS_MACOS: universal2 - name: Save artifact - uses: actions/upload-artifact@v3 + uses: actions/upload-artifact@v4 with: path: wheelhouse @@ -70,7 +70,7 @@ jobs: id-token: write steps: - name: Download build artifacts - uses: actions/download-artifact@v3 + uses: actions/download-artifact@v4 with: name: artifact path: dist diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 72275eec..25760a2d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_install_hook_types: [pre-commit, commit-msg] repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.3.2 + rev: v0.3.4 hooks: - id: ruff args: [--fix] @@ -46,7 +46,7 @@ repos: - svelte - repo: https://github.com/pre-commit/mirrors-eslint - rev: v9.0.0-beta.2 + rev: v9.0.0-rc.0 hooks: - id: eslint types: [file] diff --git a/chgnet/utils/vasp_utils.py b/chgnet/utils/vasp_utils.py index 05f61d69..6cead739 100644 --- a/chgnet/utils/vasp_utils.py +++ b/chgnet/utils/vasp_utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os.path import re from typing import TYPE_CHECKING @@ -22,28 +23,29 @@ def parse_vasp_dir( check_electronic_convergence (bool): if set to True, this function will raise Exception to VASP calculation that did not achieve electronic convergence. """ - try: - oszicar = Oszicar(f"{file_root}/OSZICAR") - vasprun_orig = Vasprun( - f"{file_root}/vasprun.xml", - parse_dos=False, - parse_eigen=False, - parse_projected_eigen=False, - parse_potcar_file=False, - exception_on_bad_xml=False, - ) - outcar_filename = f"{file_root}/OUTCAR" - except Exception: - oszicar = Oszicar(f"{file_root}/OSZICAR.gz") - vasprun_orig = Vasprun( - f"{file_root}/vasprun.xml.gz", - parse_dos=False, - parse_eigen=False, - parse_projected_eigen=False, - parse_potcar_file=False, - exception_on_bad_xml=False, - ) - outcar_filename = f"{file_root}/OUTCAR.gz" + if os.path.exists(file_root) is False: + raise FileNotFoundError("No such file or directory") + + if os.path.exists(f"{file_root}/OSZICAR"): + oszicar_path = f"{file_root}/OSZICAR" + vasprun_path = f"{file_root}/vasprun.xml" + outcar_path = f"{file_root}/OUTCAR" + elif os.path.exists(f"{file_root}/OSZICAR"): + oszicar_path = f"{file_root}/OSZICAR.gz" + vasprun_path = f"{file_root}/vasprun.xml.gz" + outcar_path = f"{file_root}/OUTCAR.gz" + else: + raise RuntimeError(f"No data parsed from {file_root}!") + + oszicar = Oszicar(oszicar_path) + vasprun_orig = Vasprun( + vasprun_path, + parse_dos=False, + parse_eigen=False, + parse_projected_eigen=False, + parse_potcar_file=False, + exception_on_bad_xml=False, + ) charge = [] mag_x = [] @@ -52,7 +54,7 @@ def parse_vasp_dir( header = [] all_lines = [] - for line in reverse_readfile(outcar_filename): + for line in reverse_readfile(outcar_path): clean = line.strip() all_lines.append(clean) @@ -132,7 +134,7 @@ def parse_vasp_dir( "stress": None if "stress" not in vasprun_orig.ionic_steps[0] else [], } - for ionic_step, mag_step in zip(vasprun_orig.ionic_steps, mag_x_all): + for index, ionic_step in enumerate(vasprun_orig.ionic_steps): if ( check_electronic_convergence and len(ionic_step["electronic_steps"]) >= vasprun_orig.parameters["NELM"] @@ -143,7 +145,8 @@ def parse_vasp_dir( dataset["uncorrected_total_energy"].append(ionic_step["e_0_energy"]) dataset["energy_per_atom"].append(ionic_step["e_0_energy"] / n_atoms) dataset["force"].append(ionic_step["forces"]) - dataset["magmom"].append([site["tot"] for site in mag_step]) + if mag_x_all != []: + dataset["magmom"].append([site["tot"] for site in mag_x_all[index]]) if "stress" in ionic_step: dataset["stress"].append(ionic_step["stress"]) diff --git a/tests/files/parse-vasp-no-magmoms.zip b/tests/files/parse-vasp-no-magmoms.zip new file mode 100644 index 00000000..aabd938c Binary files /dev/null and b/tests/files/parse-vasp-no-magmoms.zip differ diff --git a/tests/files/parse-vasp-with-magmoms.zip b/tests/files/parse-vasp-with-magmoms.zip new file mode 100644 index 00000000..1fbafb8e Binary files /dev/null and b/tests/files/parse-vasp-with-magmoms.zip differ diff --git a/tests/test_vasp_utils.py b/tests/test_vasp_utils.py new file mode 100644 index 00000000..e94a1af3 --- /dev/null +++ b/tests/test_vasp_utils.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import os.path +from typing import TYPE_CHECKING +from zipfile import ZipFile + +import pytest +from pymatgen.core import Structure + +from chgnet import ROOT +from chgnet.utils import parse_vasp_dir + +if TYPE_CHECKING: + from pathlib import Path + + +def test_parse_vasp_dir_with_magmoms(tmp_path: Path): + with ZipFile(f"{ROOT}/tests/files/parse-vasp-with-magmoms.zip") as zip_ref: + zip_ref.extractall(tmp_path) + vasp_path = os.path.join(tmp_path, "parse-vasp-with-magmoms") + dataset_dict = parse_vasp_dir(vasp_path) + + assert isinstance(dataset_dict, dict) + assert len(dataset_dict["structure"]) > 0 + assert len(dataset_dict["uncorrected_total_energy"]) > 0 + assert len(dataset_dict["energy_per_atom"]) > 0 + assert len(dataset_dict["force"]) > 0 + assert len(dataset_dict["magmom"]) == len(dataset_dict["force"]) + assert len(dataset_dict["stress"]) > 0 + + for structure in dataset_dict["structure"]: + assert isinstance(structure, Structure) + + for magmom in dataset_dict["magmom"]: + assert len(magmom) == len(dataset_dict["structure"][0]) + + +def test_parse_vasp_dir_without_magmoms(tmp_path: Path): + # using test.zip shared for error repro in + # https://github.com/CederGroupHub/chgnet/issues/147 + with ZipFile(f"{ROOT}/tests/files/parse-vasp-no-magmoms.zip") as zip_ref: + zip_ref.extractall(tmp_path) + vasp_path = os.path.join(tmp_path, "parse-vasp-no-magmoms") + dataset_dict = parse_vasp_dir(vasp_path) + + assert isinstance(dataset_dict, dict) + assert len(dataset_dict["structure"]) > 0 + assert len(dataset_dict["uncorrected_total_energy"]) > 0 + assert len(dataset_dict["energy_per_atom"]) > 0 + assert len(dataset_dict["force"]) > 0 + assert len(dataset_dict["magmom"]) == 0 + assert len(dataset_dict["stress"]) > 0 + + for structure in dataset_dict["structure"]: + assert isinstance(structure, Structure) + + for magmom in dataset_dict["magmom"]: + assert len(magmom) == len(dataset_dict["structure"][0]) + assert all(mag == 0.0 for mag in magmom) + + +def test_parse_vasp_dir_no_data(): + # test non-existing directory + with pytest.raises(FileNotFoundError, match="No such file or directory"): + parse_vasp_dir(f"{ROOT}/tests/files/non-existent") + + # test existing directory without VASP files + with pytest.raises(RuntimeError, match="No data parsed from"): + parse_vasp_dir(f"{ROOT}/tests/files")