Fix parse_vasp_dir if no magmoms (#148)

* bump GH upload-artifact action to v4 * WIP tests for parsing VASP dirs with and without magmoms * comment on test ZIP origin * fixed vasp parsing and added test --------- Co-authored-by: BowenD-UCB <[email protected]>
CederGroupHub · Apr 1, 2024 · 4b4a489 · 4b4a489
1 parent dd0dd07
commit 4b4a489
Show file tree

Hide file tree

Showing 6 changed files with 102 additions and 30 deletions.
diff --git 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
@@ -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
@@ -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
diff --git a/tests/files/parse-vasp-with-magmoms.zip b/tests/files/parse-vasp-with-magmoms.zip
diff --git a/tests/test_vasp_utils.py 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")