diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index fb054071df..4cc9c431a1 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -21,7 +21,7 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" cache: pip cache-dependency-path: pyproject.toml @@ -51,7 +51,7 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" - name: Build run: | diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 9ad7cd3567..72f2b00202 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -26,7 +26,7 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" cache: pip cache-dependency-path: pyproject.toml diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 41c7bdf320..e277c0708f 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -17,7 +17,7 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" cache: pip cache-dependency-path: pyproject.toml @@ -39,7 +39,7 @@ jobs: shell: bash -l {0} # enables conda/mamba env activation by reading bash profile strategy: matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.11", "3.12"] split: [1, 2, 3] steps: @@ -48,10 +48,11 @@ jobs: - name: Set up micromamba uses: mamba-org/setup-micromamba@main - - - name: Create mamba environment - run: | - micromamba create -n a2 python=${{ matrix.python-version }} --yes + with: + environment-name: a2 + cache-environment: false + create-args: >- + python=${{ matrix.python-version }} - name: Install uv run: micromamba run -n a2 pip install uv @@ -91,7 +92,7 @@ jobs: - uses: codecov/codecov-action@v1 - if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2' + if: matrix.python-version == '3.11' && github.repository == 'materialsproject/atomate2' with: token: ${{ secrets.CODECOV_TOKEN }} name: coverage${{ matrix.split }} @@ -113,7 +114,7 @@ jobs: shell: bash -l {0} # enables conda/mamba env activation by reading bash profile strategy: matrix: - python-version: ["3.10","3.11","3.12"] + python-version: ["3.11","3.12"] steps: - name: Check out repo @@ -121,10 +122,11 @@ jobs: - name: Set up micromamba uses: mamba-org/setup-micromamba@main - - - name: Create mamba environment - run: | - micromamba create -n a2 python=${{ matrix.python-version }} --yes + with: + environment-name: a2 + cache-environment: false + create-args: >- + python=${{ matrix.python-version }} - name: Install uv run: micromamba run -n a2 pip install uv @@ -174,7 +176,7 @@ jobs: shell: bash -l {0} # enables conda/mamba env activation by reading bash profile strategy: matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.11", "3.12"] steps: - name: Check out repo @@ -182,10 +184,11 @@ jobs: - name: Set up micromamba uses: mamba-org/setup-micromamba@main - - - name: Create mamba environment - run: | - micromamba create -n a2 python=${{ matrix.python-version }} --yes + with: + environment-name: a2 + cache-environment: false + create-args: >- + python=${{ matrix.python-version }} - name: Install uv run: micromamba run -n a2 pip install uv @@ -220,7 +223,7 @@ jobs: pytest -n auto --splits 1 --group 1 --cov=atomate2 --cov-report=xml tests/ase - uses: codecov/codecov-action@v1 - if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2' + if: matrix.python-version == '3.11' && github.repository == 'materialsproject/atomate2' with: token: ${{ secrets.CODECOV_TOKEN }} file: ./coverage.xml @@ -241,7 +244,7 @@ jobs: shell: bash -l {0} # enables conda/mamba env activation by reading bash profile strategy: matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.11", "3.12"] steps: - name: Check out repo @@ -249,10 +252,11 @@ jobs: - name: Set up micromamba uses: mamba-org/setup-micromamba@main - - - name: Create mamba environment - run: | - micromamba create -n a2 python=${{ matrix.python-version }} --yes + with: + environment-name: a2 + cache-environment: false + create-args: >- + python=${{ matrix.python-version }} - name: Install uv run: micromamba run -n a2 pip install uv @@ -292,7 +296,7 @@ jobs: - uses: codecov/codecov-action@v1 - if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2' + if: matrix.python-version == '3.11' && github.repository == 'materialsproject/atomate2' with: token: ${{ secrets.CODECOV_TOKEN }} name: coverage @@ -308,7 +312,7 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" cache: pip cache-dependency-path: pyproject.toml diff --git a/.github/workflows/update-precommit.yml b/.github/workflows/update-precommit.yml index 6959999283..fecd66a849 100644 --- a/.github/workflows/update-precommit.yml +++ b/.github/workflows/update-precommit.yml @@ -15,7 +15,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" - name: Install pre-commit run: pip install pre-commit diff --git a/.gitignore b/.gitignore index 2500a708d8..e92ea67967 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,9 @@ docs_build/* docs/tutorials/* docs/tutorials/*/* +# tutorial output files +tutorials/*.json.gz + # C extensions *.so diff --git a/docs/user/install.md b/docs/user/install.md index fead021e70..8e48bf6321 100644 --- a/docs/user/install.md +++ b/docs/user/install.md @@ -354,19 +354,34 @@ and then return to this tutorial. ### Materials Project API key You can get an API key from the [Materials Project] by logging in and going to your -[Dashboard](materials project). Add this also to -your `~/.config/.pmgrc.yaml` so that it looks like the following +[Dashboard](materials project). -```yaml -PMG_VASP_PSP_DIR: <>/pps -PMG_MAPI_KEY: <> +Add this to your `.bashrc`, `.zshrc`, etc. as an environment variable: +```console +export MP_API_KEY=your_api_key_here ``` -You can generate this file and set these values using the `pymatgen` CLI: +## Emmet-core setup -```bash -pmg config --add PMG_VASP_PSP_DIR /abs/path/to/psp PMG_MAPI_KEY your_api_key +The `emmet-core` package is used to define data schemas for parsing outputs of workflows. +It is also used in building the Materials Project data, therefore its use in `atomate2` is to broadly ensure compatibility with the Materials Project's data structures. +`emmet-core` allows you to use either `pymatgen` or `emmet-core`-defined models for larger data objects, such as charge densities (`CHGCAR`, `AECCAR*`), or trajectories (relaxation, MD, etc.). +The `pymatgen` objects have long been the default in workflows, and are structured to be output as JSON files. +The `emmet-core` objects have been designed with both JSON and Apache parquet as storage formats. + +If you will be storing data in the cloud, or would like to use these newer data models which may use less storage, you can either add a line to your `atomate2.yaml` file: +```yaml +VASP_USE_EMMET_MODELS: true ``` +or use an environment variable in your `.bashrc`: +```console +export ATOMATE2_VASP_USE_EMMET_MODEL=true +``` +Note that there is an equivalent `emmet-core` setting, which can be set by the environment variable `EMMET_USE_EMMET_MODELS`. + +For ASE and machine learning forcefield jobs, you can use the `ASE_FORCEFIELD_USE_EMMET_MODELS` flag in `atomate2.yaml` to toggle the same functionality. + +The default in `atomate2` is to use `pymatgen` models. [materials project]: https://materialsproject.org/dashboard diff --git a/pyproject.toml b/pyproject.toml index 50e46261bf..cbba565764 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "PyYAML", "click", "custodian>=2024.4.18", - "emmet-core>=v0.84.10", + "emmet-core>=0.86.1", "jobflow>=0.1.11", "monty>=2024.12.10", "numpy", @@ -57,11 +57,10 @@ forcefields = [ "chgnet>=0.2.2", "mace-torch>=0.3.3", "matgl>=1.2.1", - "torchdata<=0.7.1", + "torchdata<=0.7.1", # TODO: remove when issue fixed # quippy-ase support for py3.12 tracked in https://github.com/libAtoms/QUIP/issues/645 "quippy-ase>=0.9.14; python_version < '3.12'", "sevenn>=0.9.3", - "torchdata<=0.7.1", # TODO: remove when issue fixed ] approxneb = ["pymatgen-analysis-diffusion>=2024.7.15"] ase = ["ase>=3.26.0"] @@ -107,7 +106,7 @@ strict-forcefields = [ "calorine==3.1; python_version < '3.12'", "chgnet==0.3.8", "mace-torch==0.3.14", - "matgl==1.3.0", + "matgl==2.0.2", "quippy-ase==0.9.14; python_version < '3.12'", "sevenn==0.10.4", "torch==2.2.0", @@ -115,6 +114,7 @@ strict-forcefields = [ ] strict = [ "atomate2[strict-forcefields, docs, cclib, phonons, lobster, openmm, mp, defects, ase, ase-ext]", + "numpy<2.0", ] [project.scripts] diff --git a/src/atomate2/abinit/schemas/calculation.py b/src/atomate2/abinit/schemas/calculation.py index 374dc6fa88..1373c7f195 100644 --- a/src/atomate2/abinit/schemas/calculation.py +++ b/src/atomate2/abinit/schemas/calculation.py @@ -11,11 +11,7 @@ from abipy.flowtk import events from abipy.flowtk.utils import File from emmet.core.math import Matrix3D, Vector3D - -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum from pydantic import BaseModel, Field from pymatgen.core import Structure from typing_extensions import Self diff --git a/src/atomate2/aims/schemas/calculation.py b/src/atomate2/aims/schemas/calculation.py index 880b317b79..b4c5dada14 100644 --- a/src/atomate2/aims/schemas/calculation.py +++ b/src/atomate2/aims/schemas/calculation.py @@ -12,11 +12,7 @@ import numpy as np from ase.spectrum.band_structure import BandStructure from emmet.core.math import Matrix3D, Vector3D - -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum from pydantic import BaseModel, Field from pymatgen.core import Molecule, Structure from pymatgen.core.trajectory import Trajectory diff --git a/src/atomate2/ase/jobs.py b/src/atomate2/ase/jobs.py index 458ad892a3..786516209f 100644 --- a/src/atomate2/ase/jobs.py +++ b/src/atomate2/ase/jobs.py @@ -8,15 +8,9 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING -from ase.io import Trajectory as AseTrajectory - -try: - from emmet.core.types.enums import StoreTrajectoryOption -except ImportError: - from emmet.core.vasp.calculation import StoreTrajectoryOption +from emmet.core.types.enums import StoreTrajectoryOption from jobflow import Maker, job from pymatgen.core import Molecule, Structure -from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor from atomate2.ase.schemas import AseResult, AseTaskDoc @@ -31,7 +25,7 @@ from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc -_ASE_DATA_OBJECTS = [PmgTrajectory, AseTrajectory] +_ASE_DATA_OBJECTS = ["trajectory"] @dataclass diff --git a/src/atomate2/ase/md.py b/src/atomate2/ase/md.py index 7c6abb6f54..670b5322e4 100644 --- a/src/atomate2/ase/md.py +++ b/src/atomate2/ase/md.py @@ -24,16 +24,13 @@ Stationary, ZeroRotation, ) - -try: - from emmet.core.types.enums import StoreTrajectoryOption -except ImportError: - from emmet.core.vasp.calculation import StoreTrajectoryOption +from emmet.core.types.enums import StoreTrajectoryOption from jobflow import job from pymatgen.core.structure import Molecule, Structure from pymatgen.io.ase import AseAtomsAdaptor from scipy.interpolate import interp1d +from atomate2 import SETTINGS from atomate2.ase.jobs import _ASE_DATA_OBJECTS, AseMaker from atomate2.ase.schemas import AseResult, AseTaskDoc from atomate2.ase.utils import TrajectoryObserver @@ -165,6 +162,9 @@ class AseMDMaker(AseMaker, ABC): Whether to initialize the atomic velocities with zero angular momentum verbose : bool = False Whether to print stdout to screen during the MD run. + use_emmet_models : bool = False + Whether to use emmet-core (True) or pymatgen (False) + data models for larger objects, e.g., trajectories. """ name: str = "ASE MD" @@ -185,6 +185,7 @@ class AseMDMaker(AseMaker, ABC): zero_linear_momentum: bool = False zero_angular_momentum: bool = False verbose: bool = False + use_emmet_models: bool = SETTINGS.ASE_FORCEFIELD_USE_EMMET_MODELS def __post_init__(self) -> None: """Ensure that ensemble is an enum.""" @@ -433,7 +434,12 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: return AseResult( final_mol_or_struct=mol_or_struct, - trajectory=md_observer.to_emmet_trajectory(filename=None), + trajectory=getattr( + md_observer, + "to_emmet_trajectory" + if self.use_emmet_models + else "to_pymatgen_trajectory", + )(filename=None), dir_name=os.getcwd(), elapsed_time=t_f - t_i, ) diff --git a/src/atomate2/ase/schemas.py b/src/atomate2/ase/schemas.py index 84a0782d3f..8652b8b678 100644 --- a/src/atomate2/ase/schemas.py +++ b/src/atomate2/ase/schemas.py @@ -11,20 +11,22 @@ from __future__ import annotations from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any +import numpy as np +from ase.stress import voigt_6_to_full_3x3_stress +from ase.units import GPa from emmet.core.math import Matrix3D, Vector3D from emmet.core.structure import MoleculeMetadata, StructureMetadata from emmet.core.trajectory import AtomTrajectory - -try: - from emmet.core.types.enums import StoreTrajectoryOption, TaskState, ValueEnum -except ImportError: - from emmet.core.tasks import TaskState - from emmet.core.utils import ValueEnum - from emmet.core.vasp.calculation import StoreTrajectoryOption +from emmet.core.types.enums import StoreTrajectoryOption, TaskState, ValueEnum from pydantic import BaseModel, Field from pymatgen.core import Molecule, Structure +from pymatgen.core.trajectory import Trajectory as PmgTrajectory +from pymatgen.entries.computed_entries import ComputedEntry + +if TYPE_CHECKING: + from emmet.core.math import Matrix3D, Vector6D _task_doc_translation_keys = { "input", @@ -40,6 +42,18 @@ } +def convert_stress_from_voigt_to_symm(voigt: Vector6D) -> Matrix3D: + """Convert Voigt representation stress in GPa to 3x3 tensor in kilobar. + + Converts stress units from eV/A³ to kBar (* -1 from standard output) + and to 3x3 matrix to comply with MP convention. + + voigt : Vector6D + The Voigt representation of the stress in eV/A³ + """ + return tuple(voigt_6_to_full_3x3_stress(np.array(voigt) * -10 / GPa).tolist()) + + class AseResult(BaseModel): """Schema to store outputs in AseTaskDocument.""" @@ -51,7 +65,7 @@ class AseResult(BaseModel): None, description="The final total energy from the calculation." ) - trajectory: AtomTrajectory | None = Field( + trajectory: PmgTrajectory | AtomTrajectory | None = Field( None, description="The relaxation or molecular dynamics trajectory." ) @@ -244,6 +258,10 @@ class AseStructureTaskDoc(StructureMetadata): tags: list[str] | None = Field(None, description="List of tags for the task.") + entry: ComputedEntry | None = Field( + None, description="The computed entry summarizing this calculation." + ) + @classmethod def from_ase_task_doc( cls, ase_task_doc: AseTaskDoc, **task_document_kwargs @@ -261,6 +279,11 @@ def from_ase_task_doc( {k: getattr(ase_task_doc, k) for k in _task_doc_translation_keys}, structure=ase_task_doc.mol_or_struct, ) + if not task_document_kwargs.get("entry"): + task_document_kwargs["entry"] = ComputedEntry( + composition=ase_task_doc.mol_or_struct.composition, + energy=ase_task_doc.output.energy, + ) return cls.from_structure( meta_structure=ase_task_doc.mol_or_struct, **task_document_kwargs ) @@ -419,7 +442,12 @@ def from_ase_compatible_result( input_mol_or_struct = None if trajectory: n_steps = len(trajectory) - input_mol_or_struct = trajectory.to_pmg(frame_props=tuple(), indices=0)[0] + if isinstance(trajectory, AtomTrajectory): + input_mol_or_struct = trajectory.to_pmg(frame_props=tuple(), indices=0)[ + 0 + ] + else: + input_mol_or_struct = trajectory[0] input_doc = InputDoc( mol_or_struct=input_mol_or_struct, @@ -438,7 +466,7 @@ def from_ase_compatible_result( n_steps = 1 if trajectory: - trajectory = trajectory[-1] + trajectory = trajectory[-1:] output_mol_or_struct = input_mol_or_struct else: output_mol_or_struct = result.final_mol_or_struct @@ -449,38 +477,68 @@ def from_ase_compatible_result( ionic_steps = None if trajectory: - final_energy = trajectory.energy[-1] - final_forces = trajectory.forces[-1] - ionic_step_props = ["energy", "forces"] - if trajectory.stress: - final_stress = trajectory.stress[-1] - ionic_step_props.append("stress") + ionic_step_props = {"energy", "forces"} + if isinstance(trajectory, AtomTrajectory): + final_energy = trajectory.energy[-1] + final_forces = trajectory.forces[-1] + + if trajectory.stress: + final_stress = trajectory.stress[-1] + ionic_step_props.add("stress") + + if trajectory.magmoms: + ionic_step_props.add("magmoms") + else: + final_energy = trajectory.frame_properties[-1]["energy"] + final_forces = trajectory.frame_properties[-1]["forces"] + if any(frame.get("stress") for frame in trajectory.frame_properties): + final_stress = convert_stress_from_voigt_to_symm( + trajectory.frame_properties[-1]["stress"] + ) + ionic_step_props.add("stress") - if trajectory.magmoms: - ionic_step_props.append("magmoms") + if any(frame.get("magmoms") for frame in trajectory.frame_properties): + ionic_step_props.add("magmoms") ionic_steps = [] - if ionic_step_data is not None and len(ionic_step_data) > 0: - for idx in range(n_steps): - _ionic_step_data = { - key: ( - getattr(trajectory, key)[idx] - if key in ionic_step_data - else None - ) - for key in ionic_step_props - } - - current_mol_or_struct = trajectory.to_pmg( - frame_props=tuple(), indices=-1 - )[0] - - ionic_step = IonicStep( - mol_or_struct=current_mol_or_struct, - **_ionic_step_data, + if ( + len( + use_ionic_step_props := ionic_step_props.intersection( + ionic_step_data or set() ) - - ionic_steps.append(ionic_step) + ) + > 0 + ): + if isinstance(trajectory, AtomTrajectory): + ionic_steps = [ + IonicStep( + mol_or_struct=trajectory.to_pmg( + frame_props=tuple(), + indices=idx, + )[0], + **{ + key: getattr(trajectory, key)[idx] + for key in use_ionic_step_props + }, + ) + for idx in range(n_steps) + ] + + else: + ionic_steps = [ + IonicStep( + mol_or_struct=atoms, + **{ + key: convert_stress_from_voigt_to_symm( + trajectory.frame_properties[idx].get(key) + ) + if key == "stress" + else trajectory.frame_properties[idx].get(key) + for key in use_ionic_step_props + }, + ) + for idx, atoms in enumerate(trajectory) + ] objects: dict[AseObject, Any] = {} if store_trajectory != StoreTrajectoryOption.NO: diff --git a/src/atomate2/ase/utils.py b/src/atomate2/ase/utils.py index c2eef68ddf..63dedefdfe 100644 --- a/src/atomate2/ase/utils.py +++ b/src/atomate2/ase/utils.py @@ -22,8 +22,6 @@ from ase.mep.neb import NEB from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG -from ase.stress import voigt_6_to_full_3x3_stress -from ase.units import GPa from emmet.core.neb import NebMethod, NebResult from emmet.core.trajectory import AtomTrajectory from monty.serialization import dumpfn @@ -31,7 +29,8 @@ from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor -from atomate2.ase.schemas import AseResult +from atomate2 import SETTINGS +from atomate2.ase.schemas import AseResult, convert_stress_from_voigt_to_symm if TYPE_CHECKING: from os import PathLike @@ -300,10 +299,8 @@ def to_emmet_trajectory( ionic_step_data = {v: getattr(self, k) for k, v in frame_props.items()} if self._calc_kwargs["stresses"]: - # NOTE: convert stress units from eV/A³ to kBar (* -1 from standard output) - # and to 3x3 matrix to comply with MP convention ionic_step_data["stress"] = [ - voigt_6_to_full_3x3_stress(val * -10 / GPa) for val in self.stresses + convert_stress_from_voigt_to_symm(val) for val in self.stresses ] traj = AtomTrajectory( @@ -390,6 +387,7 @@ def relax( verbose: bool = False, cell_filter: Filter = FrechetCellFilter, filter_kwargs: dict | None = None, + use_emmet_models: bool = SETTINGS.ASE_FORCEFIELD_USE_EMMET_MODELS, **kwargs, ) -> AseResult: """ @@ -413,6 +411,9 @@ def relax( The step interval for saving the trajectories. verbose : bool If True, screen output will be shown. + use_emmet_models : bool + Whether to use emmet-core (True) or pymatgen (False) + data models for larger objects, e.g., trajectories. **kwargs Further kwargs. @@ -453,11 +454,17 @@ def relax( struct = self.ase_adaptor.get_structure( atoms, cls=Molecule if is_mol else Structure ) - traj = obs.to_emmet_trajectory(filename=None) - is_force_conv = all( - np.linalg.norm(traj.forces[-1][idx]) < abs(fmax) - for idx in range(len(struct)) - ) + + if use_emmet_models: + traj = obs.to_emmet_trajectory(filename=None) + final_forces = traj.forces[-1] + first_last_energy = [traj.energy[idx] for idx in (0, -1)] + else: + traj = obs.to_pymatgen_trajectory(filename=None, file_format="pmg") + final_forces = traj.frame_properties[-1]["forces"] + first_last_energy = [ + traj.frame_properties[idx]["energy"] for idx in (0, -1) + ] if final_atoms_object_file is not None: if steps <= 1: @@ -489,8 +496,10 @@ def relax( final_mol_or_struct=struct, trajectory=traj, converged=converged, - is_force_converged=is_force_conv, - energy_downhill=traj.energy[-1] < traj.energy[0], + is_force_converged=np.all( + np.linalg.norm(np.array(final_forces), axis=1) < abs(fmax) + ), + energy_downhill=first_last_energy[-1] < first_last_energy[0], dir_name=os.getcwd(), elapsed_time=t_f - t_i, ) diff --git a/src/atomate2/common/jobs/approx_neb.py b/src/atomate2/common/jobs/approx_neb.py index 00c7bb557a..40a7314082 100644 --- a/src/atomate2/common/jobs/approx_neb.py +++ b/src/atomate2/common/jobs/approx_neb.py @@ -6,11 +6,7 @@ import numpy as np from emmet.core.neb import HopFailureReason, NebMethod, NebPathwayResult, NebResult - -try: - from emmet.core.types.enums import TaskState -except ImportError: - from emmet.core.tasks import TaskState +from emmet.core.types.enums import TaskState from jobflow import Flow, Response, job from pymatgen.analysis.diffusion.neb.pathfinder import ChgcarPotential, NEBPathfinder from pymatgen.core import Element diff --git a/src/atomate2/common/jobs/electrode.py b/src/atomate2/common/jobs/electrode.py index e5192a7a63..e923776c50 100644 --- a/src/atomate2/common/jobs/electrode.py +++ b/src/atomate2/common/jobs/electrode.py @@ -6,18 +6,13 @@ from typing import TYPE_CHECKING, NamedTuple from emmet.core.electrode import InsertionElectrodeDoc -from emmet.core.mpid import MPID, check_ulid +from emmet.core.mpid import MPID, AlphaID, check_ulid from emmet.core.structure_group import StructureGroupDoc from jobflow import Flow, Maker, Response, job from pymatgen.analysis.defects.generators import ChargeInterstitialGenerator from pymatgen.entries.computed_entries import ComputedStructureEntry from ulid import ULID -try: - from emmet.core.mpid import AlphaID -except ImportError: - AlphaID = None - if TYPE_CHECKING: from collections.abc import Callable from pathlib import Path diff --git a/src/atomate2/cp2k/run.py b/src/atomate2/cp2k/run.py index e15cac28ca..176d4e57c2 100644 --- a/src/atomate2/cp2k/run.py +++ b/src/atomate2/cp2k/run.py @@ -21,11 +21,7 @@ ) from custodian.cp2k.jobs import Cp2kJob from custodian.cp2k.validators import Cp2kOutputValidator - -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum from atomate2 import SETTINGS diff --git a/src/atomate2/cp2k/schemas/calc_types/enums.py b/src/atomate2/cp2k/schemas/calc_types/enums.py index 47459c82be..55757b2d4d 100644 --- a/src/atomate2/cp2k/schemas/calc_types/enums.py +++ b/src/atomate2/cp2k/schemas/calc_types/enums.py @@ -4,10 +4,7 @@ Do not edit this by hand. Edit generate.py or run_types.yaml instead. """ -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum class RunType(ValueEnum): diff --git a/src/atomate2/cp2k/schemas/calculation.py b/src/atomate2/cp2k/schemas/calculation.py index c080f318c0..05ed454f84 100644 --- a/src/atomate2/cp2k/schemas/calculation.py +++ b/src/atomate2/cp2k/schemas/calculation.py @@ -7,10 +7,7 @@ from shutil import which from typing import Any, Union -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum from pydantic import BaseModel, Field, field_validator from pymatgen.command_line.bader_caller import BaderAnalysis from pymatgen.core.structure import Molecule, Structure diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 50e3890935..770873d26c 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -4,10 +4,7 @@ from typing import Any -try: - from emmet.core.types.enums import StoreTrajectoryOption -except ImportError: - from emmet.core.vasp.calculation import StoreTrajectoryOption +from emmet.core.types.enums import StoreTrajectoryOption from pydantic import Field from atomate2.ase.schemas import AseObject, AseResult, AseStructureTaskDoc, AseTaskDoc diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 32980504d1..393f78fcb6 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -9,11 +9,9 @@ from pathlib import Path from typing import TYPE_CHECKING -from ase.io import Trajectory as AseTrajectory from ase.units import Bohr from ase.units import GPa as _GPa_to_eV_per_A3 from monty.json import MontyDecoder -from pymatgen.core.trajectory import Trajectory as PmgTrajectory if TYPE_CHECKING: from collections.abc import Generator @@ -23,7 +21,7 @@ from atomate2.ase.schemas import AseResult -_FORCEFIELD_DATA_OBJECTS = [PmgTrajectory, AseTrajectory, "ionic_steps"] +_FORCEFIELD_DATA_OBJECTS = ["trajectory", "ionic_steps"] class MLFF(Enum): # TODO inherit from StrEnum when 3.11+ diff --git a/src/atomate2/lobster/run.py b/src/atomate2/lobster/run.py index 4c16d35922..80ab8f28fd 100644 --- a/src/atomate2/lobster/run.py +++ b/src/atomate2/lobster/run.py @@ -11,11 +11,7 @@ from custodian import Custodian from custodian.lobster.handlers import EnoughBandsValidator, LobsterFilesValidator from custodian.lobster.jobs import LobsterJob - -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum from atomate2 import SETTINGS diff --git a/src/atomate2/openff/core.py b/src/atomate2/openff/core.py index 38fdf3aac2..4b9385c3b8 100644 --- a/src/atomate2/openff/core.py +++ b/src/atomate2/openff/core.py @@ -7,11 +7,7 @@ import openff.toolkit as tk from emmet.core.openff import ClassicalMDTaskDocument, MoleculeSpec - -try: - from emmet.core.types.enums import TaskState -except ImportError: - from emmet.core.tasks import TaskState +from emmet.core.types.enums import TaskState from jobflow import Response, job from openff.interchange import Interchange from openff.interchange.components._packmol import pack_box diff --git a/src/atomate2/openmm/jobs/generate.py b/src/atomate2/openmm/jobs/generate.py index dea2aaebca..55890307a6 100644 --- a/src/atomate2/openmm/jobs/generate.py +++ b/src/atomate2/openmm/jobs/generate.py @@ -13,11 +13,7 @@ import numpy as np from emmet.core.openff import MoleculeSpec from emmet.core.openmm import OpenMMTaskDocument - -try: - from emmet.core.types.enums import TaskState -except ImportError: - from emmet.core.tasks import TaskState +from emmet.core.types.enums import TaskState from jobflow import Response from openmm import Context, LangevinMiddleIntegrator, System, XmlSerializer from openmm.app import PME, ForceField diff --git a/src/atomate2/qchem/run.py b/src/atomate2/qchem/run.py index 060b00c4af..d48f8a4a33 100644 --- a/src/atomate2/qchem/run.py +++ b/src/atomate2/qchem/run.py @@ -10,11 +10,8 @@ from custodian import Custodian from custodian.qchem.handlers import QChemErrorHandler from custodian.qchem.jobs import QCJob +from emmet.core.types.enums import ValueEnum -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum from atomate2 import SETTINGS if TYPE_CHECKING: diff --git a/src/atomate2/settings.py b/src/atomate2/settings.py index 69e9669bba..84b146fd25 100644 --- a/src/atomate2/settings.py +++ b/src/atomate2/settings.py @@ -114,6 +114,12 @@ class Atomate2Settings(BaseSettings): "VaspInputGenerator.", ) + VASP_USE_EMMET_MODELS: bool = Field( + default=False, + description="Whether to use emmet models (True) for VASP electronic " + "structure data, or the pymatgen models (False).", + ) + LOBSTER_CMD: str = Field( default="lobster", description="Command to run standard version of VASP." ) @@ -212,6 +218,12 @@ class Atomate2Settings(BaseSettings): 5, description="Maximum number of restarts of a job." ) + ASE_FORCEFIELD_USE_EMMET_MODELS: bool = Field( + default=False, + description="Whether to use emmet-core models (False) or pymatgen (True) " + "models for larger data objects, such as trajectories.", + ) + model_config = SettingsConfigDict(env_prefix=_ENV_PREFIX) # QChem specific settings diff --git a/src/atomate2/utils/testing/vasp.py b/src/atomate2/utils/testing/vasp.py index ba89a354aa..135433646f 100644 --- a/src/atomate2/utils/testing/vasp.py +++ b/src/atomate2/utils/testing/vasp.py @@ -193,7 +193,6 @@ def fake_run_vasp( if clear_inputs: _clear_vasp_inputs() - _copy_vasp_outputs(ref_path) # pretend to run VASP by copying pre-generated outputs from reference dir diff --git a/src/atomate2/vasp/files.py b/src/atomate2/vasp/files.py index d24cefd5c7..754928db70 100644 --- a/src/atomate2/vasp/files.py +++ b/src/atomate2/vasp/files.py @@ -64,7 +64,6 @@ def copy_vasp_outputs( A file client to use for performing file operations. """ src_dir = strip_hostname(src_dir) # TODO: Handle hostnames properly. - logger.info(f"Copying VASP inputs from {src_dir}") relax_ext = get_largest_relax_extension(src_dir, src_host, file_client=file_client) @@ -196,7 +195,7 @@ def write_vasp_input_set( if clean_prev: # remove previous inputs (prevents old KPOINTS file from overriding KSPACING) - for filename in ("POSCAR", "KPOINTS", "POTCAR", "INCAR"): + for filename in ("POSCAR", "KPOINTS", "POTCAR", "POTCAR.spec", "INCAR"): if Path(filename).exists(): Path(filename).unlink() diff --git a/src/atomate2/vasp/flows/core.py b/src/atomate2/vasp/flows/core.py index eda865d205..55c24400c2 100644 --- a/src/atomate2/vasp/flows/core.py +++ b/src/atomate2/vasp/flows/core.py @@ -6,10 +6,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING -try: - from emmet.core.types.enums import VaspObject -except ImportError: - from emmet.core.vasp.calculation import VaspObject +from emmet.core.types.enums import VaspObject from jobflow import Flow, Maker from atomate2.vasp.jobs.core import ( diff --git a/src/atomate2/vasp/flows/elph.py b/src/atomate2/vasp/flows/elph.py index 6f553984f1..72fa5220d1 100644 --- a/src/atomate2/vasp/flows/elph.py +++ b/src/atomate2/vasp/flows/elph.py @@ -116,8 +116,8 @@ class ElectronPhononMaker(Maker): user_incar_settings={"LORBIT": 10}, # disable site projections ), task_document_kwargs={ - "strip_bandstructure_projections": True, - "strip_dos_projections": True, + "strip_bandstructure_projections": False, + "strip_dos_projections": False, }, ), ) diff --git a/src/atomate2/vasp/jobs/base.py b/src/atomate2/vasp/jobs/base.py index 20ed848f31..af82fdba81 100644 --- a/src/atomate2/vasp/jobs/base.py +++ b/src/atomate2/vasp/jobs/base.py @@ -10,9 +10,10 @@ from emmet.core.neb import NebIntermediateImagesDoc from emmet.core.tasks import TaskDoc +from emmet.core.types.enums import VaspObject from jobflow import Maker, Response, job from monty.serialization import dumpfn -from pymatgen.core.trajectory import Trajectory +from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.electronic_structure.bandstructure import ( BandStructure, BandStructureSymmLine, @@ -40,20 +41,31 @@ ) _DATA_OBJECTS = [ - BandStructure, - BandStructureSymmLine, - DOS, - Dos, - CompleteDos, - Locpot, - Chgcar, Wavecar, - Trajectory, "force_constants", "normalmode_eigenvecs", "bandstructure", # FIX: BandStructure is not currently MSONable ] +if SETTINGS.VASP_USE_EMMET_MODELS: + # Because the emmet-core models deserialize to JSON + # on model_dump, we just pass field names here, not object types + _DATA_OBJECTS.extend([f.value for f in VaspObject]) +else: + # Store pymatgen objects + _DATA_OBJECTS.extend( + [ + BandStructure, + BandStructureSymmLine, + DOS, + Dos, + CompleteDos, + Locpot, + Chgcar, + PmgTrajectory, + ] + ) + # Input files. Partially from https://www.vasp.at/wiki/index.php/Category:Input_files # Exclude those that are also outputs _INPUT_FILES = [ @@ -318,4 +330,7 @@ def get_vasp_task_document( stacklevel=1, ) + kwargs["use_emmet_models"] = kwargs.get( + "use_emmet_models", SETTINGS.VASP_USE_EMMET_MODELS + ) return TaskDoc.from_directory(path, **kwargs) diff --git a/src/atomate2/vasp/jobs/elph.py b/src/atomate2/vasp/jobs/elph.py index 76cb9781a0..2309201b03 100644 --- a/src/atomate2/vasp/jobs/elph.py +++ b/src/atomate2/vasp/jobs/elph.py @@ -7,8 +7,10 @@ from typing import TYPE_CHECKING import numpy as np +from emmet.core.band_theory import ElectronicBS from jobflow import Flow, Response, job +from atomate2 import SETTINGS from atomate2.vasp.jobs.base import BaseVaspMaker, vasp_job from atomate2.vasp.jobs.core import TransmuterMaker from atomate2.vasp.schemas.elph import ElectronPhononRenormalisationDoc @@ -20,7 +22,6 @@ from pymatgen.core import Structure from pymatgen.electronic_structure.bandstructure import BandStructure - DEFAULT_ELPH_TEMPERATURES = (0, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000) DEFAULT_MIN_SUPERCELL_LENGTH = 15 logger = logging.getLogger(__name__) @@ -188,11 +189,11 @@ def run_elph_displacements( @job(output_schema=ElectronPhononRenormalisationDoc) def calculate_electron_phonon_renormalisation( temperatures: list[float], - displacement_band_structures: list[BandStructure], + displacement_band_structures: list[BandStructure | ElectronicBS], displacement_structures: list[Structure], displacement_uuids: list[str], displacement_dirs: list[str], - bulk_band_structure: BandStructure, + bulk_band_structure: BandStructure | ElectronicBS, bulk_structure: Structure, bulk_uuid: str, bulk_dir: str, @@ -207,7 +208,8 @@ def calculate_electron_phonon_renormalisation( ---------- temperatures : list of float The temperatures at which electron phonon properties were calculated. - displacement_band_structures : list of BandStructure + displacement_band_structures : + list of pymatgen .BandStructure or emmet.core .ElectronicBS The electron-phonon displaced band structures. displacement_structures : list of Structure The electron-phonon displaced structures. @@ -216,7 +218,7 @@ def calculate_electron_phonon_renormalisation( displacement_dirs : list of str The calculation directories of the electron-phonon displaced band structure calculations. - bulk_band_structure : BandStructure + bulk_band_structure : pymatgen .BandStructure or emmet.core .ElectronicBS The band structure of the bulk undisplaced supercell calculation. bulk_structure : Structure The structure of the bulk undisplaced supercell. @@ -242,7 +244,17 @@ def calculate_electron_phonon_renormalisation( # filter band structures that are None (i.e., the displacement calculation failed) keep = [idx for idx, b in enumerate(displacement_band_structures) if b is not None] temperatures = [temperatures[i] for i in keep] - displacement_band_structures = [displacement_band_structures[i] for i in keep] + + if SETTINGS.VASP_USE_EMMET_MODELS: + # Convert back to pymatgen band structures + displacement_band_structures = [ + ElectronicBS(**displacement_band_structures[i]).to_pmg() for i in keep + ] + bulk_bs = ElectronicBS(**bulk_band_structure).to_pmg() + else: + displacement_band_structures = [displacement_band_structures[i] for i in keep] + bulk_bs = bulk_band_structure + displacement_structures = [displacement_structures[i] for i in keep] displacement_uuids = [displacement_uuids[i] for i in keep] displacement_dirs = [displacement_dirs[i] for i in keep] @@ -255,7 +267,7 @@ def calculate_electron_phonon_renormalisation( displacement_structures, displacement_uuids, displacement_dirs, - bulk_band_structure, + bulk_bs, bulk_structure, bulk_uuid, bulk_dir, diff --git a/src/atomate2/vasp/jobs/md.py b/src/atomate2/vasp/jobs/md.py index f7921a5fc7..24fbb10e23 100644 --- a/src/atomate2/vasp/jobs/md.py +++ b/src/atomate2/vasp/jobs/md.py @@ -15,11 +15,7 @@ StdErrHandler, VaspErrorHandler, ) - -try: - from emmet.core.types.enums import StoreTrajectoryOption -except ImportError: - from emmet.core.vasp.calculation import StoreTrajectoryOption +from emmet.core.types.enums import StoreTrajectoryOption from jobflow import Response, job from atomate2.vasp.jobs.base import BaseVaspMaker diff --git a/src/atomate2/vasp/run.py b/src/atomate2/vasp/run.py index 77455ecbf9..abc45e95a4 100644 --- a/src/atomate2/vasp/run.py +++ b/src/atomate2/vasp/run.py @@ -27,11 +27,7 @@ ) from custodian.vasp.jobs import VaspJob, VaspNEBJob from custodian.vasp.validators import VaspFilesValidator, VasprunXMLValidator - -try: - from emmet.core.types.enums import ValueEnum -except ImportError: - from emmet.core.utils import ValueEnum +from emmet.core.types.enums import ValueEnum from atomate2 import SETTINGS diff --git a/tests/ase/test_utils.py b/tests/ase/test_utils.py index 3ab97417c6..786a10f8e2 100644 --- a/tests/ase/test_utils.py +++ b/tests/ase/test_utils.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from pymatgen.core import Structure +from atomate2.ase.schemas import convert_stress_from_voigt_to_symm from atomate2.ase.utils import AseRelaxer, TrajectoryObserver @@ -51,13 +52,12 @@ def test_trajectory_observer(si_structure: Structure, test_dir, tmp_dir): @pytest.mark.parametrize( - ("optimizer", "traj_file"), - [("BFGS", None), (None, None), (BFGS, "log_file.traj")], + ("optimizer", "traj_file", "use_emmet_models"), + [("BFGS", None, True), (None, None, False), (BFGS, "log_file.traj", False)], ) -def test_relaxer(si_structure, test_dir, tmp_dir, optimizer, traj_file): - from ase.stress import voigt_6_to_full_3x3_stress - from ase.units import GPa - +def test_relaxer( + si_structure, test_dir, tmp_dir, optimizer, traj_file, use_emmet_models +): expected_lattice = { "a": 3.866974, "b": 3.866974, @@ -83,10 +83,15 @@ def test_relaxer(si_structure, test_dir, tmp_dir, optimizer, traj_file): AseRelaxer(calculator=LennardJones(), optimizer=optimizer) return - relaxer = AseRelaxer(calculator=LennardJones(), optimizer=optimizer) + relaxer = AseRelaxer( + calculator=LennardJones(), + optimizer=optimizer, + ) try: - relax_output = relaxer.relax(atoms=si_structure, traj_file=traj_file) + relax_output = relaxer.relax( + atoms=si_structure, traj_file=traj_file, use_emmet_models=use_emmet_models + ) except TypeError: return @@ -95,17 +100,29 @@ def test_relaxer(si_structure, test_dir, tmp_dir, optimizer, traj_file): for key in expected_lattice } == pytest.approx(expected_lattice) - assert relax_output.trajectory.energy[-1] == pytest.approx(expected_energy) + if use_emmet_models: + final_frame_attrs = { + k: getattr(relax_output.trajectory, k)[-1] + for k in ("energy", "forces", "stress") + } + else: + final_frame_attrs = { + k: relax_output.trajectory.frame_properties[-1].get(k) + for k in ("energy", "forces", "stress") + } + assert final_frame_attrs["energy"] == pytest.approx(expected_energy) assert_allclose( - relax_output["trajectory"].forces[-1], + final_frame_attrs["forces"], expected_forces, atol=1e-11, ) assert_allclose( - relax_output["trajectory"].stress[-1], - voigt_6_to_full_3x3_stress(expected_stresses) * -10 / GPa, + final_frame_attrs["stress"], + convert_stress_from_voigt_to_symm(expected_stresses) + if use_emmet_models + else expected_stresses, atol=1e-11, ) diff --git a/tests/conftest.py b/tests/conftest.py index 798cfcb3dc..1439a27420 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -43,7 +43,7 @@ def clean_dir(debug_mode): os.chdir(new_path) yield if debug_mode: - print(f"Tests ran in {new_path}") # noqa: T201 + initialize_logger().log(f"Tests ran in {new_path}") else: os.chdir(old_cwd) shutil.rmtree(new_path) diff --git a/tests/forcefields/flows/test_approx_neb.py b/tests/forcefields/flows/test_approx_neb.py index 10fd50b2fe..910acaa470 100644 --- a/tests/forcefields/flows/test_approx_neb.py +++ b/tests/forcefields/flows/test_approx_neb.py @@ -33,13 +33,13 @@ def test_approx_neb_from_endpoints(test_dir, clean_dir): output["collate_images_single_hop"].energies[i] == pytest.approx(energy) for i, energy in enumerate( [ - -1559.5150146484375, - -1554.3154296875, - -1529.771484375, - -1525.8846435546875, - -1533.1453857421875, - -1554.561767578125, - -1559.5150146484375, + -1558.1566162109375, + -1552.53369140625, + -1518.686767578125, + -1534.1644287109375, + -1523.787109375, + -1552.8035888671875, + -1558.1566162109375, ] ) ) diff --git a/tests/forcefields/flows/test_mpmorph.py b/tests/forcefields/flows/test_mpmorph.py index b03c78f7d2..401637e508 100644 --- a/tests/forcefields/flows/test_mpmorph.py +++ b/tests/forcefields/flows/test_mpmorph.py @@ -3,6 +3,7 @@ import pytest from jobflow import run_locally +from atomate2 import SETTINGS from atomate2.forcefields.flows.mpmorph import ( FastQuenchMLFFMDMaker, MPMorphMLFFMDMaker, @@ -162,14 +163,22 @@ def test_mpmorph_mlff_maker(ff_name, si_structure, test_dir, clean_dir): # this is designed to check if MD is injected with approximately # the right temperature, and that's why tolerance is so high assert all( - doc.forcefield_objects["trajectory"].temperature[0] + ( + doc.forcefield_objects["trajectory"].temperature[0] + if SETTINGS.ASE_FORCEFIELD_USE_EMMET_MODELS + else doc.forcefield_objects["trajectory"].frame_properties[0]["temperature"] + ) == pytest.approx(temp, abs=50) for name, doc in task_docs.items() if "MD Maker" in name ) - assert task_docs["production run"].forcefield_objects["trajectory"].temperature[ - 0 - ] == pytest.approx(temp, abs=50) + assert ( + task_docs["production run"].forcefield_objects["trajectory"].temperature[0] + if SETTINGS.ASE_FORCEFIELD_USE_EMMET_MODELS + else task_docs["production run"] + .forcefield_objects["trajectory"] + .frame_properties[0]["temperature"] + ) == pytest.approx(temp, abs=50) # check that MD Maker Energies are close # TODO: This may be unnecessary because it changes from model to model @@ -211,7 +220,13 @@ def test_mpmorph_mlff_maker(ff_name, si_structure, test_dir, clean_dir): assert all( any( - doc.forcefield_objects["trajectory"].temperature[0] + ( + doc.forcefield_objects["trajectory"].temperature[0] + if SETTINGS.ASE_FORCEFIELD_USE_EMMET_MODELS + else doc.forcefield_objects["trajectory"].frame_properties[0][ + "temperature" + ] + ) == pytest.approx(T, abs=100) for name, doc in task_docs.items() if "K" in name diff --git a/tests/forcefields/flows/test_qha.py b/tests/forcefields/flows/test_qha.py index 69499f16fc..adbdd1b754 100644 --- a/tests/forcefields/flows/test_qha.py +++ b/tests/forcefields/flows/test_qha.py @@ -63,8 +63,6 @@ def test_qha_dir_change_defaults(clean_dir, si_structure: Structure, tmp_path: P # run the flow or job and ensure that it finished running successfully responses = run_locally(flow, create_folders=True, ensure_success=True) - # print(responses) - # # validate the outputs ph_bs_dos_doc = responses[flow[-1].uuid][1].output assert isinstance(ph_bs_dos_doc, PhononQHADoc) @@ -95,8 +93,6 @@ def test_qha_dir_manual_supercell(clean_dir, si_structure: Structure, tmp_path: # run the flow or job and ensure that it finished running successfully responses = run_locally(flow, create_folders=True, ensure_success=True) - # print(responses) - # # validate the outputs ph_bs_dos_doc = responses[flow[-1].uuid][1].output assert isinstance(ph_bs_dos_doc, PhononQHADoc) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 49c8a8977b..c9194b8ac0 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -532,8 +532,8 @@ def test_matpes_relax_makers( refs = { "PBE": { - "energy_per_atom": -7.969418334960937, - "volume": 61.5685434322787, + "energy_per_atom": -7.9611351013183596, + "volume": 60.91639399282195, "forces": [ [ -1.48095100627188e-08, @@ -554,14 +554,14 @@ def test_matpes_relax_makers( ], ], "stress": [ - [-3.2316584962663115, -5.655957247906253e-07, -1.2974634469118903e-06], - [-5.655957247906253e-07, -3.2316376062607346, -3.4183954413422158e-06], - [-1.2974634469118903e-06, -3.4183954413422158e-06, -3.23162268482818], + [6.150300775936876, -5.854866356979066e-07, -6.522582661838942e-06], + [-5.854866356979066e-07, 6.150316070405244, -3.0104131342606253e-06], + [-6.522582661838942e-06, -3.0104131342606253e-06, 6.150302268080131], ], }, "r2SCAN": { - "energy_per_atom": -12.618433380126953, - "volume": 59.608148084043876, + "energy_per_atom": -12.588912963867188, + "volume": 59.30895984045571, "forces": [ [1.1260409849001007e-07, 1.4873557496741796e-08, 6.234344596123265e-09], [ @@ -582,9 +582,9 @@ def test_matpes_relax_makers( [-7.171183824539185e-08, 3.3614934835668464e-08, 9.266178579991902e-08], ], "stress": [ - [11.739250544520468, 9.398965578128745e-07, -3.1735166915189553e-06], - [9.398965578128745e-07, 11.739264719881394, -1.220071863449669e-06], - [-3.1735166915189553e-06, -1.220071863449669e-06, 11.739276657027439], + [12.034191310755238, -1.21893513832506e-06, -6.9246067896272225e-06], + [-1.21893513832506e-06, 12.03422712219337, -8.57680763083222e-06], + [-6.9246067896272225e-06, -8.57680763083222e-06, 12.03421369290407], ], }, } @@ -602,7 +602,7 @@ def test_matpes_relax_makers( assert isinstance(output, ForceFieldTaskDocument) ref = refs[ref_func] - assert output.output.energy_per_atom == approx(ref["energy_per_atom"]) + assert output.output.energy_per_atom == approx(ref["energy_per_atom"], rel=1e-3) assert output.structure.volume == approx(ref["volume"]) assert np.all( np.abs(np.array(output.output.ionic_steps[-1].forces) - np.array(ref["forces"])) diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 95885ba0c4..f5b1ddcc1d 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -2,17 +2,20 @@ import sys from contextlib import nullcontext +from itertools import product from pathlib import Path import numpy as np import pytest from ase import units -from ase.io import Trajectory +from ase.io import Trajectory as AseTrajectory from ase.md.verlet import VelocityVerlet +from emmet.core.trajectory import AtomTrajectory from jobflow import run_locally from monty.serialization import loadfn from pymatgen.analysis.structure_matcher import StructureMatcher from pymatgen.core import Structure +from pymatgen.core.trajectory import Trajectory as PmgTrajectory from atomate2.forcefields import MLFF from atomate2.forcefields.md import ForceFieldMDMaker @@ -37,9 +40,15 @@ def test_maker_initialization(): ) == ForceFieldMDMaker(force_field_name=mlff) -@pytest.mark.parametrize("ff_name", MLFF) +@pytest.mark.parametrize("ff_name, use_emmet_models", product(MLFF, [True, False])) def test_ml_ff_md_maker( - ff_name, si_structure, sr_ti_o3_structure, al2_au_structure, test_dir, clean_dir + ff_name, + use_emmet_models, + si_structure, + sr_ti_o3_structure, + al2_au_structure, + test_dir, + clean_dir, ): if ff_name in map(MLFF, ("Forcefield", "MACE")): return # nothing to test here, MLFF.Forcefield is just a generic placeholder @@ -62,8 +71,8 @@ def test_ml_ff_md_maker( MLFF.NEP: -3.966232215741286, MLFF.Nequip: -8.84670181274414, MLFF.SevenNet: -5.394115447998047, - MLFF.MATPES_PBE: -5.4188947677612305, - MLFF.MATPES_R2SCAN: -8.707625389099121, + MLFF.MATPES_PBE: -5.230762481689453, + MLFF.MATPES_R2SCAN: -8.561729431152344, } # ASE can slightly change tolerances on structure positions @@ -103,6 +112,7 @@ def test_ml_ff_md_maker( store_trajectory="partial", ionic_step_data=("energy", "forces", "stress", "mol_or_struct"), calculator_kwargs=calculator_kwargs, + use_emmet_models=use_emmet_models, ).make(structure) response = run_locally(job, ensure_success=True) task_doc = response[next(iter(response))][1].output @@ -128,10 +138,20 @@ def test_ml_ff_md_maker( assert task_doc.included_objects == ["trajectory"] assert len(task_doc.objects["trajectory"]) == n_steps + 1 assert task_doc.objects == task_doc.forcefield_objects # test legacy alias - assert all( - getattr(task_doc.objects["trajectory"], key, None) is not None - for key in ("energy", "forces", "stress", "velocities", "temperature") - ) + + if use_emmet_models: + assert all( + getattr(task_doc.objects["trajectory"], key, None) is not None + for key in ("energy", "forces", "stress", "velocities", "temperature") + ) + assert isinstance(task_doc.objects["trajectory"], AtomTrajectory) + else: + assert all( + frame.get(key) is not None + for key in ("energy", "forces", "stress", "velocities", "temperature") + for frame in task_doc.objects["trajectory"].frame_properties + ) + assert isinstance(task_doc.objects["trajectory"], PmgTrajectory) @pytest.mark.parametrize( @@ -148,7 +168,7 @@ def test_traj_file(traj_file, ff_name, si_structure, clean_dir): traj_file_loader = loadfn else: traj_file_fmt = "ase" - traj_file_loader = Trajectory + traj_file_loader = AseTrajectory structure = si_structure.to_conventional() * (2, 2, 2) job = ForceFieldMDMaker( @@ -156,6 +176,7 @@ def test_traj_file(traj_file, ff_name, si_structure, clean_dir): n_steps=n_steps, traj_file=traj_file, traj_file_fmt=traj_file_fmt, + use_emmet_models=True, ).make(structure) response = run_locally(job, ensure_success=True) task_doc = response[next(iter(response))][1].output @@ -264,6 +285,7 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): dynamics="nose-hoover", temperature=temp_schedule, ase_md_kwargs={"ttime": 50.0 * units.fs, "pfactor": None}, + use_emmet_models=True, ).make(structure) response = run_locally(job, ensure_success=True) task_doc = response[next(iter(response))][1].output diff --git a/tests/forcefields/test_neb.py b/tests/forcefields/test_neb.py index d71118b6df..b98de63686 100644 --- a/tests/forcefields/test_neb.py +++ b/tests/forcefields/test_neb.py @@ -51,11 +51,11 @@ def test_neb_from_images(test_dir, clean_dir): output.energies[i] == pytest.approx(energy) for i, energy in enumerate( [ - -339.3764953613281, - -339.2301025390625, - -338.865234375, - -339.23004150390625, - -339.37652587890625, + -328.3260803222656, + -328.3229064941406, + -327.90411376953125, + -328.3229064941406, + -328.3260803222656, ] ) ) diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_1/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_2/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_3/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_4/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_5/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_6/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/dft_phonon_static_1_1_eos_deformation_7/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d0/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d1/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d2/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d3/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d4/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_1_d5/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d0/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d1/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d2/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d3/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d4/outputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/inputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/inputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/inputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/inputs/POTCAR.spec.gz differ diff --git a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/outputs/POTCAR.spec.gz b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/outputs/POTCAR.spec.gz index 44393de246..d64f15ee28 100644 Binary files a/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/outputs/POTCAR.spec.gz and b/tests/test_data/vasp/Si_qha_2/tight_relax_2_d5/outputs/POTCAR.spec.gz differ diff --git a/tests/vasp/flows/test_core.py b/tests/vasp/flows/test_core.py index 97cf028b5f..9bdddb035b 100644 --- a/tests/vasp/flows/test_core.py +++ b/tests/vasp/flows/test_core.py @@ -1,16 +1,9 @@ import pytest from emmet.core.tasks import TaskDoc - -try: - from emmet.core.types.enums import VaspObject -except ImportError: - from emmet.core.vasp.calculation import VaspObject +from emmet.core.types.enums import VaspObject from jobflow import run_locally -from pymatgen.electronic_structure.bandstructure import ( - BandStructure, - BandStructureSymmLine, -) +from atomate2 import SETTINGS from atomate2.vasp.flows.core import ( BandStructureMaker, DoubleRelaxMaker, @@ -25,6 +18,20 @@ from atomate2.vasp.jobs.core import RelaxMaker from atomate2.vasp.sets.core import RelaxSetGenerator +if SETTINGS.VASP_USE_EMMET_MODELS: + from emmet.core.band_theory import ElectronicBS + + BS_UNIF_TYPE = ElectronicBS + BS_LINE_TYPE = ElectronicBS +else: + from pymatgen.electronic_structure.bandstructure import ( + BandStructure, + BandStructureSymmLine, + ) + + BS_UNIF_TYPE = BandStructure + BS_LINE_TYPE = BandStructureSymmLine + def test_double_relax(mock_vasp, clean_dir, si_structure): # mapping from job name to directory containing test files @@ -128,11 +135,9 @@ def test_band_structure(mock_vasp, clean_dir, si_structure): VaspObject.DOS, } + assert isinstance(line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_LINE_TYPE) assert isinstance( - line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine - ) - assert isinstance( - uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure + uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_UNIF_TYPE ) @@ -169,7 +174,7 @@ def test_uniform_band_structure(mock_vasp, clean_dir, si_structure): VaspObject.DOS, } assert isinstance( - uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure + uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_UNIF_TYPE ) @@ -201,9 +206,7 @@ def test_line_mode_band_structure(mock_vasp, clean_dir, si_structure): assert static_output.output.energy == pytest.approx(-10.85037078) assert static_output.included_objects is None - assert isinstance( - line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine - ) + assert isinstance(line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_LINE_TYPE) def test_hse_band_structure(mock_vasp, clean_dir, si_structure): @@ -244,11 +247,9 @@ def test_hse_band_structure(mock_vasp, clean_dir, si_structure): VaspObject.DOS, } + assert isinstance(line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_LINE_TYPE) assert isinstance( - line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine - ) - assert isinstance( - uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure + uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_UNIF_TYPE ) @@ -287,7 +288,7 @@ def test_hse_uniform_band_structure(mock_vasp, clean_dir, si_structure): VaspObject.DOS, } assert isinstance( - uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructure + uniform_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_UNIF_TYPE ) @@ -320,9 +321,7 @@ def test_hse_line_mode_band_structure(mock_vasp, clean_dir, si_structure): assert static_output.output.energy == pytest.approx(-12.52887403) assert static_output.included_objects is None - assert isinstance( - line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BandStructureSymmLine - ) + assert isinstance(line_output.vasp_objects[VaspObject.BANDSTRUCTURE], BS_LINE_TYPE) def test_optics(mock_vasp, clean_dir, si_structure): diff --git a/tests/vasp/flows/test_defect.py b/tests/vasp/flows/test_defect.py index 225d0d31e4..2d63c7d19a 100644 --- a/tests/vasp/flows/test_defect.py +++ b/tests/vasp/flows/test_defect.py @@ -174,7 +174,7 @@ def test_formation_energy_maker(mock_vasp, clean_dir, test_dir, monkeypatch): ) def _check_plnr_locpot(name): - job = SETTINGS.JOB_STORE.query_one({"output.task_label": name}) + job = SETTINGS.JOB_STORE.query_one({"output.task_label": name}, load=True) plnr_locpot = job["output"]["calcs_reversed"][0]["output"]["locpot"] assert set(plnr_locpot) == {"0", "1", "2"} diff --git a/tests/vasp/flows/test_ferroelectric.py b/tests/vasp/flows/test_ferroelectric.py index 9aadccac15..6b64e5298b 100644 --- a/tests/vasp/flows/test_ferroelectric.py +++ b/tests/vasp/flows/test_ferroelectric.py @@ -1,15 +1,14 @@ import pytest +from emmet.core.tasks import TaskDoc +from jobflow import run_locally +from pymatgen.core import Structure +from atomate2.vasp.flows.ferroelectric import FerroelectricMaker +from atomate2.vasp.powerups import update_user_incar_settings +from atomate2.vasp.schemas.ferroelectric import PolarizationDocument -def test_my_flow(mock_vasp, clean_dir, test_dir): - from emmet.core.tasks import TaskDoc - from jobflow import run_locally - from pymatgen.core import Structure - - from atomate2.vasp.flows.ferroelectric import FerroelectricMaker - from atomate2.vasp.powerups import update_user_incar_settings - from atomate2.vasp.schemas.ferroelectric import PolarizationDocument +def test_ferroelectric_flow(mock_vasp, clean_dir, test_dir): # mapping from job name to directory containing test files ref_paths = { "polarization interpolation_0": "KNbO3_ferroelectric/polarization_interpolation_0", # noqa: E501 diff --git a/tests/vasp/flows/test_md.py b/tests/vasp/flows/test_md.py index 0c0959ad8c..f7dde2d582 100644 --- a/tests/vasp/flows/test_md.py +++ b/tests/vasp/flows/test_md.py @@ -1,10 +1,20 @@ from emmet.core.types.enums import VaspObject from jobflow import Flow +from atomate2 import SETTINGS from atomate2.vasp.flows.md import MultiMDMaker from atomate2.vasp.powerups import update_user_kpoints_settings from atomate2.vasp.schemas.md import MultiMDOutput +if SETTINGS.VASP_USE_EMMET_MODELS: + from emmet.core.trajectory import RelaxTrajectory + + TRAJ_TYPE = RelaxTrajectory +else: + from pymatgen.core.trajectory import Trajectory + + TRAJ_TYPE = Trajectory + def test_multi_md_flow(mock_vasp, clean_dir, si_structure): from emmet.core.tasks import TaskDoc @@ -48,7 +58,15 @@ def test_multi_md_flow(mock_vasp, clean_dir, si_structure): # validate the outputs output_md_1 = responses[flow.jobs[0].jobs[0].uuid][1].output traj = output_md_1.vasp_objects[VaspObject.TRAJECTORY] - assert len(traj.frame_properties) == 3 + assert isinstance(traj, TRAJ_TYPE) + if SETTINGS.VASP_USE_EMMET_MODELS: + assert traj.num_ionic_steps == 3 + assert all( + len(getattr(traj, k)) == traj.num_ionic_steps + for k in ("energy", "forces", "lattice", "stress") + ) + else: + assert len(traj.frame_properties) == 3 assert isinstance(output_md_1, TaskDoc) output_recap_1 = responses[flow.jobs[0].jobs[2].uuid][1].output diff --git a/tests/vasp/flows/test_phonons.py b/tests/vasp/flows/test_phonons.py index 057a803988..887f6fcc72 100644 --- a/tests/vasp/flows/test_phonons.py +++ b/tests/vasp/flows/test_phonons.py @@ -358,7 +358,6 @@ def test_phonon_wf_vasp_only_displacements_add_inputs( responses = run_locally(job, create_folders=True, ensure_success=True) # validate the outputs - # print(type(responses)) assert isinstance(responses[job.jobs[-1].uuid][1].output, PhononBSDOSDoc) assert_allclose( diff --git a/tests/vasp/jobs/test_core.py b/tests/vasp/jobs/test_core.py index 55d4bf0a73..fda32acbb8 100644 --- a/tests/vasp/jobs/test_core.py +++ b/tests/vasp/jobs/test_core.py @@ -1,10 +1,11 @@ -import jobflow import numpy as np from emmet.core.tasks import TaskDoc -from jobflow import run_locally +from jobflow import JobStore, run_locally +from maggma.stores import MemoryStore from numpy.testing import assert_allclose from pytest import approx +from atomate2 import SETTINGS from atomate2.vasp.jobs.core import ( DielectricMaker, HSERelaxMaker, @@ -14,10 +15,12 @@ TransmuterMaker, ) +if SETTINGS.VASP_USE_EMMET_MODELS: + from emmet.core.vasp.models import ChgcarLike -def test_static_maker(mock_vasp, clean_dir, si_structure): - job_store = jobflow.SETTINGS.JOB_STORE +def test_static_maker(mock_vasp, clean_dir, si_structure): + job_store = JobStore(MemoryStore(), additional_stores={"data": MemoryStore()}) # mapping from job name to directory containing test files ref_paths = {"static": "Si_band_structure/static"} @@ -35,7 +38,9 @@ def test_static_maker(mock_vasp, clean_dir, si_structure): ) # run the flow or job and ensure that it finished running successfully - responses = run_locally(job, create_folders=True, ensure_success=True) + responses = run_locally( + job, create_folders=True, ensure_success=True, store=job_store + ) # validate job outputs output1 = responses[job.uuid][1].output @@ -44,7 +49,11 @@ def test_static_maker(mock_vasp, clean_dir, si_structure): with job_store.additional_stores["data"] as store: doc = store.query_one({"job_uuid": job.uuid}) - assert doc["data"]["@class"] == "Chgcar" + + if SETTINGS.VASP_USE_EMMET_MODELS: + assert all(k in doc["data"] for k in ChgcarLike.model_fields) + else: + assert doc["data"]["@class"] == "Chgcar" def test_relax_maker(mock_vasp, clean_dir, si_structure): diff --git a/tests/vasp/jobs/test_md.py b/tests/vasp/jobs/test_md.py index 05a8fab83c..2205c8b794 100644 --- a/tests/vasp/jobs/test_md.py +++ b/tests/vasp/jobs/test_md.py @@ -3,7 +3,9 @@ from emmet.core.types.enums import VaspObject from emmet.core.vasp.calculation import IonicStep from jobflow import run_locally +from pymatgen.core import Element, Structure +from atomate2 import SETTINGS from atomate2.vasp.jobs.md import MDMaker @@ -46,9 +48,30 @@ def test_molecular_dynamics(mock_vasp, clean_dir, si_structure): # check ionic steps stored as pymatgen Trajectory assert output1.calcs_reversed[0].output.ionic_steps is None traj = output1.vasp_objects[VaspObject.TRAJECTORY] - assert len(traj.frame_properties) == nsw + if SETTINGS.VASP_USE_EMMET_MODELS: + assert all( + len(getattr(traj, k)) == nsw + for k in ("energy", "forces", "lattice", "stress") + ) + else: + assert len(traj.frame_properties) == nsw + # check that a frame property can be converted to an IonicStep energies = [-11.47041923, -11.46905352, -11.46520398] - for idx, frame in enumerate(traj.frame_properties): - ionic_step = IonicStep(**frame) - assert ionic_step.e_wo_entrp == pytest.approx(energies[idx]) + for idx, ref_energy in enumerate(energies): + if SETTINGS.VASP_USE_EMMET_MODELS: + ionic_step = IonicStep( + **{ + k: getattr(traj, k)[idx] + for k in ("energy", "forces", "stress", "e_wo_entrp") + }, + structure=Structure( + traj.lattice[idx], + [Element.from_Z(z) for z in traj.elements], + traj.cart_coords[idx], + coords_are_cartesian=True, + ), + ) + else: + ionic_step = IonicStep(**traj.frame_properties[idx]) + assert ionic_step.e_wo_entrp == pytest.approx(ref_energy)