From 7637392deffbd996f8e5c86246bc71b08a3fc5ee Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 15 Feb 2024 15:01:50 -0800 Subject: [PATCH 01/48] FF MD fully implemented --- src/atomate2/forcefields/md.py | 232 ++++++++++++++++++++++++++++ src/atomate2/forcefields/schemas.py | 4 +- 2 files changed, 235 insertions(+), 1 deletion(-) create mode 100644 src/atomate2/forcefields/md.py diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py new file mode 100644 index 0000000000..6b91115223 --- /dev/null +++ b/src/atomate2/forcefields/md.py @@ -0,0 +1,232 @@ +"""Makers to perform MD with forcefields.""" +from __future__ import annotations + +import contextlib +import io +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from ase.md.andersen import Andersen +from ase.md.langevin import Langevin +from ase.md.npt import NPT +from ase.md.nptberendsen import NPTBerendsen +from ase.md.nvtberendsen import NVTBerendsen +from ase.md.velocitydistribution import ( + MaxwellBoltzmannDistribution, + Stationary, + ZeroRotation, +) +from ase.md.verlet import VelocityVerlet +from jobflow import Maker, job +from pymatgen.io.ase import AseAtomsAdaptor + +from atomate2.forcefields.schemas import ForceFieldTaskDocument +from atomate2.forcefields.utils import TrajectoryObserver + +if TYPE_CHECKING: + from pathlib import Path + from typing import Literal + + from ase.calculators.calculator import Calculator + from pymatgen.core.structure import Structure + + +_valid_thermostats: dict[str, tuple[str, ...]] = { + "nve": ("velocityverlet",), + "nvt": ("nose-hoover", "langevin", "andersen", "berendsen"), + "npt": ("nose-hoover", "berendsen"), +} + +_thermostats: dict = { + "nve_velocityverlet": VelocityVerlet, + "nvt_andersen": Andersen, + "nvt_berendsen": NVTBerendsen, + "nvt_langevin": Langevin, + "nvt_nose-hoover": NPT, + "npt_berendsen": NPTBerendsen, + "npt_nose-hoover": NPT, +} + + +@dataclass +class ForceFieldMDMaker(Maker): + """ + Perform MD with a forcefield. + + Parameters + ---------- + name : str + The name of the MD Maker + force_field_name : str + The name of the forcefield (for provenance) + timestep : float | None = 0.2 + The timestep of the MD run in ase time units + md_steps : int = 500 + The number of MD steps to run + ensemble : str = "nve" + The ensemble to use. Valid ensembles are nve, nvt, or npt + temperature : float | None = 300. + The temperature in Kelvin + thermostat : str = "velocityverlet" + The thermostat to use. See _valid_thermostats for a list of options + ase_md_kwargs : dict | None = None + Options to pass to the ASE MD function + calculator_args : Sequence | None = None + args to pass to the ASE calculator class + calculator_kwargs : dict | None = None + kwargs to pass to the ASE calculator class + traj_file : str | Path | None = None + If a str or Path, the name of the file to save the MD trajectory to. + If None, the trajectory is not written to disk + traj_interval : int + The step interval for saving the trajectories. + zero_linear_momentum : bool = False + Whether to initialize the atomic velocities with zero linear momentum + zero_angular_momentum : bool = False + Whether to initialize the atomic velocities with zero angular momentum + task_document_kwargs: dict + Options to pass to the TaskDoc. Default choice + {"ionic_step_data": ("energy", "forces", "magmoms", "stress",)} + is consistent with atomate2.vasp.md.MDMaker + """ + + name: str = "Forcefield MD" + force_field_name: str = "Forcefield" + timestep: float | None = 0.2 # approx 2 fs + md_steps: int = 500 + ensemble: Literal["nve", "nvt", "npt"] = "nve" + temperature: float | None = 300.0 + thermostat: str = "velocityverlet" + ase_md_kwargs: dict | None = None + calculator_args: list | tuple | None = None + calculator_kwargs: dict | None = None + traj_file: str | Path | None = None + traj_interval: int = 1 + zero_linear_momentum: bool = False + zero_angular_momentum: bool = False + task_document_kwargs: dict = field( + default_factory=lambda: { + "ionic_step_data": ("energy", "forces", "magmoms", "stress") + } + ) + + @job(output_schema=ForceFieldTaskDocument) + def make( + self, + structure: Structure, + prev_dir: str | Path | None = None, + ) -> ForceFieldTaskDocument: + """ + Perform MD on a structure using a force field. + + Parameters + ---------- + structure: .Structure + pymatgen structure. + prev_dir : str or Path or None + A previous calculation directory to copy output files from. Unused, just + added to match the method signature of other makers. + """ + initial_velocities = structure.site_properties.get("velocities") + + if self.thermostat.lower() not in _valid_thermostats[self.ensemble]: + raise ValueError( + f"{self.thermostat} thermostat not available for {self.ensemble}." + f"Available {self.ensemble} thermostats are:" + " ".join(_valid_thermostats[self.ensemble]) + ) + + if self.ensemble == "nve" and self.thermostat is None: + self.thermostat = "velocityverlet" + md_func = _thermostats[f"{self.ensemble}_{self.thermostat.lower()}"] + + atoms = AseAtomsAdaptor.get_atoms(structure) + if initial_velocities: + atoms.set_velocities(initial_velocities) + elif self.temperature: + MaxwellBoltzmannDistribution(atoms=atoms, temperature_K=self.temperature) + + if self.zero_linear_momentum: + Stationary(atoms) + + if self.zero_angular_momentum: + ZeroRotation(atoms) + + else: + atoms.set_velocities( + [[0.0 for _ in range(3)] for _ in range(len(structure))] + ) + + self.calculator_args = self.calculator_args or [] + self.calculator_kwargs = self.calculator_kwargs or {} + atoms.set_calculator(self._calculator()) + + with contextlib.redirect_stdout(io.StringIO()): + md_observer = TrajectoryObserver(atoms) + self.ase_md_kwargs = self.ase_md_kwargs or {} + md_runner = md_func( + atoms=atoms, timestep=self.timestep, **self.ase_md_kwargs + ) + md_runner.attach(md_observer, interval=self.traj_interval) + + md_runner.run(steps=self.md_steps) + md_observer() + + if self.traj_file is not None: + md_observer.save(self.traj_file) + + structure = AseAtomsAdaptor.get_structure(atoms) + + self.task_document_kwargs = self.task_document_kwargs or {} + return ForceFieldTaskDocument.from_ase_compatible_result( + self.force_field_name, + {"final_structure": structure, "trajectory": md_observer}, + relax_cell=(self.ensemble == "npt"), + steps=self.md_steps, + relax_kwargs=None, + optimizer_kwargs=None, + **self.task_document_kwargs, + ) + + def _calculator(self) -> Calculator: + """To be implemented by the user.""" + return NotImplementedError + + +class MACEMDMaker(ForceFieldMDMaker): + """Perform an MD run with MACE.""" + + name: str = "MACE MD" + force_field_name: str = "MACE" + calculator_kwargs: dict = field(default_factory=lambda: {"dtype": "float32"}) + + def _calculator(self) -> Calculator: + from mace.calculators import mace_mp + + return mace_mp(*self.calculator_args, **self.calculator_kwargs) + + +class M3GNetMDMaker(ForceFieldMDMaker): + """Perform an MD run with M3GNet.""" + + name: str = "M3GNet MD" + force_field_name: str = "M3GNet" + + def _calculator(self) -> Calculator: + import matgl + from matgl.ext.ase import PESCalculator + + potential = matgl.load_model("M3GNet-MP-2021.2.8-PES") + return PESCalculator(potential, **self.calculator_kwargs) + + +class CHGNetMDMaker(ForceFieldMDMaker): + """Perform an MD run with CHGNet.""" + + name: str = "CHGNet MD" + force_field_name: str = "CHGNet" + + def _calculator(self) -> Calculator: + from chgnet.model.dynamics import CHGNetCalculator + + return CHGNetCalculator(*self.calculator_args, **self.calculator_kwargs) diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 28d0e013db..d5f4f0455e 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -19,7 +19,9 @@ class IonicStep(BaseModel, extra="allow"): # type: ignore[call-arg] None, description="The forces on each atom." ) stress: Optional[Matrix3D] = Field(None, description="The stress on the lattice.") - structure: Structure = Field(None, description="The structure at this step.") + structure: Optional[Structure] = Field( + None, description="The structure at this step." + ) class InputDoc(BaseModel): From 659ad4a63e9fd0e0b7942e6665c26c4bc149c42a Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 15 Feb 2024 15:23:43 -0800 Subject: [PATCH 02/48] change MD defaults to be NVT + Langevin dynamics --- src/atomate2/forcefields/md.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 6b91115223..528df7d5d4 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -63,11 +63,11 @@ class ForceFieldMDMaker(Maker): The timestep of the MD run in ase time units md_steps : int = 500 The number of MD steps to run - ensemble : str = "nve" + ensemble : str = "nvt" The ensemble to use. Valid ensembles are nve, nvt, or npt temperature : float | None = 300. The temperature in Kelvin - thermostat : str = "velocityverlet" + thermostat : str = "langevin" The thermostat to use. See _valid_thermostats for a list of options ase_md_kwargs : dict | None = None Options to pass to the ASE MD function @@ -94,9 +94,9 @@ class ForceFieldMDMaker(Maker): force_field_name: str = "Forcefield" timestep: float | None = 0.2 # approx 2 fs md_steps: int = 500 - ensemble: Literal["nve", "nvt", "npt"] = "nve" + ensemble: Literal["nve", "nvt", "npt"] = "nvt" temperature: float | None = 300.0 - thermostat: str = "velocityverlet" + thermostat: str = "langevin" ase_md_kwargs: dict | None = None calculator_args: list | tuple | None = None calculator_kwargs: dict | None = None From ade233e3cb3a6fa70725353470e95f51f7cd9aca Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 15 Feb 2024 16:33:18 -0800 Subject: [PATCH 03/48] add md defaults that are as consistent with vasp implementation as possible --- src/atomate2/forcefields/md.py | 76 +++++++++++++++++++++++++++------- 1 file changed, 62 insertions(+), 14 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 528df7d5d4..75535f6cb4 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -17,6 +17,7 @@ ZeroRotation, ) from ase.md.verlet import VelocityVerlet +from ase.units import bar, fs from jobflow import Maker, job from pymatgen.io.ase import AseAtomsAdaptor @@ -53,20 +54,31 @@ class ForceFieldMDMaker(Maker): """ Perform MD with a forcefield. + Note that units are consistent with the VASP MD implementation. + - Temperature is in Kelvin (TEBEG and TEEND) + - Time steps are in femtoseconds (POTIM) + - Langevin NVT friction coefficients are in picoseconds^-1 (LANGEVIN_GAMMA) + - Pressure in kB (PSTRESS) + Parameters ---------- name : str The name of the MD Maker force_field_name : str The name of the forcefield (for provenance) - timestep : float | None = 0.2 - The timestep of the MD run in ase time units - md_steps : int = 500 + timestep : float | None = 2. + The timestep of the MD run in femtoseconds + md_steps : int = 1000 The number of MD steps to run ensemble : str = "nvt" The ensemble to use. Valid ensembles are nve, nvt, or npt - temperature : float | None = 300. - The temperature in Kelvin + start_temperature : float | None = 300. + The temperature to initialize the system, in Kelvin. + end_temperature : float | None = 300. + The temperature to equilibrate towards, in Kelvin. + If start_temperature is a float and end_temperature is None, + the system will be initialized at and equilibrated towards + the start_temperature. thermostat : str = "langevin" The thermostat to use. See _valid_thermostats for a list of options ase_md_kwargs : dict | None = None @@ -92,10 +104,12 @@ class ForceFieldMDMaker(Maker): name: str = "Forcefield MD" force_field_name: str = "Forcefield" - timestep: float | None = 0.2 # approx 2 fs - md_steps: int = 500 + timestep: float | None = 2.0 + md_steps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" - temperature: float | None = 300.0 + start_temperature: float | None = 300.0 + end_temperature: float | None = 300.0 + pressure: float | None = None thermostat: str = "langevin" ase_md_kwargs: dict | None = None calculator_args: list | tuple | None = None @@ -110,6 +124,34 @@ class ForceFieldMDMaker(Maker): } ) + def _get_ensemble_defaults(self, structure: Structure) -> None: + """Update ASE MD kwargs with defaults consistent with VASP MD.""" + self.ase_md_kwargs = self.ase_md_kwargs or {} + if self.ensemble in ("nvt", "npt") and all( + self.ase_md_kwargs.get(key) is None + for key in ("temperature_K", "temperature") + ): + self.ase_md_kwargs["temperature_K"] = ( + self.end_temperature if self.end_temperature else self.start_temperature + ) + + if self.ensemble == "npt": + # convert from kilobar to eV/Ang**3 + self.ase_md_kwargs["pressure_au"] = ( + self.pressure * 1.0e-3 / bar if self.pressure else 0.0 + ) + + if self.thermostat == "langevin": + # Same default as in VASP + self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get("friction", 10.0) + # friction coefficient(s) specified in ev^-1, convert from picoseconds + if isinstance(self.ase_md_kwargs["friction"], list): + self.ase_md_kwargs["friction"] = [ + coeff * 1.0e-3 / fs for coeff in self.ase_md_kwargs["friction"] + ] + else: + self.ase_md_kwargs["friction"] *= 1.0e-3 / fs + @job(output_schema=ForceFieldTaskDocument) def make( self, @@ -127,9 +169,12 @@ def make( A previous calculation directory to copy output files from. Unused, just added to match the method signature of other makers. """ + self._get_ensemble_defaults(structure=structure) + initial_velocities = structure.site_properties.get("velocities") - if self.thermostat.lower() not in _valid_thermostats[self.ensemble]: + self.thermostat = self.thermostat.lower() + if self.thermostat not in _valid_thermostats[self.ensemble]: raise ValueError( f"{self.thermostat} thermostat not available for {self.ensemble}." f"Available {self.ensemble} thermostats are:" @@ -138,13 +183,16 @@ def make( if self.ensemble == "nve" and self.thermostat is None: self.thermostat = "velocityverlet" - md_func = _thermostats[f"{self.ensemble}_{self.thermostat.lower()}"] + md_func = _thermostats[f"{self.ensemble}_{self.thermostat}"] atoms = AseAtomsAdaptor.get_atoms(structure) if initial_velocities: atoms.set_velocities(initial_velocities) - elif self.temperature: - MaxwellBoltzmannDistribution(atoms=atoms, temperature_K=self.temperature) + + elif self.start_temperature: + MaxwellBoltzmannDistribution( + atoms=atoms, temperature_K=self.start_temperature + ) if self.zero_linear_momentum: Stationary(atoms) @@ -163,9 +211,9 @@ def make( with contextlib.redirect_stdout(io.StringIO()): md_observer = TrajectoryObserver(atoms) - self.ase_md_kwargs = self.ase_md_kwargs or {} + md_runner = md_func( - atoms=atoms, timestep=self.timestep, **self.ase_md_kwargs + atoms=atoms, timestep=self.timestep * fs, **self.ase_md_kwargs ) md_runner.attach(md_observer, interval=self.traj_interval) From db51e233b1743fff5ce6684bb4fb8a001bf2e26e Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 15 Feb 2024 16:47:22 -0800 Subject: [PATCH 04/48] correct typing of langevin friction coeff --- src/atomate2/forcefields/md.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 75535f6cb4..f461b7871e 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -145,7 +145,7 @@ def _get_ensemble_defaults(self, structure: Structure) -> None: # Same default as in VASP self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get("friction", 10.0) # friction coefficient(s) specified in ev^-1, convert from picoseconds - if isinstance(self.ase_md_kwargs["friction"], list): + if isinstance(self.ase_md_kwargs["friction"], (list, tuple)): self.ase_md_kwargs["friction"] = [ coeff * 1.0e-3 / fs for coeff in self.ase_md_kwargs["friction"] ] From 1dcad1d4306f9d2ef7df24055ad578b5d96c7fcd Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 16 Feb 2024 12:35:51 -0800 Subject: [PATCH 05/48] Add option to ForceFieldTaskDocument to store trajectory information as pmg object --- src/atomate2/forcefields/md.py | 8 ++- src/atomate2/forcefields/schemas.py | 77 ++++++++++++++++++++--------- src/atomate2/forcefields/utils.py | 15 +++++- 3 files changed, 72 insertions(+), 28 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index f461b7871e..4f9bf00f38 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -98,7 +98,7 @@ class ForceFieldMDMaker(Maker): Whether to initialize the atomic velocities with zero angular momentum task_document_kwargs: dict Options to pass to the TaskDoc. Default choice - {"ionic_step_data": ("energy", "forces", "magmoms", "stress",)} + {"store_trajectory": "partial"} is consistent with atomate2.vasp.md.MDMaker """ @@ -119,9 +119,7 @@ class ForceFieldMDMaker(Maker): zero_linear_momentum: bool = False zero_angular_momentum: bool = False task_document_kwargs: dict = field( - default_factory=lambda: { - "ionic_step_data": ("energy", "forces", "magmoms", "stress") - } + default_factory=lambda: {"store_trajectory": "partial"} ) def _get_ensemble_defaults(self, structure: Structure) -> None: @@ -210,7 +208,7 @@ def make( atoms.set_calculator(self._calculator()) with contextlib.redirect_stdout(io.StringIO()): - md_observer = TrajectoryObserver(atoms) + md_observer = TrajectoryObserver(atoms, store_md_outputs=True) md_runner = md_func( atoms=atoms, timestep=self.timestep * fs, **self.ase_md_kwargs diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index d5f4f0455e..8bbe3da7cf 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -1,16 +1,25 @@ """Schema definitions for force field tasks.""" -from typing import Optional +from typing import Any, Optional 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 StructureMetadata +from emmet.core.utils import ValueEnum +from emmet.core.vasp.calculation import StoreTrajectoryOption from pydantic import BaseModel, Field from pymatgen.core.structure import Structure +from pymatgen.core.trajectory import Trajectory from pymatgen.io.ase import AseAtomsAdaptor +class ForcefieldObject(ValueEnum): + """Types of forcefield data objects.""" + + TRAJECTORY = "trajectory" + + class IonicStep(BaseModel, extra="allow"): # type: ignore[call-arg] """Document defining the information at each ionic step.""" @@ -104,6 +113,13 @@ class ForceFieldTaskDocument(StructureMetadata): None, description="Directory where the force field calculations are performed." ) + included_objects: Optional[list[ForcefieldObject]] = Field( + None, description="list of forcefield objects included with this task document" + ) + forcefield_objects: Optional[dict[ForcefieldObject, Any]] = Field( + None, description="Forcefield objects associated with this task" + ) + @classmethod def from_ase_compatible_result( cls, @@ -114,6 +130,7 @@ def from_ase_compatible_result( relax_kwargs: dict = None, optimizer_kwargs: dict = None, ionic_step_data: tuple = ("energy", "forces", "magmoms", "stress", "structure"), + store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO, ) -> "ForceFieldTaskDocument": """ Create a ForceFieldTaskDocument for a Task that has ASE-compatible outputs. @@ -206,30 +223,44 @@ def from_ase_compatible_result( cur_structure = None # include "magmoms" in :obj:`ionic_step` if the trajectory has "magmoms" - if "magmoms" in trajectory: - ionic_step = IonicStep( - energy=energy, - forces=forces, - magmoms=( - trajectory["magmoms"][idx].tolist() - if "magmoms" in ionic_step_data - else None - ), - stress=stress, - structure=cur_structure, - ) - - # otherwise do not include "magmoms" in :obj:`ionic_step` - elif "magmoms" not in trajectory: - ionic_step = IonicStep( - energy=energy, - forces=forces, - stress=stress, - structure=cur_structure, - ) + magmom_data = {} + if all("magmoms" in obj for obj in (trajectory, ionic_step_data)): + magmom_data = {"magmoms": trajectory["magmoms"][idx].tolist()} + + ionic_step = IonicStep( + energy=energy, + forces=forces, + stress=stress, + structure=cur_structure, + **magmom_data, + ) ionic_steps.append(ionic_step) + forcefield_objects: dict[ForcefieldObject, Any] = {} + if store_trajectory != StoreTrajectoryOption.NO: + # For VASP calculations, the PARTIAL trajectory option removes + # electronic step info. There is no equivalent for forcefields, + # so we just save the same info for FULL and PARTIAL options. + + traj_steps = [ + step.model_dump(exclude=("structure",)) for step in ionic_steps + ] + for idx in range(n_steps): + if trajectory.get("temperatures"): + traj_steps[idx]["temperature"] = trajectory.get("temperatures")[idx] + if trajectory.get("velocities"): + traj_steps[idx]["velocities"] = trajectory["velocities"][ + idx + ].tolist() + + traj = Trajectory.from_structures( + [step.structure for step in ionic_steps], + frame_properties=traj_steps, + constant_lattice=False, + ) + forcefield_objects[ForcefieldObject.TRAJECTORY] = traj # type: ignore[index] + output_doc = OutputDoc( structure=output_structure, energy=final_energy, @@ -257,4 +288,6 @@ def from_ase_compatible_result( output=output_doc, forcefield_name=forcefield_name, forcefield_version=version, + included_objects=list(forcefield_objects.keys()), + forcefield_objects=forcefield_objects, ) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 95b24fe75f..ba5f8178ab 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -64,7 +64,7 @@ class TrajectoryObserver: This is a hook in the relaxation process that saves the intermediate structures. """ - def __init__(self, atoms: Atoms) -> None: + def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: """ Initialize the Observer. @@ -83,6 +83,11 @@ def __init__(self, atoms: Atoms) -> None: self.atom_positions: list[np.ndarray] = [] self.cells: list[np.ndarray] = [] + self._store_md_outputs = store_md_outputs + if self._store_md_outputs: + self.velocities: list[np.ndarray] = [] + self.temperatures: list[float] = [] + def __call__(self) -> None: """Save the properties of an Atoms during the relaxation.""" # TODO: maybe include magnetic moments @@ -92,6 +97,10 @@ def __call__(self) -> None: self.atom_positions.append(self.atoms.get_positions()) self.cells.append(self.atoms.get_cell()[:]) + if self._store_md_outputs: + self.velocities.append(self.atoms.get_velocities()) + self.temperatures.append(self.atoms.get_temperature()) + def compute_energy(self) -> float: """ Calculate the energy, here we just use the potential energy. @@ -122,6 +131,10 @@ def save(self, filename: str | PathLike) -> None: "cell": self.cells, "atomic_number": self.atoms.get_atomic_numbers(), } + if self._store_md_outputs: + traj_dict.update( + {"velocities": self.velocities, "temperature": self.temperatures} + ) with open(filename, "wb") as file: pickle.dump(traj_dict, file) From 35fa20d2fe2ac0389d3e311310d8e42c68835b7c Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 16 Feb 2024 15:28:26 -0800 Subject: [PATCH 06/48] Add tests, remove dependence on pickle for trajectory i/o --- src/atomate2/forcefields/md.py | 1 - src/atomate2/forcefields/utils.py | 19 ++++++-- tests/forcefields/test_md.py | 78 +++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 tests/forcefields/test_md.py diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 4f9bf00f38..eeb8901e9f 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -216,7 +216,6 @@ def make( md_runner.attach(md_observer, interval=self.traj_interval) md_runner.run(steps=self.md_steps) - md_observer() if self.traj_file is not None: md_observer.save(self.traj_file) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index ba5f8178ab..d3735d6528 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -11,13 +11,14 @@ import contextlib import io -import pickle import sys import warnings from typing import TYPE_CHECKING +import numpy as np from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG +from monty.serialization import dumpfn from pymatgen.core.structure import Molecule, Structure from pymatgen.io.ase import AseAtomsAdaptor @@ -39,7 +40,6 @@ from os import PathLike from typing import Any - import numpy as np from ase import Atoms from ase.calculators.calculator import Calculator from ase.filters import Filter @@ -113,7 +113,7 @@ def compute_energy(self) -> float: def save(self, filename: str | PathLike) -> None: """ - Save the trajectory file. + Save the trajectory file using monty.serialization. Parameters ---------- @@ -123,6 +123,10 @@ def save(self, filename: str | PathLike) -> None: ------- None """ + dumpfn(self.as_dict(), filename) + + def as_dict(self) -> dict: + """Make JSONable dict representation of the Trajectory.""" traj_dict = { "energy": self.energies, "forces": self.forces, @@ -135,8 +139,13 @@ def save(self, filename: str | PathLike) -> None: traj_dict.update( {"velocities": self.velocities, "temperature": self.temperatures} ) - with open(filename, "wb") as file: - pickle.dump(traj_dict, file) + # sanitize dict + for key in traj_dict: + if all(isinstance(val, np.ndarray) for val in traj_dict[key]): + traj_dict[key] = [val.tolist() for val in traj_dict[key]] + elif isinstance(traj_dict[key], np.ndarray): + traj_dict[key] = traj_dict[key].tolist() + return traj_dict class Relaxer: diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py new file mode 100644 index 0000000000..dbd5244d53 --- /dev/null +++ b/tests/forcefields/test_md.py @@ -0,0 +1,78 @@ +""" Tests for forcefield MD flows. """ + +import numpy as np +import pytest +import torch +from jobflow import run_locally +from monty.serialization import loadfn + +from atomate2.forcefields.md import CHGNetMDMaker, M3GNetMDMaker, MACEMDMaker + +_to_maker = {"CHGNet": CHGNetMDMaker, "M3GNet": M3GNetMDMaker, "MACE": MACEMDMaker} + + +@pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE"]) +def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): + md_steps = 5 + + ref_energies_per_atom = { + "CHGNet": -5.30686092376709, + "M3GNet": -5.417105674743652, + "MACE": -5.33246374130249, + } + + structure = si_structure.to_conventional() * (2, 2, 2) + + # MACE changes the default dtype, ensure consistent dtype here + torch.set_default_dtype(torch.float32) + + job = _to_maker[ff_name](md_steps=md_steps, traj_file="md_traj.json.gz").make( + structure + ) + response = run_locally(job, ensure_success=True) + taskdoc = response[next(iter(response))][1].output + + # Check that energies are reasonably close to reference values + assert taskdoc.output.energy / len(structure) == pytest.approx( + ref_energies_per_atom[ff_name], abs=0.1 + ) + + # Check that we have the right number of MD steps + # ASE logs the initial structure energy, and doesn"t count this as an MD step + assert taskdoc.output.ionic_steps[0].structure == structure + assert len(taskdoc.output.ionic_steps) == md_steps + 1 + + # Check that the ionic steps have the expected physical properties + assert all( + key in step.model_dump() + for key in ("energy", "forces", "stress", "structure") + for step in taskdoc.output.ionic_steps + ) + + # Check that the trajectory has expected physical properties + assert taskdoc.included_objects == ["trajectory"] + assert len(taskdoc.forcefield_objects["trajectory"]) == md_steps + 1 + assert all( + key in step + for key in ("energy", "forces", "stress", "velocities", "temperature") + for step in taskdoc.forcefield_objects["trajectory"].frame_properties + ) + + # Check that traj file written to disk is consistent with trajectory + # stored to the task document + traj_from_file = loadfn("md_traj.json.gz") + + assert len(traj_from_file["energy"]) == md_steps + 1 + _traj_key_to_object_key = { + "stresses": "stress", + } + assert all( + np.all( + traj_from_file[key][idx] + == taskdoc.forcefield_objects["trajectory"].frame_properties[idx][ + _traj_key_to_object_key.get(key, key) + ] + ) + for key in ("energy", "temperature", "forces", "velocities") + for idx in range(md_steps + 1) + ) From 0d8fdef7d496691eab51055aef62e04f89263688 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 16 Feb 2024 16:26:14 -0800 Subject: [PATCH 07/48] fix tests, update matgl dependence --- pyproject.toml | 4 ++-- src/atomate2/common/jobs/utils.py | 2 +- src/atomate2/forcefields/jobs.py | 2 +- src/atomate2/forcefields/schemas.py | 14 +++++++++++--- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9c92c67946..37e486316e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,7 @@ forcefields = [ "ase>=3.22.1", "chgnet>=0.2.2", "mace-torch>=0.3.3", - "matgl>=0.9.0", + "matgl>=1.0.0", "quippy-ase>=0.9.14", ] docs = [ @@ -80,7 +80,7 @@ strict = [ "jobflow==0.1.17", "lobsterpy==0.3.6", "mace-torch>=0.3.3", - "matgl==0.9.2", + "matgl==1.0.0", "monty==2023.9.25", "mp-api==0.39.5", "numpy", diff --git a/src/atomate2/common/jobs/utils.py b/src/atomate2/common/jobs/utils.py index dffcec0d4e..f0e99d2fc1 100644 --- a/src/atomate2/common/jobs/utils.py +++ b/src/atomate2/common/jobs/utils.py @@ -40,7 +40,7 @@ def structure_to_conventional( structure: Structure, symprec: float = SETTINGS.SYMPREC ) -> Structure: """ - Job hat creates a standard conventional structure. + Job that creates a standard conventional structure. Parameters ---------- diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 48a741f8a5..1c9b5df5c7 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -91,7 +91,7 @@ def _relax(self, structure: Structure) -> dict: @dataclass class ForceFieldStaticMaker(ForceFieldRelaxMaker): """ - Maker to calculate forces and stresses using the CHGNet force field. + Maker to calculate forces and stresses using any force field. Parameters ---------- diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 8bbe3da7cf..963a382f7f 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -152,7 +152,11 @@ def from_ase_compatible_result( ionic_step_data : tuple Which data to save from each ionic step. """ - trajectory = result["trajectory"].__dict__ + trajectory = { + key: value + for key, value in result["trajectory"].__dict__.items() + if not key.startswith("_") + } # NOTE: units for stresses were converted from eV/Angstrom³ to kBar # (* -1 from standard output) @@ -224,8 +228,12 @@ def from_ase_compatible_result( # include "magmoms" in :obj:`ionic_step` if the trajectory has "magmoms" magmom_data = {} - if all("magmoms" in obj for obj in (trajectory, ionic_step_data)): - magmom_data = {"magmoms": trajectory["magmoms"][idx].tolist()} + if trajectory.get("magmoms"): + magmom_data = { + "magmoms": trajectory["magmoms"][idx].tolist() + if "magmoms" in ionic_step_data + else None + } ionic_step = IonicStep( energy=energy, From 013d11366c18c9b95a600fe515b9be501c893467 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Tue, 20 Feb 2024 15:00:45 -0800 Subject: [PATCH 08/48] rename MD forcefield class vars to be consistent with VASP MDSetGenerator --- src/atomate2/forcefields/md.py | 36 ++++++++++++++++------------------ tests/forcefields/test_md.py | 14 ++++++------- 2 files changed, 23 insertions(+), 27 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index eeb8901e9f..a1f4389787 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -66,19 +66,19 @@ class ForceFieldMDMaker(Maker): The name of the MD Maker force_field_name : str The name of the forcefield (for provenance) - timestep : float | None = 2. - The timestep of the MD run in femtoseconds - md_steps : int = 1000 + time_step : float | None = 2. + The time_step of the MD run in femtoseconds + nsteps : int = 1000 The number of MD steps to run ensemble : str = "nvt" The ensemble to use. Valid ensembles are nve, nvt, or npt - start_temperature : float | None = 300. + start_temp : float | None = 300. The temperature to initialize the system, in Kelvin. - end_temperature : float | None = 300. + end_temp : float | None = 300. The temperature to equilibrate towards, in Kelvin. - If start_temperature is a float and end_temperature is None, + If start_temp is a float and end_temp is None, the system will be initialized at and equilibrated towards - the start_temperature. + the start_temp. thermostat : str = "langevin" The thermostat to use. See _valid_thermostats for a list of options ase_md_kwargs : dict | None = None @@ -104,11 +104,11 @@ class ForceFieldMDMaker(Maker): name: str = "Forcefield MD" force_field_name: str = "Forcefield" - timestep: float | None = 2.0 - md_steps: int = 1000 + time_step: float | None = 2.0 + nsteps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" - start_temperature: float | None = 300.0 - end_temperature: float | None = 300.0 + start_temp: float | None = 300.0 + end_temp: float | None = 300.0 pressure: float | None = None thermostat: str = "langevin" ase_md_kwargs: dict | None = None @@ -130,7 +130,7 @@ def _get_ensemble_defaults(self, structure: Structure) -> None: for key in ("temperature_K", "temperature") ): self.ase_md_kwargs["temperature_K"] = ( - self.end_temperature if self.end_temperature else self.start_temperature + self.end_temp if self.end_temp else self.start_temp ) if self.ensemble == "npt": @@ -187,10 +187,8 @@ def make( if initial_velocities: atoms.set_velocities(initial_velocities) - elif self.start_temperature: - MaxwellBoltzmannDistribution( - atoms=atoms, temperature_K=self.start_temperature - ) + elif self.start_temp: + MaxwellBoltzmannDistribution(atoms=atoms, temperature_K=self.start_temp) if self.zero_linear_momentum: Stationary(atoms) @@ -211,11 +209,11 @@ def make( md_observer = TrajectoryObserver(atoms, store_md_outputs=True) md_runner = md_func( - atoms=atoms, timestep=self.timestep * fs, **self.ase_md_kwargs + atoms=atoms, timestep=self.time_step * fs, **self.ase_md_kwargs ) md_runner.attach(md_observer, interval=self.traj_interval) - md_runner.run(steps=self.md_steps) + md_runner.run(steps=self.nsteps) if self.traj_file is not None: md_observer.save(self.traj_file) @@ -227,7 +225,7 @@ def make( self.force_field_name, {"final_structure": structure, "trajectory": md_observer}, relax_cell=(self.ensemble == "npt"), - steps=self.md_steps, + steps=self.nsteps, relax_kwargs=None, optimizer_kwargs=None, **self.task_document_kwargs, diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index dbd5244d53..015b568000 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -13,7 +13,7 @@ @pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE"]) def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): - md_steps = 5 + nsteps = 5 ref_energies_per_atom = { "CHGNet": -5.30686092376709, @@ -26,9 +26,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): # MACE changes the default dtype, ensure consistent dtype here torch.set_default_dtype(torch.float32) - job = _to_maker[ff_name](md_steps=md_steps, traj_file="md_traj.json.gz").make( - structure - ) + job = _to_maker[ff_name](nsteps=nsteps, traj_file="md_traj.json.gz").make(structure) response = run_locally(job, ensure_success=True) taskdoc = response[next(iter(response))][1].output @@ -40,7 +38,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): # Check that we have the right number of MD steps # ASE logs the initial structure energy, and doesn"t count this as an MD step assert taskdoc.output.ionic_steps[0].structure == structure - assert len(taskdoc.output.ionic_steps) == md_steps + 1 + assert len(taskdoc.output.ionic_steps) == nsteps + 1 # Check that the ionic steps have the expected physical properties assert all( @@ -51,7 +49,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): # Check that the trajectory has expected physical properties assert taskdoc.included_objects == ["trajectory"] - assert len(taskdoc.forcefield_objects["trajectory"]) == md_steps + 1 + assert len(taskdoc.forcefield_objects["trajectory"]) == nsteps + 1 assert all( key in step for key in ("energy", "forces", "stress", "velocities", "temperature") @@ -62,7 +60,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): # stored to the task document traj_from_file = loadfn("md_traj.json.gz") - assert len(traj_from_file["energy"]) == md_steps + 1 + assert len(traj_from_file["energy"]) == nsteps + 1 _traj_key_to_object_key = { "stresses": "stress", } @@ -74,5 +72,5 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): ] ) for key in ("energy", "temperature", "forces", "velocities") - for idx in range(md_steps + 1) + for idx in range(nsteps + 1) ) From 1ec5910da2121079501918ad22f42cc5aad21735 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 22 Feb 2024 15:23:50 -0800 Subject: [PATCH 09/48] Add option to explicitly specify thermostat as ASE MolecularDynamics object --- src/atomate2/forcefields/md.py | 57 ++++++++++++++++++++-------------- 1 file changed, 34 insertions(+), 23 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index a1f4389787..71cc37a09e 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -8,6 +8,7 @@ from ase.md.andersen import Andersen from ase.md.langevin import Langevin +from ase.md.md import MolecularDynamics from ase.md.npt import NPT from ase.md.nptberendsen import NPTBerendsen from ase.md.nvtberendsen import NVTBerendsen @@ -25,6 +26,7 @@ from atomate2.forcefields.utils import TrajectoryObserver if TYPE_CHECKING: + from collections.abc import Sequence from pathlib import Path from typing import Literal @@ -72,15 +74,20 @@ class ForceFieldMDMaker(Maker): The number of MD steps to run ensemble : str = "nvt" The ensemble to use. Valid ensembles are nve, nvt, or npt - start_temp : float | None = 300. + start_temp : float | Sequence | None = 300. The temperature to initialize the system, in Kelvin. end_temp : float | None = 300. The temperature to equilibrate towards, in Kelvin. If start_temp is a float and end_temp is None, the system will be initialized at and equilibrated towards the start_temp. - thermostat : str = "langevin" - The thermostat to use. See _valid_thermostats for a list of options + pressure: float | Sequence | None = None + The pressure in kilobar. + thermostat : str | ASE .MolecularDynamics = "langevin" + The thermostat to use. If thermostat is an ASE .MolecularDynamics + object, this uses the option specified explicitly by the user. + See _valid_thermostats for a list of pre-defined options when + specifying thermostat as a string. ase_md_kwargs : dict | None = None Options to pass to the ASE MD function calculator_args : Sequence | None = None @@ -107,10 +114,10 @@ class ForceFieldMDMaker(Maker): time_step: float | None = 2.0 nsteps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" - start_temp: float | None = 300.0 + start_temp: float | Sequence | None = 300.0 end_temp: float | None = 300.0 - pressure: float | None = None - thermostat: str = "langevin" + pressure: float | Sequence | None = None + thermostat: str | MolecularDynamics = "langevin" ase_md_kwargs: dict | None = None calculator_args: list | tuple | None = None calculator_kwargs: dict | None = None @@ -122,7 +129,7 @@ class ForceFieldMDMaker(Maker): default_factory=lambda: {"store_trajectory": "partial"} ) - def _get_ensemble_defaults(self, structure: Structure) -> None: + def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" self.ase_md_kwargs = self.ase_md_kwargs or {} if self.ensemble in ("nvt", "npt") and all( @@ -133,11 +140,9 @@ def _get_ensemble_defaults(self, structure: Structure) -> None: self.end_temp if self.end_temp else self.start_temp ) - if self.ensemble == "npt": + if self.ensemble == "npt" and isinstance(self.pressure, float): # convert from kilobar to eV/Ang**3 - self.ase_md_kwargs["pressure_au"] = ( - self.pressure * 1.0e-3 / bar if self.pressure else 0.0 - ) + self.ase_md_kwargs["pressure_au"] = self.pressure * 1.0e-3 / bar if self.thermostat == "langevin": # Same default as in VASP @@ -167,21 +172,27 @@ def make( A previous calculation directory to copy output files from. Unused, just added to match the method signature of other makers. """ - self._get_ensemble_defaults(structure=structure) + self._get_ensemble_defaults() initial_velocities = structure.site_properties.get("velocities") - self.thermostat = self.thermostat.lower() - if self.thermostat not in _valid_thermostats[self.ensemble]: - raise ValueError( - f"{self.thermostat} thermostat not available for {self.ensemble}." - f"Available {self.ensemble} thermostats are:" - " ".join(_valid_thermostats[self.ensemble]) - ) - - if self.ensemble == "nve" and self.thermostat is None: - self.thermostat = "velocityverlet" - md_func = _thermostats[f"{self.ensemble}_{self.thermostat}"] + if isinstance(self.thermostat, MolecularDynamics): + # Allow user to explicitly set dynamics run via thermostat + md_func = self.thermostat + + elif isinstance(self.thermostat, str): + # Otherwise, use known thermostat + self.thermostat = self.thermostat.lower() + if self.thermostat not in _valid_thermostats[self.ensemble]: + raise ValueError( + f"{self.thermostat} thermostat not available for {self.ensemble}." + f"Available {self.ensemble} thermostats are:" + " ".join(_valid_thermostats[self.ensemble]) + ) + + if self.ensemble == "nve" and self.thermostat is None: + self.thermostat = "velocityverlet" + md_func = _thermostats[f"{self.ensemble}_{self.thermostat}"] atoms = AseAtomsAdaptor.get_atoms(structure) if initial_velocities: From 1e4fcb76bc3b5c290f5534d425e1b8a95a05d6d8 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Fri, 23 Feb 2024 13:44:34 -0800 Subject: [PATCH 10/48] refactoring and add T/P schedule --- src/atomate2/forcefields/md.py | 65 ++++++++++++++++++++++++---------- 1 file changed, 46 insertions(+), 19 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 71cc37a09e..50b5bbac99 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -3,9 +3,11 @@ import contextlib import io +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING +import numpy as np from ase.md.andersen import Andersen from ase.md.langevin import Langevin from ase.md.md import MolecularDynamics @@ -21,12 +23,12 @@ from ase.units import bar, fs from jobflow import Maker, job from pymatgen.io.ase import AseAtomsAdaptor +from scipy.interpolate import interp1d from atomate2.forcefields.schemas import ForceFieldTaskDocument from atomate2.forcefields.utils import TrajectoryObserver if TYPE_CHECKING: - from collections.abc import Sequence from pathlib import Path from typing import Literal @@ -68,8 +70,8 @@ class ForceFieldMDMaker(Maker): The name of the MD Maker force_field_name : str The name of the forcefield (for provenance) - time_step : float | None = 2. - The time_step of the MD run in femtoseconds + timestep : float | None = 2. + The timestep of the MD run in femtoseconds nsteps : int = 1000 The number of MD steps to run ensemble : str = "nvt" @@ -111,13 +113,13 @@ class ForceFieldMDMaker(Maker): name: str = "Forcefield MD" force_field_name: str = "Forcefield" - time_step: float | None = 2.0 + timestep: float | None = 2.0 nsteps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" - start_temp: float | Sequence | None = 300.0 + dynamics: str | MolecularDynamics = "langevin" + temperature: float | Sequence | None = 300.0 end_temp: float | None = 300.0 - pressure: float | Sequence | None = None - thermostat: str | MolecularDynamics = "langevin" + pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None calculator_args: list | tuple | None = None calculator_kwargs: dict | None = None @@ -129,6 +131,31 @@ class ForceFieldMDMaker(Maker): default_factory=lambda: {"store_trajectory": "partial"} ) + def _get_ensemble_schedule(self) -> None: + if isinstance(self.temperature, Sequence): + self.tschedule = np.interp( + np.arange(self.nsteps), + np.arange(len(self.temperature)), + self.temperature + ) + else: + self.tschedule = np.full(self.nsteps, self.temperature) + + if isinstance(self.pressure, Sequence): + self.pschedule = np.interp( + np.arange(self.nsteps), + np.arange(len(self.pressure)), + self.pressure + ) + elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: + self.pschedule = interp1d( + np.arange(self.nsteps), + self.pressure, + kind="linear" + ) + else: + self.pschedule = np.full(self.nsteps, self.pressure) + def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" self.ase_md_kwargs = self.ase_md_kwargs or {} @@ -144,7 +171,7 @@ def _get_ensemble_defaults(self) -> None: # convert from kilobar to eV/Ang**3 self.ase_md_kwargs["pressure_au"] = self.pressure * 1.0e-3 / bar - if self.thermostat == "langevin": + if self.dynamics == "langevin": # Same default as in VASP self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get("friction", 10.0) # friction coefficient(s) specified in ev^-1, convert from picoseconds @@ -176,25 +203,25 @@ def make( initial_velocities = structure.site_properties.get("velocities") - if isinstance(self.thermostat, MolecularDynamics): + if isinstance(self.dynamics, MolecularDynamics): # Allow user to explicitly set dynamics run via thermostat - md_func = self.thermostat + md_func = self.dynamics - elif isinstance(self.thermostat, str): + elif isinstance(self.dynamics, str): # Otherwise, use known thermostat - self.thermostat = self.thermostat.lower() - if self.thermostat not in _valid_thermostats[self.ensemble]: + self.dynamics = self.dynamics.lower() + if self.dynamics not in _valid_thermostats[self.ensemble]: raise ValueError( - f"{self.thermostat} thermostat not available for {self.ensemble}." + f"{self.dynamics} thermostat not available for {self.ensemble}." f"Available {self.ensemble} thermostats are:" " ".join(_valid_thermostats[self.ensemble]) ) - if self.ensemble == "nve" and self.thermostat is None: - self.thermostat = "velocityverlet" - md_func = _thermostats[f"{self.ensemble}_{self.thermostat}"] + if self.ensemble == "nve" and self.dynamics is None: + self.dynamics = "velocityverlet" + md_func = _thermostats[f"{self.ensemble}_{self.dynamics}"] - atoms = AseAtomsAdaptor.get_atoms(structure) + atoms = structure.to_ase_atoms() if initial_velocities: atoms.set_velocities(initial_velocities) @@ -220,7 +247,7 @@ def make( md_observer = TrajectoryObserver(atoms, store_md_outputs=True) md_runner = md_func( - atoms=atoms, timestep=self.time_step * fs, **self.ase_md_kwargs + atoms=atoms, timestep=self.timestep * fs, **self.ase_md_kwargs ) md_runner.attach(md_observer, interval=self.traj_interval) From 317412c6ebcac2f50288e9f4ec9ee294d7fe148a Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Fri, 23 Feb 2024 19:21:29 -0800 Subject: [PATCH 11/48] implement T/P schedule and attach callback func --- src/atomate2/forcefields/md.py | 135 +++++++++++++++++++++++---------- 1 file changed, 93 insertions(+), 42 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 50b5bbac99..891f28bb78 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING import numpy as np +from ase import units from ase.md.andersen import Andersen from ase.md.langevin import Langevin from ase.md.md import MolecularDynamics @@ -20,7 +21,6 @@ ZeroRotation, ) from ase.md.verlet import VelocityVerlet -from ase.units import bar, fs from jobflow import Maker, job from pymatgen.io.ase import AseAtomsAdaptor from scipy.interpolate import interp1d @@ -35,8 +35,7 @@ from ase.calculators.calculator import Calculator from pymatgen.core.structure import Structure - -_valid_thermostats: dict[str, tuple[str, ...]] = { +_valid_dynamics: dict[str, tuple[str, ...]] = { "nve": ("velocityverlet",), "nvt": ("nose-hoover", "langevin", "andersen", "berendsen"), "npt": ("nose-hoover", "berendsen"), @@ -88,10 +87,10 @@ class ForceFieldMDMaker(Maker): thermostat : str | ASE .MolecularDynamics = "langevin" The thermostat to use. If thermostat is an ASE .MolecularDynamics object, this uses the option specified explicitly by the user. - See _valid_thermostats for a list of pre-defined options when + See _valid_dynamics for a list of pre-defined options when specifying thermostat as a string. ase_md_kwargs : dict | None = None - Options to pass to the ASE MD function + Options except for temperature and pressure to pass into the ASE MD function calculator_args : Sequence | None = None args to pass to the ASE calculator class calculator_kwargs : dict | None = None @@ -117,7 +116,7 @@ class ForceFieldMDMaker(Maker): nsteps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" dynamics: str | MolecularDynamics = "langevin" - temperature: float | Sequence | None = 300.0 + temperature: float | Sequence | np.ndarray | None = 300.0 end_temp: float | None = 300.0 pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None @@ -132,15 +131,34 @@ class ForceFieldMDMaker(Maker): ) def _get_ensemble_schedule(self) -> None: + + if self.ensemble == "nve": + # Distable thermostat and barostat + self.temperature = np.nan + self.pressure = np.nan + self.tschedule = np.full(self.nsteps, self.temperature) + self.pschedule = np.full(self.nsteps, self.pressure) + return + if isinstance(self.temperature, Sequence): self.tschedule = np.interp( np.arange(self.nsteps), np.arange(len(self.temperature)), self.temperature ) + elif isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1: + self.tschedule = self.temperature + # NOTE: In ASE Langevin dynamics, the temperature are normally + # scalars, but in principle one quantity per atom could be specified by giving + # an array. This is not implemented yet here. else: self.tschedule = np.full(self.nsteps, self.temperature) + if self.ensemble == "nvt": + self.pressure = None + self.pschedule = np.full(self.nsteps, self.pressure) + return + if isinstance(self.pressure, Sequence): self.pschedule = np.interp( np.arange(self.nsteps), @@ -159,28 +177,51 @@ def _get_ensemble_schedule(self) -> None: def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" self.ase_md_kwargs = self.ase_md_kwargs or {} - if self.ensemble in ("nvt", "npt") and all( - self.ase_md_kwargs.get(key) is None - for key in ("temperature_K", "temperature") - ): - self.ase_md_kwargs["temperature_K"] = ( - self.end_temp if self.end_temp else self.start_temp + + if self.ensemble == "nve": + self.ase_md_kwargs.pop("temperature", None) + self.ase_md_kwargs.pop("temperature_K", None) + self.ase_md_kwargs.pop("externalstress", None) + elif self.ensemble == "nvt": + self.ase_md_kwargs["temperature_K"] = self.tschedule[0] + self.ase_md_kwargs.pop("externalstress", None) + elif self.ensemble == "npt": + self.ase_md_kwargs["temperature_K"] = self.tschedule[0] + self.ase_md_kwargs["externalstress"] = ( + self.pschedule[0] * 1.0e-3 / units.bar ) - if self.ensemble == "npt" and isinstance(self.pressure, float): - # convert from kilobar to eV/Ang**3 - self.ase_md_kwargs["pressure_au"] = self.pressure * 1.0e-3 / bar - - if self.dynamics == "langevin": - # Same default as in VASP - self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get("friction", 10.0) - # friction coefficient(s) specified in ev^-1, convert from picoseconds - if isinstance(self.ase_md_kwargs["friction"], (list, tuple)): - self.ase_md_kwargs["friction"] = [ - coeff * 1.0e-3 / fs for coeff in self.ase_md_kwargs["friction"] - ] - else: - self.ase_md_kwargs["friction"] *= 1.0e-3 / fs + + # NOTE: We take of the temperature in _get_ensemble_schedule instead + # if self.ensemble in ("nvt", "npt") and all( + # self.ase_md_kwargs.get(key) is None + # for key in ("temperature_K", "temperature") + # ): + # self.ase_md_kwargs["temperature_K"] = ( + # self.end_temp if self.end_temp else self.start_temp + # ) + + # NOTE: We take care of the units when passing into the ASE MD function + # if self.ensemble == "npt" and isinstance(self.pressure, float): + # # convert from kilobar to eV/Ang**3 + # self.ase_md_kwargs["pressure_au"] = self.pressure * 1.0e-3 / bar + + if self.dynamics.lower() == "langevin": + # NOTE: Unless we have a detailed documentation on the conversion of all + # parameters for all different ASE dyanmics, it is not a good idea to + # convert the parameters here. It is better to let the user to consult + # the ASE documentation and set the parameters themselves. + self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get( + "friction", 10.0 * 1e-3 / units.fs # Same default as in VASP: 10 ps^-1 + ) + # NOTE: same as above, user can specify per atom friction but we don't + # expect to change their intention to pass into ASE MD function + # if isinstance(self.ase_md_kwargs["friction"], (list, tuple)): + # self.ase_md_kwargs["friction"] = [ + # coeff * 1.0e-3 / fs for coeff in self.ase_md_kwargs["friction"] + # ] + # else: + # self.ase_md_kwargs["friction"] *= 1.0e-3 / fs @job(output_schema=ForceFieldTaskDocument) def make( @@ -199,22 +240,23 @@ def make( A previous calculation directory to copy output files from. Unused, just added to match the method signature of other makers. """ + self._get_ensemble_schedule() self._get_ensemble_defaults() initial_velocities = structure.site_properties.get("velocities") if isinstance(self.dynamics, MolecularDynamics): - # Allow user to explicitly set dynamics run via thermostat + # Allow user to explicitly run ASE Dynamics class md_func = self.dynamics elif isinstance(self.dynamics, str): - # Otherwise, use known thermostat + # Otherwise, use known dynamics self.dynamics = self.dynamics.lower() - if self.dynamics not in _valid_thermostats[self.ensemble]: + if self.dynamics not in _valid_dynamics[self.ensemble]: raise ValueError( f"{self.dynamics} thermostat not available for {self.ensemble}." f"Available {self.ensemble} thermostats are:" - " ".join(_valid_thermostats[self.ensemble]) + " ".join(_valid_dynamics[self.ensemble]) ) if self.ensemble == "nve" and self.dynamics is None: @@ -224,33 +266,42 @@ def make( atoms = structure.to_ase_atoms() if initial_velocities: atoms.set_velocities(initial_velocities) - - elif self.start_temp: - MaxwellBoltzmannDistribution(atoms=atoms, temperature_K=self.start_temp) - + elif not np.isnan(self.tschedule).any(): + MaxwellBoltzmannDistribution( + atoms=atoms, + temperature_K=self.tschedule[0], + rng=None # TODO: we might want to use a seed for reproducibility + ) if self.zero_linear_momentum: Stationary(atoms) - if self.zero_angular_momentum: ZeroRotation(atoms) - else: - atoms.set_velocities( - [[0.0 for _ in range(3)] for _ in range(len(structure))] - ) + # NOTE: ase has zero velocities by default already + pass self.calculator_args = self.calculator_args or [] self.calculator_kwargs = self.calculator_kwargs or {} - atoms.set_calculator(self._calculator()) + atoms.calc = self._calculator() + + def callback(dyn: MolecularDynamics = md_func) -> None: + if self.ensemble == "nve": + return + dyn.set_temperature(self.tschedule[dyn.nsteps]) + if self.ensemble == "nvt": + return + dyn.set_stress(self.pschedule[dyn.nsteps]) with contextlib.redirect_stdout(io.StringIO()): md_observer = TrajectoryObserver(atoms, store_md_outputs=True) md_runner = md_func( - atoms=atoms, timestep=self.timestep * fs, **self.ase_md_kwargs + atoms=atoms, + timestep=self.timestep * units.fs, + **self.ase_md_kwargs ) md_runner.attach(md_observer, interval=self.traj_interval) - + md_runner.attach(callback, interval=1) md_runner.run(steps=self.nsteps) if self.traj_file is not None: From 76ffa66d879898a521ba9c22415eee93d840a9e2 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Fri, 23 Feb 2024 19:27:07 -0800 Subject: [PATCH 12/48] np.nan --- src/atomate2/forcefields/md.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 891f28bb78..b8f4ed095d 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -155,7 +155,7 @@ def _get_ensemble_schedule(self) -> None: self.tschedule = np.full(self.nsteps, self.temperature) if self.ensemble == "nvt": - self.pressure = None + self.pressure = np.nan self.pschedule = np.full(self.nsteps, self.pressure) return From 8253a35a1d3d74397b51bf7817bd4dcedb5cc4fc Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Fri, 23 Feb 2024 19:27:07 -0800 Subject: [PATCH 13/48] np.nan --- src/atomate2/forcefields/md.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 891f28bb78..b8f4ed095d 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -155,7 +155,7 @@ def _get_ensemble_schedule(self) -> None: self.tschedule = np.full(self.nsteps, self.temperature) if self.ensemble == "nvt": - self.pressure = None + self.pressure = np.nan self.pschedule = np.full(self.nsteps, self.pressure) return From 298250d9f1379ac1479895a93b72e138e4c22123 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Mon, 26 Feb 2024 19:56:18 -0800 Subject: [PATCH 14/48] small fixes --- src/atomate2/forcefields/md.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index b8f4ed095d..d0f6483d38 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -41,7 +41,7 @@ "npt": ("nose-hoover", "berendsen"), } -_thermostats: dict = { +_preset_dynamics: dict = { "nve_velocityverlet": VelocityVerlet, "nvt_andersen": Andersen, "nvt_berendsen": NVTBerendsen, @@ -261,7 +261,7 @@ def make( if self.ensemble == "nve" and self.dynamics is None: self.dynamics = "velocityverlet" - md_func = _thermostats[f"{self.ensemble}_{self.dynamics}"] + md_func = _preset_dynamics[f"{self.ensemble}_{self.dynamics}"] atoms = structure.to_ase_atoms() if initial_velocities: @@ -290,7 +290,7 @@ def callback(dyn: MolecularDynamics = md_func) -> None: dyn.set_temperature(self.tschedule[dyn.nsteps]) if self.ensemble == "nvt": return - dyn.set_stress(self.pschedule[dyn.nsteps]) + dyn.set_stress(self.pschedule[dyn.nsteps] * 1.0e-3 / units.bar) with contextlib.redirect_stdout(io.StringIO()): md_observer = TrajectoryObserver(atoms, store_md_outputs=True) @@ -330,7 +330,7 @@ class MACEMDMaker(ForceFieldMDMaker): name: str = "MACE MD" force_field_name: str = "MACE" - calculator_kwargs: dict = field(default_factory=lambda: {"dtype": "float32"}) + calculator_kwargs: dict = field(default_factory=lambda: {"default_dtype": "float32"}) def _calculator(self) -> Calculator: from mace.calculators import mace_mp From 0045b2bab7fed8f717e6de82444e97e842102b26 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Tue, 27 Feb 2024 15:18:26 -0800 Subject: [PATCH 15/48] pass ASE dynamics object instead of class, update docstring --- src/atomate2/forcefields/md.py | 72 +++++++++++++++++++--------------- 1 file changed, 41 insertions(+), 31 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index d0f6483d38..50446e4c7a 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -57,12 +57,20 @@ class ForceFieldMDMaker(Maker): """ Perform MD with a forcefield. - Note that units are consistent with the VASP MD implementation. + Note the the following units are consistent with the VASP MD implementation: - Temperature is in Kelvin (TEBEG and TEEND) - Time steps are in femtoseconds (POTIM) - - Langevin NVT friction coefficients are in picoseconds^-1 (LANGEVIN_GAMMA) - Pressure in kB (PSTRESS) + The default dynamics is Langevin NVT consistent with VASP MD, with the friction + coefficient set to 10 ps^-1 (LANGEVIN_GAMMA). + + For the rest of preset dynamics (`_valid_dynamics`) and custom dynamics inherited + from ASE (`MolecularDynamics`), the user can specify the dynamics as a string or an + ASE class into the `dynamics` attribute. In this case, please consult the ASE + documentation for the parameters and units to pass into the ASE MD function through + `ase_md_kwargs`. + Parameters ---------- name : str @@ -75,15 +83,14 @@ class ForceFieldMDMaker(Maker): The number of MD steps to run ensemble : str = "nvt" The ensemble to use. Valid ensembles are nve, nvt, or npt - start_temp : float | Sequence | None = 300. - The temperature to initialize the system, in Kelvin. - end_temp : float | None = 300. - The temperature to equilibrate towards, in Kelvin. - If start_temp is a float and end_temp is None, - the system will be initialized at and equilibrated towards - the start_temp. + temperature: float | Sequence | np.ndarray | None. + The temperature in Kelvin. If a sequence or 1D array, the temperature + schedule will be interpolated linearly between the given values. If a + float, the temperature will be constant throughout the run. pressure: float | Sequence | None = None - The pressure in kilobar. + The pressure in kilobar. If a sequence or 1D array, the pressure + schedule will be interpolated linearly between the given values. If a + float, the pressure will be constant throughout the run. thermostat : str | ASE .MolecularDynamics = "langevin" The thermostat to use. If thermostat is an ASE .MolecularDynamics object, this uses the option specified explicitly by the user. @@ -117,7 +124,7 @@ class ForceFieldMDMaker(Maker): ensemble: Literal["nve", "nvt", "npt"] = "nvt" dynamics: str | MolecularDynamics = "langevin" temperature: float | Sequence | np.ndarray | None = 300.0 - end_temp: float | None = 300.0 + # end_temp: float | None = 300.0 pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None calculator_args: list | tuple | None = None @@ -136,43 +143,45 @@ def _get_ensemble_schedule(self) -> None: # Distable thermostat and barostat self.temperature = np.nan self.pressure = np.nan - self.tschedule = np.full(self.nsteps, self.temperature) - self.pschedule = np.full(self.nsteps, self.pressure) + self.tschedule = np.full(self.nsteps+1, self.temperature) + self.pschedule = np.full(self.nsteps+1, self.pressure) return - if isinstance(self.temperature, Sequence): + if isinstance(self.temperature, Sequence) or ( + isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1 + ): self.tschedule = np.interp( - np.arange(self.nsteps), + np.arange(self.nsteps+1), np.arange(len(self.temperature)), self.temperature ) - elif isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1: - self.tschedule = self.temperature # NOTE: In ASE Langevin dynamics, the temperature are normally # scalars, but in principle one quantity per atom could be specified by giving # an array. This is not implemented yet here. else: - self.tschedule = np.full(self.nsteps, self.temperature) + self.tschedule = np.full(self.nsteps+1, self.temperature) if self.ensemble == "nvt": self.pressure = np.nan - self.pschedule = np.full(self.nsteps, self.pressure) + self.pschedule = np.full(self.nsteps+1, self.pressure) return - if isinstance(self.pressure, Sequence): + if isinstance(self.pressure, Sequence) or ( + isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 1 + ): self.pschedule = np.interp( - np.arange(self.nsteps), + np.arange(self.nsteps+1), np.arange(len(self.pressure)), self.pressure ) elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: self.pschedule = interp1d( - np.arange(self.nsteps), + np.arange(self.nsteps+1), self.pressure, kind="linear" ) else: - self.pschedule = np.full(self.nsteps, self.pressure) + self.pschedule = np.full(self.nsteps+1, self.pressure) def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" @@ -284,14 +293,6 @@ def make( self.calculator_kwargs = self.calculator_kwargs or {} atoms.calc = self._calculator() - def callback(dyn: MolecularDynamics = md_func) -> None: - if self.ensemble == "nve": - return - dyn.set_temperature(self.tschedule[dyn.nsteps]) - if self.ensemble == "nvt": - return - dyn.set_stress(self.pschedule[dyn.nsteps] * 1.0e-3 / units.bar) - with contextlib.redirect_stdout(io.StringIO()): md_observer = TrajectoryObserver(atoms, store_md_outputs=True) @@ -300,6 +301,15 @@ def callback(dyn: MolecularDynamics = md_func) -> None: timestep=self.timestep * units.fs, **self.ase_md_kwargs ) + + def callback(dyn: MolecularDynamics = md_runner) -> None: + if self.ensemble == "nve": + return + dyn.set_temperature(temperature_K=self.tschedule[dyn.nsteps]) + if self.ensemble == "nvt": + return + dyn.set_stress(self.pschedule[dyn.nsteps] * 1.0e-3 / units.bar) + md_runner.attach(md_observer, interval=self.traj_interval) md_runner.attach(callback, interval=1) md_runner.run(steps=self.nsteps) From b870c0407f85cccc0f84928ce41c504e66b9bfe6 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Tue, 27 Feb 2024 17:20:29 -0800 Subject: [PATCH 16/48] add temp ramp test and make sure upper tria matrix cell --- src/atomate2/forcefields/md.py | 6 ++++++ tests/forcefields/test_md.py | 27 +++++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 50446e4c7a..1a3a5c4c70 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -24,6 +24,7 @@ from jobflow import Maker, job from pymatgen.io.ase import AseAtomsAdaptor from scipy.interpolate import interp1d +from scipy.linalg import schur from atomate2.forcefields.schemas import ForceFieldTaskDocument from atomate2.forcefields.utils import TrajectoryObserver @@ -273,6 +274,10 @@ def make( md_func = _preset_dynamics[f"{self.ensemble}_{self.dynamics}"] atoms = structure.to_ase_atoms() + + u, q = schur(atoms.get_cell(complete=True)) + atoms.set_cell(u, scale_atoms=True) + if initial_velocities: atoms.set_velocities(initial_velocities) elif not np.isnan(self.tschedule).any(): @@ -299,6 +304,7 @@ def make( md_runner = md_func( atoms=atoms, timestep=self.timestep * units.fs, + # trajectory="trajectory.traj", # TODO: we might want implement taskdoc to save ase binary traj **self.ase_md_kwargs ) diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 015b568000..e9ca10d21f 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -3,6 +3,7 @@ import numpy as np import pytest import torch +from ase import units from jobflow import run_locally from monty.serialization import loadfn @@ -74,3 +75,29 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): for key in ("energy", "temperature", "forces", "velocities") for idx in range(nsteps + 1) ) + +@pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) +def test_temp_schedule(ff_name, si_structure, clean_dir): + nsteps = 200 + temp_schedule = [300, 3000] + + structure = si_structure.to_conventional() * (2, 2, 2) + + # MACE changes the default dtype, ensure consistent dtype here + torch.set_default_dtype(torch.float32) + + job = _to_maker[ff_name]( + nsteps=nsteps, traj_file="md_traj.json.gz", + dynamics="nose-hoover", + temperature=temp_schedule, + ase_md_kwargs=dict(ttime=50.0 * units.fs, pfactor=None) + ).make(structure) + response = run_locally(job, ensure_success=True) + taskdoc = response[next(iter(response))][1].output + + + temp_history = [ + step["temperature"] for step in taskdoc.forcefield_objects["trajectory"].frame_properties + ] + + assert temp_history[-1] > temp_schedule[0] From 71728a34a7d3172bbc61fe4a8f95aea488fd5a47 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Wed, 28 Feb 2024 15:17:09 -0800 Subject: [PATCH 17/48] add stresses to trajeoctroy observer, add pressure schedule test --- src/atomate2/forcefields/utils.py | 10 +++++++--- tests/forcefields/test_md.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index d3735d6528..ad49f42ecc 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -87,6 +87,7 @@ def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: if self._store_md_outputs: self.velocities: list[np.ndarray] = [] self.temperatures: list[float] = [] + self.stresses: list[np.ndarray] = [] def __call__(self) -> None: """Save the properties of an Atoms during the relaxation.""" @@ -100,6 +101,7 @@ def __call__(self) -> None: if self._store_md_outputs: self.velocities.append(self.atoms.get_velocities()) self.temperatures.append(self.atoms.get_temperature()) + self.stresses.append(self.atoms.get_stress(voigt=True, include_ideal_gas=True)) def compute_energy(self) -> float: """ @@ -136,9 +138,11 @@ def as_dict(self) -> dict: "atomic_number": self.atoms.get_atomic_numbers(), } if self._store_md_outputs: - traj_dict.update( - {"velocities": self.velocities, "temperature": self.temperatures} - ) + traj_dict.update({ + "velocities": self.velocities, + "temperature": self.temperatures, + "stresses": self.stresses, + }) # sanitize dict for key in traj_dict: if all(isinstance(val, np.ndarray) for val in traj_dict[key]): diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index e9ca10d21f..a1693c2af3 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -101,3 +101,32 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): ] assert temp_history[-1] > temp_schedule[0] + + +@pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) +def test_press_schedule(ff_name, si_structure, clean_dir): + nsteps = 200 + press_schedule = [0, 1e4] + + structure = si_structure.to_conventional() * (2, 2, 2) + + # MACE changes the default dtype, ensure consistent dtype here + torch.set_default_dtype(torch.float32) + + job = _to_maker[ff_name]( + nsteps=nsteps, traj_file="md_traj.json.gz", + dynamics="nose-hoover", + pressure=press_schedule, + ase_md_kwargs=dict( + ttime=50.0 * units.fs, + pfactor=(75.0 * units.fs)**2 * units.GPa, + ) + ).make(structure) + response = run_locally(job, ensure_success=True) + taskdoc = response[next(iter(response))][1].output + + stress_history = [ + step["stresses"] for step in taskdoc.forcefield_objects["trajectory"].frame_properties + ] + + assert stress_history[-1].trace() > stress_history[0].trace() From 45ddaa1c246c2e1c3d2dc3ffe8b52622d266483a Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Wed, 28 Feb 2024 20:05:32 -0800 Subject: [PATCH 18/48] avoid `stresses` key not in taskdoc bug --- src/atomate2/forcefields/md.py | 4 ++-- src/atomate2/forcefields/utils.py | 9 +++++---- tests/forcefields/test_md.py | 19 +++++++++++-------- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 1a3a5c4c70..67375ddf27 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -198,7 +198,7 @@ def _get_ensemble_defaults(self) -> None: elif self.ensemble == "npt": self.ase_md_kwargs["temperature_K"] = self.tschedule[0] self.ase_md_kwargs["externalstress"] = ( - self.pschedule[0] * 1.0e-3 / units.bar + self.pschedule[0] * 1e3 * units.bar ) @@ -314,7 +314,7 @@ def callback(dyn: MolecularDynamics = md_runner) -> None: dyn.set_temperature(temperature_K=self.tschedule[dyn.nsteps]) if self.ensemble == "nvt": return - dyn.set_stress(self.pschedule[dyn.nsteps] * 1.0e-3 / units.bar) + dyn.set_stress(self.pschedule[dyn.nsteps] * 1e3 * units.bar) md_runner.attach(md_observer, interval=self.traj_interval) md_runner.attach(callback, interval=1) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index ad49f42ecc..b662469bbb 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -87,21 +87,23 @@ def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: if self._store_md_outputs: self.velocities: list[np.ndarray] = [] self.temperatures: list[float] = [] - self.stresses: list[np.ndarray] = [] def __call__(self) -> None: """Save the properties of an Atoms during the relaxation.""" # TODO: maybe include magnetic moments self.energies.append(self.compute_energy()) self.forces.append(self.atoms.get_forces()) - self.stresses.append(self.atoms.get_stress()) + # TODO: MD needs kinetic energy parts of stress, we might just need to refactor TrajectoryObserver + # Now it is safe for 0K relaxation as the atoms don't have momenta + self.stresses.append(self.atoms.get_stress(include_ideal_gas=True)) self.atom_positions.append(self.atoms.get_positions()) self.cells.append(self.atoms.get_cell()[:]) if self._store_md_outputs: self.velocities.append(self.atoms.get_velocities()) self.temperatures.append(self.atoms.get_temperature()) - self.stresses.append(self.atoms.get_stress(voigt=True, include_ideal_gas=True)) + # self.stresses.append(self.atoms.get_stress(voigt=True, include_ideal_gas=True)) + def compute_energy(self) -> float: """ @@ -141,7 +143,6 @@ def as_dict(self) -> dict: traj_dict.update({ "velocities": self.velocities, "temperature": self.temperatures, - "stresses": self.stresses, }) # sanitize dict for key in traj_dict: diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index a1693c2af3..446889148c 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -78,7 +78,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_temp_schedule(ff_name, si_structure, clean_dir): - nsteps = 200 + nsteps = 100 temp_schedule = [300, 3000] structure = si_structure.to_conventional() * (2, 2, 2) @@ -105,15 +105,16 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_press_schedule(ff_name, si_structure, clean_dir): - nsteps = 200 - press_schedule = [0, 1e4] + nsteps = 100 + press_schedule = [0, 10] # kbar - structure = si_structure.to_conventional() * (2, 2, 2) + structure = si_structure.to_conventional() * (3, 3, 3) # MACE changes the default dtype, ensure consistent dtype here torch.set_default_dtype(torch.float32) job = _to_maker[ff_name]( + ensemble="npt", nsteps=nsteps, traj_file="md_traj.json.gz", dynamics="nose-hoover", pressure=press_schedule, @@ -122,11 +123,13 @@ def test_press_schedule(ff_name, si_structure, clean_dir): pfactor=(75.0 * units.fs)**2 * units.GPa, ) ).make(structure) - response = run_locally(job, ensure_success=True) - taskdoc = response[next(iter(response))][1].output + run_locally(job, ensure_success=True) + # taskdoc = response[next(iter(response))][1].output + + traj_from_file = loadfn("md_traj.json.gz") stress_history = [ - step["stresses"] for step in taskdoc.forcefield_objects["trajectory"].frame_properties + (stress[0] + stress[1] + stress[2])/3.0 for stress in traj_from_file["stresses"] ] - assert stress_history[-1].trace() > stress_history[0].trace() + assert stress_history[-1] < stress_history[0] From 74bf6a0f9d463223b50fb8791dfc1f41c917525f Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 1 Mar 2024 12:03:30 -0800 Subject: [PATCH 19/48] Refactory TrajectoryObserver, add option to save trajectory as either pmg or ase trajectory objects, add test for trajectory parsing --- src/atomate2/forcefields/jobs.py | 44 +++++++++- src/atomate2/forcefields/md.py | 77 +++++++++--------- src/atomate2/forcefields/schemas.py | 105 +++++++----------------- src/atomate2/forcefields/utils.py | 109 ++++++++++++++++++++++--- tests/forcefields/test_md.py | 122 +++++++++++++++++++--------- 5 files changed, 290 insertions(+), 167 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 8ca0403b0c..294db171ae 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING from jobflow import Maker, job +from pymatgen.core.trajectory import Trajectory from atomate2.forcefields import MLFF from atomate2.forcefields.schemas import ForceFieldTaskDocument @@ -15,11 +16,50 @@ if TYPE_CHECKING: from collections.abc import Sequence from pathlib import Path + from typing import Callable from pymatgen.core.structure import Structure logger = logging.getLogger(__name__) +_FORCEFIELD_DATA_OBJECTS = [Trajectory] + + +def forcefield_job(method: Callable) -> job: + """ + Decorate the ``make`` method of forcefield job makers. + + This is a thin wrapper around :obj:`~jobflow.core.job.Job` that configures common + settings for all forcefield jobs. For example, it ensures that large data objects + (currently only trajectories) are all stored in the atomate2 data store. + It also configures the output schema to be a ForceFieldTaskDocument :obj:`.TaskDoc`. + + Any makers that return forcefield jobs (not flows) should decorate the + ``make`` method with @forcefield_job. For example: + + .. code-block:: python + + class MyForcefieldMaker(Maker): + @forcefield_job + def make(structure): + # code to run forcefield job. + pass + + Parameters + ---------- + method : callable + A Maker.make method. This should not be specified directly and is + implied by the decorator. + + Returns + ------- + callable + A decorated version of the make function that will generate forcefield jobs. + """ + return job( + method, data=_FORCEFIELD_DATA_OBJECTS, output_schema=ForceFieldTaskDocument + ) + @dataclass class ForceFieldRelaxMaker(Maker): @@ -54,7 +94,7 @@ class ForceFieldRelaxMaker(Maker): optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - @job(output_schema=ForceFieldTaskDocument) + @forcefield_job def make( self, structure: Structure, prev_dir: str | Path | None = None ) -> ForceFieldTaskDocument: @@ -110,7 +150,7 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker): force_field_name: str = "Force field" task_document_kwargs: dict = field(default_factory=dict) - @job(output_schema=ForceFieldTaskDocument) + @forcefield_job def make( self, structure: Structure, prev_dir: str | Path | None = None ) -> ForceFieldTaskDocument: diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 67375ddf27..13ca4864bb 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -21,11 +21,12 @@ ZeroRotation, ) from ase.md.verlet import VelocityVerlet -from jobflow import Maker, job +from jobflow import Maker from pymatgen.io.ase import AseAtomsAdaptor from scipy.interpolate import interp1d from scipy.linalg import schur +from atomate2.forcefields.jobs import forcefield_job from atomate2.forcefields.schemas import ForceFieldTaskDocument from atomate2.forcefields.utils import TrajectoryObserver @@ -131,6 +132,7 @@ class ForceFieldMDMaker(Maker): calculator_args: list | tuple | None = None calculator_kwargs: dict | None = None traj_file: str | Path | None = None + traj_file_fmt: Literal["pmg", "ase"] = "ase" traj_interval: int = 1 zero_linear_momentum: bool = False zero_angular_momentum: bool = False @@ -139,50 +141,45 @@ class ForceFieldMDMaker(Maker): ) def _get_ensemble_schedule(self) -> None: - if self.ensemble == "nve": # Distable thermostat and barostat self.temperature = np.nan self.pressure = np.nan - self.tschedule = np.full(self.nsteps+1, self.temperature) - self.pschedule = np.full(self.nsteps+1, self.pressure) + self.tschedule = np.full(self.nsteps + 1, self.temperature) + self.pschedule = np.full(self.nsteps + 1, self.pressure) return if isinstance(self.temperature, Sequence) or ( isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1 - ): + ): self.tschedule = np.interp( - np.arange(self.nsteps+1), + np.arange(self.nsteps + 1), np.arange(len(self.temperature)), - self.temperature + self.temperature, ) # NOTE: In ASE Langevin dynamics, the temperature are normally # scalars, but in principle one quantity per atom could be specified by giving # an array. This is not implemented yet here. else: - self.tschedule = np.full(self.nsteps+1, self.temperature) + self.tschedule = np.full(self.nsteps + 1, self.temperature) if self.ensemble == "nvt": self.pressure = np.nan - self.pschedule = np.full(self.nsteps+1, self.pressure) + self.pschedule = np.full(self.nsteps + 1, self.pressure) return if isinstance(self.pressure, Sequence) or ( isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 1 - ): + ): self.pschedule = np.interp( - np.arange(self.nsteps+1), - np.arange(len(self.pressure)), - self.pressure + np.arange(self.nsteps + 1), np.arange(len(self.pressure)), self.pressure ) elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: self.pschedule = interp1d( - np.arange(self.nsteps+1), - self.pressure, - kind="linear" + np.arange(self.nsteps + 1), self.pressure, kind="linear" ) else: - self.pschedule = np.full(self.nsteps+1, self.pressure) + self.pschedule = np.full(self.nsteps + 1, self.pressure) def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" @@ -197,12 +194,9 @@ def _get_ensemble_defaults(self) -> None: self.ase_md_kwargs.pop("externalstress", None) elif self.ensemble == "npt": self.ase_md_kwargs["temperature_K"] = self.tschedule[0] - self.ase_md_kwargs["externalstress"] = ( - self.pschedule[0] * 1e3 * units.bar - ) - + self.ase_md_kwargs["externalstress"] = self.pschedule[0] * 1e3 * units.bar - # NOTE: We take of the temperature in _get_ensemble_schedule instead + # NOTE: We take in the temperature in _get_ensemble_schedule instead # if self.ensemble in ("nvt", "npt") and all( # self.ase_md_kwargs.get(key) is None # for key in ("temperature_K", "temperature") @@ -218,11 +212,12 @@ def _get_ensemble_defaults(self) -> None: if self.dynamics.lower() == "langevin": # NOTE: Unless we have a detailed documentation on the conversion of all - # parameters for all different ASE dyanmics, it is not a good idea to + # parameters for all different ASE dynamics, it is not a good idea to # convert the parameters here. It is better to let the user to consult # the ASE documentation and set the parameters themselves. self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get( - "friction", 10.0 * 1e-3 / units.fs # Same default as in VASP: 10 ps^-1 + "friction", + 10.0 * 1e-3 / units.fs, # Same default as in VASP: 10 ps^-1 ) # NOTE: same as above, user can specify per atom friction but we don't # expect to change their intention to pass into ASE MD function @@ -233,7 +228,7 @@ def _get_ensemble_defaults(self) -> None: # else: # self.ase_md_kwargs["friction"] *= 1.0e-3 / fs - @job(output_schema=ForceFieldTaskDocument) + @forcefield_job def make( self, structure: Structure, @@ -248,7 +243,7 @@ def make( pymatgen structure. prev_dir : str or Path or None A previous calculation directory to copy output files from. Unused, just - added to match the method signature of other makers. + added to match the method signature of other makers. """ self._get_ensemble_schedule() self._get_ensemble_defaults() @@ -275,7 +270,7 @@ def make( atoms = structure.to_ase_atoms() - u, q = schur(atoms.get_cell(complete=True)) + u, _ = schur(atoms.get_cell(complete=True)) atoms.set_cell(u, scale_atoms=True) if initial_velocities: @@ -284,15 +279,12 @@ def make( MaxwellBoltzmannDistribution( atoms=atoms, temperature_K=self.tschedule[0], - rng=None # TODO: we might want to use a seed for reproducibility - ) + rng=None, # TODO: we might want to use a seed for reproducibility + ) if self.zero_linear_momentum: Stationary(atoms) if self.zero_angular_momentum: ZeroRotation(atoms) - else: - # NOTE: ase has zero velocities by default already - pass self.calculator_args = self.calculator_args or [] self.calculator_kwargs = self.calculator_kwargs or {} @@ -304,11 +296,12 @@ def make( md_runner = md_func( atoms=atoms, timestep=self.timestep * units.fs, - # trajectory="trajectory.traj", # TODO: we might want implement taskdoc to save ase binary traj - **self.ase_md_kwargs + **self.ase_md_kwargs, ) - def callback(dyn: MolecularDynamics = md_runner) -> None: + md_runner.attach(md_observer, interval=self.traj_interval) + + def _callback(dyn: MolecularDynamics = md_runner) -> None: if self.ensemble == "nve": return dyn.set_temperature(temperature_K=self.tschedule[dyn.nsteps]) @@ -316,19 +309,21 @@ def callback(dyn: MolecularDynamics = md_runner) -> None: return dyn.set_stress(self.pschedule[dyn.nsteps] * 1e3 * units.bar) - md_runner.attach(md_observer, interval=self.traj_interval) - md_runner.attach(callback, interval=1) + md_runner.attach(_callback, interval=1) md_runner.run(steps=self.nsteps) if self.traj_file is not None: - md_observer.save(self.traj_file) + md_observer.save(filename=self.traj_file, fmt=self.traj_file_fmt) structure = AseAtomsAdaptor.get_structure(atoms) self.task_document_kwargs = self.task_document_kwargs or {} return ForceFieldTaskDocument.from_ase_compatible_result( self.force_field_name, - {"final_structure": structure, "trajectory": md_observer}, + { + "final_structure": structure, + "trajectory": md_observer.to_pymatgen_trajectory(None), + }, relax_cell=(self.ensemble == "npt"), steps=self.nsteps, relax_kwargs=None, @@ -346,7 +341,9 @@ class MACEMDMaker(ForceFieldMDMaker): name: str = "MACE MD" force_field_name: str = "MACE" - calculator_kwargs: dict = field(default_factory=lambda: {"default_dtype": "float32"}) + calculator_kwargs: dict = field( + default_factory=lambda: {"default_dtype": "float32"} + ) def _calculator(self) -> Calculator: from mace.calculators import mace_mp diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 184f61b5cc..bdfdf497ba 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -10,8 +10,6 @@ from emmet.core.vasp.calculation import StoreTrajectoryOption from pydantic import BaseModel, Field from pymatgen.core.structure import Structure -from pymatgen.core.trajectory import Trajectory -from pymatgen.io.ase import AseAtomsAdaptor from atomate2.forcefields import MLFF @@ -154,27 +152,17 @@ def from_ase_compatible_result( ionic_step_data : tuple Which data to save from each ionic step. """ - trajectory = { - key: value - for key, value in result["trajectory"].__dict__.items() - if not key.startswith("_") - } + trajectory = result["trajectory"] + n_steps = len(trajectory) # NOTE: convert stress units from eV/A³ to kBar (* -1 from standard output) # and to 3x3 matrix to comply with MP convention - for idx in range(len(trajectory["stresses"])): - trajectory["stresses"][idx] = voigt_6_to_full_3x3_stress( - trajectory["stresses"][idx] * -10 / GPa + for idx in range(n_steps): + trajectory.frame_properties[idx]["stress"] = voigt_6_to_full_3x3_stress( + [val * -10 / GPa for val in trajectory.frame_properties[idx]["stress"]] ) - species = AseAtomsAdaptor.get_structure(trajectory["atoms"]).species - - input_structure = Structure( - lattice=trajectory["cells"][0], - coords=trajectory["atom_positions"][0], - species=species, - coords_are_cartesian=True, - ) + input_structure = trajectory[0] input_doc = InputDoc( structure=input_structure, @@ -188,60 +176,42 @@ def from_ase_compatible_result( # number of steps for static calculations. if steps <= 1: steps = 1 - for key in trajectory: - trajectory[key] = [trajectory[key][0]] + trajectory = trajectory[0] + for key in trajectory.frame_properties: + trajectory.frame_properties[key] = [trajectory.frame_properties[key][0]] output_structure = input_structure else: output_structure = result["final_structure"] - final_energy = trajectory["energies"][-1] - final_energy_per_atom = trajectory["energies"][-1] / input_structure.num_sites - final_forces = trajectory["forces"][-1].tolist() - final_stress = trajectory["stresses"][-1].tolist() - - n_steps = len(trajectory["energies"]) + final_energy = trajectory.frame_properties[-1]["energy"] + final_energy_per_atom = final_energy / input_structure.num_sites + final_forces = trajectory.frame_properties[-1]["forces"] + final_stress = trajectory.frame_properties[-1]["stress"] ionic_steps = [] for idx in range(n_steps): - energy = ( - trajectory["energies"][idx] if "energy" in ionic_step_data else None - ) - forces = ( - trajectory["forces"][idx].tolist() - if "forces" in ionic_step_data - else None - ) - stress = ( - trajectory["stresses"][idx].tolist() - if "stress" in ionic_step_data + _ionic_step_data = { + key: trajectory.frame_properties[idx][key] + if key in ionic_step_data else None - ) + for key in ("energy", "forces", "stress") + } - if "structure" in ionic_step_data: - cur_structure = Structure( - lattice=trajectory["cells"][idx], - coords=trajectory["atom_positions"][idx], - species=species, - coords_are_cartesian=True, - ) - else: - cur_structure = None + cur_structure = trajectory[idx] if "structure" in ionic_step_data else None # include "magmoms" in :obj:`ionic_step` if the trajectory has "magmoms" - magmom_data = {} - if trajectory.get("magmoms"): - magmom_data = { - "magmoms": trajectory["magmoms"][idx].tolist() - if "magmoms" in ionic_step_data - else None - } + if trajectory.frame_properties[idx].get("magmoms"): + _ionic_step_data.update( + { + "magmoms": trajectory.frame_properties[idx]["magmoms"] + if "magmoms" in ionic_step_data + else None + } + ) ionic_step = IonicStep( - energy=energy, - forces=forces, - stress=stress, structure=cur_structure, - **magmom_data, + **_ionic_step_data, ) ionic_steps.append(ionic_step) @@ -251,24 +221,7 @@ def from_ase_compatible_result( # For VASP calculations, the PARTIAL trajectory option removes # electronic step info. There is no equivalent for forcefields, # so we just save the same info for FULL and PARTIAL options. - - traj_steps = [ - step.model_dump(exclude=("structure",)) for step in ionic_steps - ] - for idx in range(n_steps): - if trajectory.get("temperatures"): - traj_steps[idx]["temperature"] = trajectory.get("temperatures")[idx] - if trajectory.get("velocities"): - traj_steps[idx]["velocities"] = trajectory["velocities"][ - idx - ].tolist() - - traj = Trajectory.from_structures( - [step.structure for step in ionic_steps], - frame_properties=traj_steps, - constant_lattice=False, - ) - forcefield_objects[ForcefieldObject.TRAJECTORY] = traj # type: ignore[index] + forcefield_objects[ForcefieldObject.TRAJECTORY] = trajectory # type: ignore[index] output_doc = OutputDoc( structure=output_structure, diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index b662469bbb..a32e2b661b 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -16,10 +16,13 @@ from typing import TYPE_CHECKING import numpy as np +from ase.calculators.singlepoint import SinglePointCalculator +from ase.io import Trajectory as AseTrajectory from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG from monty.serialization import dumpfn from pymatgen.core.structure import Molecule, Structure +from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor try: @@ -38,7 +41,7 @@ if TYPE_CHECKING: from os import PathLike - from typing import Any + from typing import Any, Literal from ase import Atoms from ase.calculators.calculator import Calculator @@ -93,7 +96,8 @@ def __call__(self) -> None: # TODO: maybe include magnetic moments self.energies.append(self.compute_energy()) self.forces.append(self.atoms.get_forces()) - # TODO: MD needs kinetic energy parts of stress, we might just need to refactor TrajectoryObserver + # TODO: MD needs kinetic energy parts of stress, + # we might just need to refactor TrajectoryObserver # Now it is safe for 0K relaxation as the atoms don't have momenta self.stresses.append(self.atoms.get_stress(include_ideal_gas=True)) self.atom_positions.append(self.atoms.get_positions()) @@ -102,8 +106,9 @@ def __call__(self) -> None: if self._store_md_outputs: self.velocities.append(self.atoms.get_velocities()) self.temperatures.append(self.atoms.get_temperature()) - # self.stresses.append(self.atoms.get_stress(voigt=True, include_ideal_gas=True)) - + # self.stresses.append(self.atoms.get_stress( + # voigt=True, include_ideal_gas=True) + # ) def compute_energy(self) -> float: """ @@ -115,7 +120,9 @@ def compute_energy(self) -> float: """ return self.atoms.get_potential_energy() - def save(self, filename: str | PathLike) -> None: + def save( + self, filename: str | PathLike | None, fmt: Literal["pmg", "ase"] = "ase" + ) -> None: """ Save the trajectory file using monty.serialization. @@ -127,23 +134,103 @@ def save(self, filename: str | PathLike) -> None: ------- None """ - dumpfn(self.as_dict(), filename) + filename = str(filename) if filename is not None else None + if fmt == "pmg": + self.to_pymatgen_trajectory(filename=filename) + elif fmt == "ase": + self.to_ase_trajectory(filename=filename) + + def to_ase_trajectory(self, filename: str | None = "atoms.traj") -> AseTrajectory: + """ + Convert to an ASE .Trajectory. + + Parameters + ---------- + filename : str | None + Name of the file to write the ASE trajectory to. + If None, no file is written. + """ + for idx in range(len(self.cells)): + atoms = self.atoms.copy() + atoms.set_positions(self.atom_positions[idx]) + atoms.set_cell(self.cells[idx]) + atoms.set_velocities(self.velocities[idx]) + atoms.calc = SinglePointCalculator( + atoms=atoms, + energy=self.energies[idx], + forces=self.forces[idx], + stress=self.stresses[idx], + ) + with AseTrajectory(filename, "a" if idx > 0 else "w", atoms=atoms) as f: + f.write() + + return AseTrajectory(filename, "r") + + def to_pymatgen_trajectory( + self, filename: str | None = "trajectory.json.gz" + ) -> PmgTrajectory: + """ + Convert the trajectory to a pymatgen .Trajectory object. + + Parameters + ---------- + filename : str or None + Name of the file to write the pymatgen trajectory to. + If None, no file is written. + """ + n_md_steps = len(self.cells) + species = AseAtomsAdaptor.get_structure(self.atoms).species + + structures = [ + Structure( + lattice=self.cells[idx], + coords=self.atom_positions[idx], + species=species, + coords_are_cartesian=True, + ) + for idx in range(n_md_steps) + ] + + # sanitize trajectory entries + traj = self.as_dict() + + frame_property_keys = ["energy", "forces", "stress"] + if self._store_md_outputs: + frame_property_keys += ["velocities", "temperature"] + + frame_properties = [ + {key: traj[key][idx] for key in frame_property_keys} + for idx in range(n_md_steps) + ] + + traj = PmgTrajectory.from_structures( + structures, + frame_properties=frame_properties, + constant_lattice=False, + ) + + if filename: + dumpfn(traj, filename) + + return traj def as_dict(self) -> dict: """Make JSONable dict representation of the Trajectory.""" traj_dict = { "energy": self.energies, "forces": self.forces, - "stresses": self.stresses, + "stress": self.stresses, "atom_positions": self.atom_positions, "cell": self.cells, "atomic_number": self.atoms.get_atomic_numbers(), } if self._store_md_outputs: - traj_dict.update({ - "velocities": self.velocities, - "temperature": self.temperatures, - }) + traj_dict.update( + { + "velocities": self.velocities, + "temperature": self.temperatures, + } + ) # sanitize dict for key in traj_dict: if all(isinstance(val, np.ndarray) for val in traj_dict[key]): diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 446889148c..8f17c4fb1e 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -4,30 +4,36 @@ import pytest import torch from ase import units +from ase.io import Trajectory from jobflow import run_locally from monty.serialization import loadfn +from pymatgen.analysis.structure_matcher import StructureMatcher from atomate2.forcefields.md import CHGNetMDMaker, M3GNetMDMaker, MACEMDMaker _to_maker = {"CHGNet": CHGNetMDMaker, "M3GNet": M3GNetMDMaker, "MACE": MACEMDMaker} +# MACE changes the default dtype, ensure consistent dtype here +torch.set_default_dtype(torch.float32) + @pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE"]) def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): nsteps = 5 ref_energies_per_atom = { - "CHGNet": -5.30686092376709, - "M3GNet": -5.417105674743652, - "MACE": -5.33246374130249, + "CHGNet": -5.280157089233398, + "M3GNet": -5.387282371520996, + "MACE": -5.311369895935059, } structure = si_structure.to_conventional() * (2, 2, 2) + # ASE can slightly change tolerances on structure positions + matcher = StructureMatcher() - # MACE changes the default dtype, ensure consistent dtype here - torch.set_default_dtype(torch.float32) - - job = _to_maker[ff_name](nsteps=nsteps, traj_file="md_traj.json.gz").make(structure) + job = _to_maker[ff_name]( + nsteps=nsteps, traj_file="md_traj.json.gz", traj_file_fmt="pmg" + ).make(structure) response = run_locally(job, ensure_success=True) taskdoc = response[next(iter(response))][1].output @@ -37,8 +43,8 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): ) # Check that we have the right number of MD steps - # ASE logs the initial structure energy, and doesn"t count this as an MD step - assert taskdoc.output.ionic_steps[0].structure == structure + # ASE logs the initial structure energy, and doesn't count this as an MD step + assert matcher.fit(taskdoc.output.ionic_steps[0].structure, structure) assert len(taskdoc.output.ionic_steps) == nsteps + 1 # Check that the ionic steps have the expected physical properties @@ -57,24 +63,66 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): for step in taskdoc.forcefield_objects["trajectory"].frame_properties ) + +@pytest.mark.parametrize("traj_file", ["trajectory.json.gz", "atoms.traj"]) +def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): + nsteps = 5 + # Check that traj file written to disk is consistent with trajectory # stored to the task document - traj_from_file = loadfn("md_traj.json.gz") - assert len(traj_from_file["energy"]) == nsteps + 1 - _traj_key_to_object_key = { - "stresses": "stress", - } - assert all( - np.all( - traj_from_file[key][idx] - == taskdoc.forcefield_objects["trajectory"].frame_properties[idx][ - _traj_key_to_object_key.get(key, key) - ] + if ".json.gz" in traj_file: + traj_file_fmt = "pmg" + traj_file_loader = loadfn + else: + traj_file_fmt = "ase" + traj_file_loader = Trajectory + + structure = si_structure.to_conventional() * (2, 2, 2) + job = _to_maker[ff_name]( + nsteps=nsteps, + traj_file=traj_file, + traj_file_fmt=traj_file_fmt, + ).make(structure) + response = run_locally(job, ensure_success=True) + taskdoc = response[next(iter(response))][1].output + + traj_from_file = traj_file_loader(traj_file) + + assert len(traj_from_file) == nsteps + 1 + + if traj_file_fmt == "pmg": + assert all( + np.all( + traj_from_file.frame_properties[idx][key] + == taskdoc.forcefield_objects["trajectory"] + .frame_properties[idx] + .get(key) + ) + for key in ("energy", "temperature", "forces", "velocities") + for idx in range(nsteps + 1) ) - for key in ("energy", "temperature", "forces", "velocities") - for idx in range(nsteps + 1) - ) + elif traj_file_fmt == "ase": + traj_as_dict = [ + { + "energy": traj_from_file[idx].get_potential_energy(), + "temperature": traj_from_file[idx].get_temperature(), + "forces": traj_from_file[idx].get_forces(), + "velocities": traj_from_file[idx].get_velocities(), + } + for idx in range(1, nsteps + 1) + ] + assert all( + np.all( + traj_as_dict[idx - 1][key] + == taskdoc.forcefield_objects["trajectory"] + .frame_properties[idx] + .get(key) + ) + for key in ("energy", "temperature", "forces", "velocities") + for idx in range(1, nsteps + 1) + ) + @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_temp_schedule(ff_name, si_structure, clean_dir): @@ -83,21 +131,19 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): structure = si_structure.to_conventional() * (2, 2, 2) - # MACE changes the default dtype, ensure consistent dtype here - torch.set_default_dtype(torch.float32) - job = _to_maker[ff_name]( - nsteps=nsteps, traj_file="md_traj.json.gz", + nsteps=nsteps, + traj_file=None, dynamics="nose-hoover", temperature=temp_schedule, - ase_md_kwargs=dict(ttime=50.0 * units.fs, pfactor=None) + ase_md_kwargs=dict(ttime=50.0 * units.fs, pfactor=None), ).make(structure) response = run_locally(job, ensure_success=True) taskdoc = response[next(iter(response))][1].output - temp_history = [ - step["temperature"] for step in taskdoc.forcefield_objects["trajectory"].frame_properties + step["temperature"] + for step in taskdoc.forcefield_objects["trajectory"].frame_properties ] assert temp_history[-1] > temp_schedule[0] @@ -106,22 +152,21 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_press_schedule(ff_name, si_structure, clean_dir): nsteps = 100 - press_schedule = [0, 10] # kbar + press_schedule = [0, 10] # kbar structure = si_structure.to_conventional() * (3, 3, 3) - # MACE changes the default dtype, ensure consistent dtype here - torch.set_default_dtype(torch.float32) - job = _to_maker[ff_name]( ensemble="npt", - nsteps=nsteps, traj_file="md_traj.json.gz", + nsteps=nsteps, + traj_file="md_traj.json.gz", + traj_file_fmt="pmg", dynamics="nose-hoover", pressure=press_schedule, ase_md_kwargs=dict( ttime=50.0 * units.fs, - pfactor=(75.0 * units.fs)**2 * units.GPa, - ) + pfactor=(75.0 * units.fs) ** 2 * units.GPa, + ), ).make(structure) run_locally(job, ensure_success=True) # taskdoc = response[next(iter(response))][1].output @@ -129,7 +174,8 @@ def test_press_schedule(ff_name, si_structure, clean_dir): traj_from_file = loadfn("md_traj.json.gz") stress_history = [ - (stress[0] + stress[1] + stress[2])/3.0 for stress in traj_from_file["stresses"] + sum(traj_from_file.frame_properties[idx]["stress"][:3]) / 3.0 + for idx in range(len(traj_from_file)) ] assert stress_history[-1] < stress_history[0] From 9ff7dae91d0569c7c5e23385a15256fe51584edb Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 1 Mar 2024 12:16:07 -0800 Subject: [PATCH 20/48] Change decomposition to upper triangular cell only when Nose-Hoover requested --- src/atomate2/forcefields/md.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 13ca4864bb..4a631b1c1c 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -270,8 +270,11 @@ def make( atoms = structure.to_ase_atoms() - u, _ = schur(atoms.get_cell(complete=True)) - atoms.set_cell(u, scale_atoms=True) + if md_func is NPT: + # Note that until md_func is instantiated, isinstance(md_func,NPT) is False + # ASE NPT implementation requires upper triangular cell + u, _ = schur(atoms.get_cell(complete=True)) + atoms.set_cell(u, scale_atoms=True) if initial_velocities: atoms.set_velocities(initial_velocities) From 9ea1ef95e840c114a91e25bcf320393e666f9fbd Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 1 Mar 2024 15:11:01 -0800 Subject: [PATCH 21/48] Fix failing tests --- src/atomate2/common/jobs/phonons.py | 10 ++-- src/atomate2/forcefields/schemas.py | 32 +++++++++---- src/atomate2/forcefields/utils.py | 74 ++++++++++++++++++----------- 3 files changed, 77 insertions(+), 39 deletions(-) diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index d273f1546a..e70b080d56 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -153,11 +153,13 @@ def generate_phonon_displacements( "of the symmetry of the structure and thus will be removed now.", stacklevel=1, ) - cell = get_phonopy_structure( + + # Note that `remove_site_property` returns None + if "magmom" in structure.site_properties: structure.remove_site_property(property_name="magmom") - if "magmom" in structure.site_properties - else structure - ) + + cell = get_phonopy_structure(structure) + factor = get_factor(code) # a bit of code repetition here as I currently diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index bdfdf497ba..aba5c243a1 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -10,8 +10,10 @@ from emmet.core.vasp.calculation import StoreTrajectoryOption from pydantic import BaseModel, Field from pymatgen.core.structure import Structure +from pymatgen.core.trajectory import Trajectory from atomate2.forcefields import MLFF +from atomate2.forcefields.utils import _get_pymatgen_trajectory_from_observer class ForcefieldObject(ValueEnum): @@ -152,15 +154,26 @@ def from_ase_compatible_result( ionic_step_data : tuple Which data to save from each ionic step. """ - trajectory = result["trajectory"] + if isinstance(result["trajectory"], Trajectory): + trajectory = result["trajectory"] + else: + trajectory = _get_pymatgen_trajectory_from_observer( + result["trajectory"], + frame_property_keys=["energies", "forces", "stresses", "magmoms"], + ) + n_steps = len(trajectory) # NOTE: convert stress units from eV/A³ to kBar (* -1 from standard output) # and to 3x3 matrix to comply with MP convention for idx in range(n_steps): - trajectory.frame_properties[idx]["stress"] = voigt_6_to_full_3x3_stress( - [val * -10 / GPa for val in trajectory.frame_properties[idx]["stress"]] - ) + if trajectory.frame_properties[idx].get("stress") is not None: + trajectory.frame_properties[idx]["stress"] = voigt_6_to_full_3x3_stress( + [ + val * -10 / GPa + for val in trajectory.frame_properties[idx]["stress"] + ] + ) input_structure = trajectory[0] @@ -176,9 +189,12 @@ def from_ase_compatible_result( # number of steps for static calculations. if steps <= 1: steps = 1 - trajectory = trajectory[0] - for key in trajectory.frame_properties: - trajectory.frame_properties[key] = [trajectory.frame_properties[key][0]] + n_steps = 1 + trajectory = Trajectory.from_structures( + structures=[trajectory[0]], + frame_properties=[trajectory.frame_properties[0]], + constant_lattice=False, + ) output_structure = input_structure else: output_structure = result["final_structure"] @@ -200,7 +216,7 @@ def from_ase_compatible_result( cur_structure = trajectory[idx] if "structure" in ionic_step_data else None # include "magmoms" in :obj:`ionic_step` if the trajectory has "magmoms" - if trajectory.frame_properties[idx].get("magmoms"): + if "magmoms" in trajectory.frame_properties[idx]: _ionic_step_data.update( { "magmoms": trajectory.frame_properties[idx]["magmoms"] diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 50eccaff68..f392d7ea3b 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -61,6 +61,45 @@ } +def _get_pymatgen_trajectory_from_observer( + trajectory_observer: Any, frame_property_keys: list[str] +) -> PmgTrajectory: + to_singluar = {"energies": "energy", "stresses": "stress"} + + if hasattr(trajectory_observer, "as_dict"): + traj = trajectory_observer.as_dict() + else: + traj = trajectory_observer.__dict__ + + n_md_steps = len(traj["cells"]) + species = AseAtomsAdaptor.get_structure(traj["atoms"]).species + + structures = [ + Structure( + lattice=traj["cells"][idx], + coords=traj["atom_positions"][idx], + species=species, + coords_are_cartesian=True, + ) + for idx in range(n_md_steps) + ] + + frame_properties = [ + { + to_singluar.get(key, key): traj[key][idx] + for key in frame_property_keys + if key in traj + } + for idx in range(n_md_steps) + ] + + return PmgTrajectory.from_structures( + structures, + frame_properties=frame_properties, + constant_lattice=False, + ) + + class TrajectoryObserver: """Trajectory observer. @@ -178,35 +217,12 @@ def to_pymatgen_trajectory( Name of the file to write the pymatgen trajectory to. If None, no file is written. """ - n_md_steps = len(self.cells) - species = AseAtomsAdaptor.get_structure(self.atoms).species - - structures = [ - Structure( - lattice=self.cells[idx], - coords=self.atom_positions[idx], - species=species, - coords_are_cartesian=True, - ) - for idx in range(n_md_steps) - ] - - # sanitize trajectory entries - traj = self.as_dict() - frame_property_keys = ["energy", "forces", "stress"] if self._store_md_outputs: frame_property_keys += ["velocities", "temperature"] - frame_properties = [ - {key: traj[key][idx] for key in frame_property_keys} - for idx in range(n_md_steps) - ] - - traj = PmgTrajectory.from_structures( - structures, - frame_properties=frame_properties, - constant_lattice=False, + traj = _get_pymatgen_trajectory_from_observer( + self, frame_property_keys=frame_property_keys ) if filename: @@ -221,7 +237,8 @@ def as_dict(self) -> dict: "forces": self.forces, "stress": self.stresses, "atom_positions": self.atom_positions, - "cell": self.cells, + "cells": self.cells, + "atoms": self.atoms, "atomic_number": self.atoms.get_atomic_numbers(), } if self._store_md_outputs: @@ -324,4 +341,7 @@ def relax( atoms = atoms.atoms struct = self.ase_adaptor.get_structure(atoms) - return {"final_structure": struct, "trajectory": obs} + return { + "final_structure": struct, + "trajectory": obs.to_pymatgen_trajectory(None), + } From 8727300cce3c9f163995c11e95dcdc15e9f1ef26 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 1 Mar 2024 15:21:09 -0800 Subject: [PATCH 22/48] Change MD default TaskDoc ionic step to only store mandatory info --- src/atomate2/forcefields/md.py | 5 ++++- tests/forcefields/test_md.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 4a631b1c1c..ef29410afb 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -137,7 +137,10 @@ class ForceFieldMDMaker(Maker): zero_linear_momentum: bool = False zero_angular_momentum: bool = False task_document_kwargs: dict = field( - default_factory=lambda: {"store_trajectory": "partial"} + default_factory=lambda: { + "store_trajectory": "partial", + "ionic_step_data": ("energy",), # energy is required in ionicsteps + } ) def _get_ensemble_schedule(self) -> None: diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 8f17c4fb1e..bcc77a87b8 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -32,7 +32,10 @@ def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): matcher = StructureMatcher() job = _to_maker[ff_name]( - nsteps=nsteps, traj_file="md_traj.json.gz", traj_file_fmt="pmg" + nsteps=nsteps, + traj_file="md_traj.json.gz", + traj_file_fmt="pmg", + task_document_kwargs={"store_trajectory": "partial"}, ).make(structure) response = run_locally(job, ensure_success=True) taskdoc = response[next(iter(response))][1].output From 56edc31dda4b3e19cdb21f3ae87c1c88c8eaf669 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 1 Mar 2024 15:38:32 -0800 Subject: [PATCH 23/48] update pymatgen requirement because of CifParser deprecation --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a7077edf53..6882f28625 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "numpy", "pydantic-settings>=2.0.3", "pydantic>=2.0.1", - "pymatgen>=2023.12.18", + "pymatgen>=2024.3.1", ] [project.optional-dependencies] @@ -88,7 +88,7 @@ strict = [ "pydantic-settings==2.2.1", "pydantic==2.6.2", "pymatgen-analysis-defects==2024.2.24", - "pymatgen==2024.2.23", + "pymatgen==2024.3.1", "quippy-ase==0.9.14", "seekpath==2.1.0", "typing-extensions==4.10.0", From 5bbe9d7e906bd3c47d511630a38ee2bebc6090b7 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Wed, 6 Mar 2024 08:45:00 -0800 Subject: [PATCH 24/48] Significant refactor to all forcefield jobs. Use commmon ASE calc structure, allow loading via MontyDecoder. Add check in relaxation for force convergence, attr in Forcefield taskdoc --- src/atomate2/forcefields/__init__.py | 1 + src/atomate2/forcefields/jobs.py | 147 +++++------------------- src/atomate2/forcefields/md.py | 54 ++++----- src/atomate2/forcefields/schemas.py | 31 +++-- src/atomate2/forcefields/utils.py | 129 +++++++++++++++++++-- tests/forcefields/flows/test_elastic.py | 3 +- tests/forcefields/test_jobs.py | 41 ++++--- tests/forcefields/test_utils.py | 35 +++++- 8 files changed, 251 insertions(+), 190 deletions(-) diff --git a/src/atomate2/forcefields/__init__.py b/src/atomate2/forcefields/__init__.py index 1a192060cc..0f4110eaf9 100644 --- a/src/atomate2/forcefields/__init__.py +++ b/src/atomate2/forcefields/__init__.py @@ -9,3 +9,4 @@ class MLFF(Enum): # TODO inherit from StrEnum when 3.11+ GAP = "GAP" M3GNet = "M3GNet" CHGNet = "CHGNet" + Forcefield = "Forcefield" # default placeholder option diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 294db171ae..64a48bc685 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -11,13 +11,13 @@ from atomate2.forcefields import MLFF from atomate2.forcefields.schemas import ForceFieldTaskDocument -from atomate2.forcefields.utils import Relaxer +from atomate2.forcefields.utils import Relaxer, ase_calculator if TYPE_CHECKING: - from collections.abc import Sequence from pathlib import Path from typing import Callable + from ase.calculators.calculator import Calculator from pymatgen.core.structure import Structure logger = logging.getLogger(__name__) @@ -66,7 +66,10 @@ class ForceFieldRelaxMaker(Maker): """ Base Maker to calculate forces and stresses using any force field. - Should be subclassed to use a specific force field. + Should be subclassed to use a specific force field. By default, + the code attempts to use the `self.force_field_name` attr to look + up a predefined forcefield. To overwrite this behavior, + redefine `self._calculator`. Parameters ---------- @@ -82,16 +85,19 @@ class ForceFieldRelaxMaker(Maker): Keyword arguments that will get passed to :obj:`Relaxer.relax`. optimizer_kwargs : dict Keyword arguments that will get passed to :obj:`Relaxer()`. + calculator_kwargs : dict + Keyword arguments that will get passed to the ASE calculator. task_document_kwargs : dict Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`. """ name: str = "Force field relax" - force_field_name: str = "Force field" + force_field_name: str = f"{MLFF.Forcefield}" relax_cell: bool = True steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) @forcefield_job @@ -115,7 +121,10 @@ def make( "Behavior may vary..." ) - result = self._relax(structure) + relaxer = Relaxer( + self._calculator(), relax_cell=self.relax_cell, **self.optimizer_kwargs + ) + result = relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) return ForceFieldTaskDocument.from_ase_compatible_result( self.force_field_name, @@ -127,8 +136,9 @@ def make( **self.task_document_kwargs, ) - def _relax(self, structure: Structure) -> dict: - raise NotImplementedError + def _calculator(self) -> Calculator: + """ASE calculator, can be overwritten by user.""" + return ase_calculator(self.force_field_name, self.calculator_kwargs) @dataclass @@ -148,6 +158,11 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker): name: str = "Force field static" force_field_name: str = "Force field" + relax_cell: bool = False + steps: int = 1 + relax_kwargs: dict = field(default_factory=dict) + optimizer_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) @forcefield_job @@ -155,7 +170,7 @@ def make( self, structure: Structure, prev_dir: str | Path | None = None ) -> ForceFieldTaskDocument: """ - Perform a static evaluation using a force field. + Perform a relaxation of a structure using a force field. Parameters ---------- @@ -171,11 +186,12 @@ def make( "Behavior may vary..." ) - result = self._evaluate_static(structure) + static_calc = Relaxer(self._calculator(), relax_cell=False) + result = static_calc.relax(structure, steps=1) return ForceFieldTaskDocument.from_ase_compatible_result( - self.force_field_name, - result, + forcefield_name=self.force_field_name, + result=result, relax_cell=False, steps=1, relax_kwargs=None, @@ -183,9 +199,6 @@ def make( **self.task_document_kwargs, ) - def _evaluate_static(self, structure: Structure) -> dict: - raise NotImplementedError - @dataclass class CHGNetRelaxMaker(ForceFieldRelaxMaker): @@ -209,21 +222,13 @@ class CHGNetRelaxMaker(ForceFieldRelaxMaker): """ name: str = f"{MLFF.CHGNet} relax" - force_field_name = f"{MLFF.CHGNet}" + force_field_name: str = f"{MLFF.CHGNet}" relax_cell: bool = True steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - def _relax(self, structure: Structure) -> dict: - from chgnet.model import StructOptimizer - - relaxer = StructOptimizer(**self.optimizer_kwargs) - return relaxer.relax( - structure, relax_cell=self.relax_cell, steps=self.steps, **self.relax_kwargs - ) - @dataclass class CHGNetStaticMaker(ForceFieldStaticMaker): @@ -239,15 +244,9 @@ class CHGNetStaticMaker(ForceFieldStaticMaker): """ name: str = f"{MLFF.CHGNet} static" - force_field_name = f"{MLFF.CHGNet}" + force_field_name: str = f"{MLFF.CHGNet}" task_document_kwargs: dict = field(default_factory=dict) - def _evaluate_static(self, structure: Structure) -> dict: - from chgnet.model import StructOptimizer - - relaxer = StructOptimizer() - return relaxer.relax(structure, steps=1) - @dataclass class M3GNetRelaxMaker(ForceFieldRelaxMaker): @@ -280,20 +279,6 @@ class M3GNetRelaxMaker(ForceFieldRelaxMaker): optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - def _relax(self, structure: Structure) -> dict: - import matgl - from matgl.ext.ase import Relaxer - - # Note: the below code was taken from the matgl repo examples. - # Load pre-trained M3GNet model (currently uses the MP-2021.2.8 database) - potential = matgl.load_model("M3GNet-MP-2021.2.8-PES") - - relaxer = Relaxer( - potential=potential, relax_cell=self.relax_cell, **self.optimizer_kwargs - ) - - return relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) - @dataclass class M3GNetStaticMaker(ForceFieldStaticMaker): @@ -314,18 +299,6 @@ class M3GNetStaticMaker(ForceFieldStaticMaker): force_field_name: str = f"{MLFF.M3GNet}" task_document_kwargs: dict = field(default_factory=dict) - def _evaluate_static(self, structure: Structure) -> dict: - import matgl - from matgl.ext.ase import Relaxer - - # Note: the below code was taken from the matgl repo examples. - # Load pre-trained M3GNet model (currently uses the MP-2021.2.8 database) - potential = matgl.load_model("M3GNet-MP-2021.2.8-PES") - - relaxer = Relaxer(potential=potential, relax_cell=False) - - return relaxer.relax(structure, steps=1) - @dataclass class MACERelaxMaker(ForceFieldRelaxMaker): @@ -365,17 +338,6 @@ class MACERelaxMaker(ForceFieldRelaxMaker): relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - model: str | Path | Sequence[str | Path] | None = None - model_kwargs: dict = field(default_factory=dict) - - def _relax(self, structure: Structure) -> dict: - from mace.calculators import mace_mp - - calculator = mace_mp(model=self.model, **self.model_kwargs) - relaxer = Relaxer( - calculator, relax_cell=self.relax_cell, **self.optimizer_kwargs - ) - return relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) @dataclass @@ -404,15 +366,6 @@ class MACEStaticMaker(ForceFieldStaticMaker): name: str = f"{MLFF.MACE} static" force_field_name: str = f"{MLFF.MACE}" task_document_kwargs: dict = field(default_factory=dict) - model: str | Path | Sequence[str | Path] | None = None - model_kwargs: dict = field(default_factory=dict) - - def _evaluate_static(self, structure: Structure) -> dict: - from mace.calculators import mace_mp - - calculator = mace_mp(model=self.model, **self.model_kwargs) - relaxer = Relaxer(calculator, relax_cell=False) - return relaxer.relax(structure, steps=1) @dataclass @@ -436,12 +389,6 @@ class GAPRelaxMaker(ForceFieldRelaxMaker): Keyword arguments that will get passed to :obj:`Relaxer()`. task_document_kwargs : dict Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`. - potential_args_str: str - args_str for :obj:`quippy.potential.Potential()'`. - potential_param_file_name: str | Path - param_file_name for :obj:`quippy.potential.Potential()'`. - potential_kwargs: dict - Further keywords for :obj:`quippy.potential.Potential()'`. """ name: str = f"{MLFF.GAP} relax" @@ -451,22 +398,6 @@ class GAPRelaxMaker(ForceFieldRelaxMaker): relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - potential_args_str: str | Path = "IP GAP" - potential_param_file_name: str = "gap.xml" - potential_kwargs: dict = field(default_factory=dict) - - def _relax(self, structure: Structure) -> dict: - from quippy.potential import Potential - - calculator = Potential( - args_str=self.potential_args_str, - param_filename=str(self.potential_param_file_name), - **self.potential_kwargs, - ) - relaxer = Relaxer( - calculator, **self.optimizer_kwargs, relax_cell=self.relax_cell - ) - return relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) @dataclass @@ -482,28 +413,8 @@ class GAPStaticMaker(ForceFieldStaticMaker): The name of the force field. task_document_kwargs : dict Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`. - potential_args_str: str - args_str for :obj:`quippy.potential.Potential()'`. - potential_param_file_name: str | Path - param_file_name for :obj:`quippy.potential.Potential()'`. - potential_kwargs: dict - Further keywords for :obj:`quippy.potential.Potential()'`. """ name: str = f"{MLFF.GAP} static" force_field_name: str = f"{MLFF.GAP}" task_document_kwargs: dict = field(default_factory=dict) - potential_args_str: str = "IP GAP" - potential_param_file_name: str | Path = "gap.xml" - potential_kwargs: dict = field(default_factory=dict) - - def _evaluate_static(self, structure: Structure) -> dict: - from quippy.potential import Potential - - calculator = Potential( - args_str=self.potential_args_str, - param_filename=str(self.potential_param_file_name), - **self.potential_kwargs, - ) - relaxer = Relaxer(calculator, relax_cell=False) - return relaxer.relax(structure, steps=1) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index ef29410afb..d42a80363f 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -26,9 +26,10 @@ from scipy.interpolate import interp1d from scipy.linalg import schur +from atomate2.forcefields import MLFF from atomate2.forcefields.jobs import forcefield_job -from atomate2.forcefields.schemas import ForceFieldTaskDocument -from atomate2.forcefields.utils import TrajectoryObserver +from atomate2.forcefields.schemas import ForcefieldResult, ForceFieldTaskDocument +from atomate2.forcefields.utils import TrajectoryObserver, ase_calculator if TYPE_CHECKING: from pathlib import Path @@ -100,8 +101,6 @@ class ForceFieldMDMaker(Maker): specifying thermostat as a string. ase_md_kwargs : dict | None = None Options except for temperature and pressure to pass into the ASE MD function - calculator_args : Sequence | None = None - args to pass to the ASE calculator class calculator_kwargs : dict | None = None kwargs to pass to the ASE calculator class traj_file : str | Path | None = None @@ -120,7 +119,7 @@ class ForceFieldMDMaker(Maker): """ name: str = "Forcefield MD" - force_field_name: str = "Forcefield" + force_field_name: str = f"{MLFF.Forcefield}" timestep: float | None = 2.0 nsteps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" @@ -129,7 +128,6 @@ class ForceFieldMDMaker(Maker): # end_temp: float | None = 300.0 pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None - calculator_args: list | tuple | None = None calculator_kwargs: dict | None = None traj_file: str | Path | None = None traj_file_fmt: Literal["pmg", "ase"] = "ase" @@ -292,7 +290,6 @@ def make( if self.zero_angular_momentum: ZeroRotation(atoms) - self.calculator_args = self.calculator_args or [] self.calculator_kwargs = self.calculator_kwargs or {} atoms.calc = self._calculator() @@ -324,12 +321,13 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: structure = AseAtomsAdaptor.get_structure(atoms) self.task_document_kwargs = self.task_document_kwargs or {} + return ForceFieldTaskDocument.from_ase_compatible_result( self.force_field_name, - { - "final_structure": structure, - "trajectory": md_observer.to_pymatgen_trajectory(None), - }, + ForcefieldResult( + final_structure=structure, + trajectory=md_observer.to_pymatgen_trajectory(None), + ), relax_cell=(self.ensemble == "npt"), steps=self.nsteps, relax_kwargs=None, @@ -338,46 +336,40 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: ) def _calculator(self) -> Calculator: - """To be implemented by the user.""" - return NotImplementedError + """ASE calculator, can be overwritten by user.""" + return ase_calculator(self.force_field_name, self.calculator_kwargs) +@dataclass class MACEMDMaker(ForceFieldMDMaker): """Perform an MD run with MACE.""" name: str = "MACE MD" - force_field_name: str = "MACE" + force_field_name: str = f"{MLFF.MACE}" calculator_kwargs: dict = field( default_factory=lambda: {"default_dtype": "float32"} ) - def _calculator(self) -> Calculator: - from mace.calculators import mace_mp - - return mace_mp(*self.calculator_args, **self.calculator_kwargs) - +@dataclass class M3GNetMDMaker(ForceFieldMDMaker): """Perform an MD run with M3GNet.""" name: str = "M3GNet MD" - force_field_name: str = "M3GNet" - - def _calculator(self) -> Calculator: - import matgl - from matgl.ext.ase import PESCalculator - - potential = matgl.load_model("M3GNet-MP-2021.2.8-PES") - return PESCalculator(potential, **self.calculator_kwargs) + force_field_name: str = f"{MLFF.M3GNet}" +@dataclass class CHGNetMDMaker(ForceFieldMDMaker): """Perform an MD run with CHGNet.""" name: str = "CHGNet MD" - force_field_name: str = "CHGNet" + force_field_name: str = f"{MLFF.CHGNet}" - def _calculator(self) -> Calculator: - from chgnet.model.dynamics import CHGNetCalculator - return CHGNetCalculator(*self.calculator_args, **self.calculator_kwargs) +@dataclass +class GAPMDMaker(ForceFieldMDMaker): + """Perform an MD run with GAP.""" + + name: str = "CHGNet MD" + force_field_name: str = f"{MLFF.GAP}" diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index aba5c243a1..06004ec754 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -13,7 +13,19 @@ from pymatgen.core.trajectory import Trajectory from atomate2.forcefields import MLFF -from atomate2.forcefields.utils import _get_pymatgen_trajectory_from_observer + + +class ForcefieldResult(dict): + """Schema to store outputs in ForceFieldTaskDocument.""" + + def __init__(self, **kwargs: Any) -> None: + kwargs = { + "final_structure": None, + "trajectory": None, + "is_force_converged": None, + **kwargs, + } + super().__init__(**kwargs) class ForcefieldObject(ValueEnum): @@ -122,6 +134,14 @@ class ForceFieldTaskDocument(StructureMetadata): None, description="Forcefield objects associated with this task" ) + is_force_converged: Optional[bool] = Field( + None, + description=( + "Whether the calculation is converged with respect " + "to interatomic forces." + ), + ) + @classmethod def from_ase_compatible_result( cls, @@ -154,13 +174,7 @@ def from_ase_compatible_result( ionic_step_data : tuple Which data to save from each ionic step. """ - if isinstance(result["trajectory"], Trajectory): - trajectory = result["trajectory"] - else: - trajectory = _get_pymatgen_trajectory_from_observer( - result["trajectory"], - frame_property_keys=["energies", "forces", "stresses", "magmoms"], - ) + trajectory = result["trajectory"] n_steps = len(trajectory) @@ -270,4 +284,5 @@ def from_ase_compatible_result( forcefield_version=version, included_objects=list(forcefield_objects.keys()), forcefield_objects=forcefield_objects, + is_force_converged=result.get("is_force_converged"), ) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index f392d7ea3b..0857f2bd91 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -11,20 +11,26 @@ import contextlib import io +import json import sys import warnings from typing import TYPE_CHECKING import numpy as np +from ase.calculators.calculator import PropertyNotImplementedError from ase.calculators.singlepoint import SinglePointCalculator from ase.io import Trajectory as AseTrajectory from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG +from monty.json import MontyDecoder from monty.serialization import dumpfn from pymatgen.core.structure import Molecule, Structure from pymatgen.core.trajectory import Trajectory as PmgTrajectory from pymatgen.io.ase import AseAtomsAdaptor +from atomate2.forcefields import MLFF +from atomate2.forcefields.schemas import ForcefieldResult + try: from ase.filters import FrechetCellFilter except ImportError: @@ -60,6 +66,13 @@ "BFGSLineSearch": BFGSLineSearch, } +_ase_calculator_defaults = { + "gap": { + "args_str": "IP GAP", + "param_filename": "gap.xml", + } +} + def _get_pymatgen_trajectory_from_observer( trajectory_observer: Any, frame_property_keys: list[str] @@ -122,6 +135,10 @@ def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: self.energies: list[float] = [] self.forces: list[np.ndarray] = [] self.stresses: list[np.ndarray] = [] + + self._store_magmoms = True + self.magmoms: list[np.ndarray] = [] + self.atom_positions: list[np.ndarray] = [] self.cells: list[np.ndarray] = [] @@ -139,6 +156,13 @@ def __call__(self) -> None: # we might just need to refactor TrajectoryObserver # Now it is safe for 0K relaxation as the atoms don't have momenta self.stresses.append(self.atoms.get_stress(include_ideal_gas=True)) + + if self._store_magmoms: + try: + self.magmoms.append(self.atoms.get_magnetic_moments()) + except PropertyNotImplementedError: + self._store_magmoms = False + self.atom_positions.append(self.atoms.get_positions()) self.cells.append(self.atoms.get_cell()[:]) @@ -193,13 +217,19 @@ def to_ase_trajectory(self, filename: str | None = "atoms.traj") -> AseTrajector atoms = self.atoms.copy() atoms.set_positions(self.atom_positions[idx]) atoms.set_cell(self.cells[idx]) - atoms.set_velocities(self.velocities[idx]) - atoms.calc = SinglePointCalculator( - atoms=atoms, - energy=self.energies[idx], - forces=self.forces[idx], - stress=self.stresses[idx], - ) + + if self._store_md_outputs: + atoms.set_velocities(self.velocities[idx]) + + kwargs = { + "energy": self.energies[idx], + "forces": self.forces[idx], + "stress": self.stresses[idx], + } + if self._store_magmoms: + kwargs["magmom"] = self.magmoms[idx] + + atoms.calc = SinglePointCalculator(atoms=atoms, **kwargs) with AseTrajectory(filename, "a" if idx > 0 else "w", atoms=atoms) as f: f.write() @@ -218,6 +248,8 @@ def to_pymatgen_trajectory( If None, no file is written. """ frame_property_keys = ["energy", "forces", "stress"] + if self._store_magmoms: + frame_property_keys += ["magmoms"] if self._store_md_outputs: frame_property_keys += ["velocities", "temperature"] @@ -241,6 +273,10 @@ def as_dict(self) -> dict: "atoms": self.atoms, "atomic_number": self.atoms.get_atomic_numbers(), } + + if self._store_magmoms: + traj_dict["magmoms"] = self.magmoms + if self._store_md_outputs: traj_dict.update( { @@ -298,7 +334,7 @@ def relax( verbose: bool = False, cell_filter: Filter = FrechetCellFilter, **kwargs, - ) -> dict[str, Any]: + ) -> ForcefieldResult: """ Relax the structure. @@ -341,7 +377,76 @@ def relax( atoms = atoms.atoms struct = self.ase_adaptor.get_structure(atoms) - return { - "final_structure": struct, - "trajectory": obs.to_pymatgen_trajectory(None), - } + traj = obs.to_pymatgen_trajectory(None) + is_force_conv = all( + np.linalg.norm(traj.frame_properties[-1]["forces"][i]) < abs(fmax) + for i in range(struct.num_sites) + ) + return ForcefieldResult( + final_structure=struct, trajectory=traj, is_force_converged=is_force_conv + ) + + +def ase_calculator( + calculator_meta: str | dict, calculator_kwargs: dict | None = None +) -> Calculator | None: + """ + Create an ASE calculator from a given set of metadata. + + Parameters + ---------- + calculator_meta : str or dict + If a str, should be one of `atomate2.forcefields.MLFF`. + If a dict, should be decodable by `monty.json.MontyDecoder`. + For example, one can also call the CHGNet calculator as follows + ``` + calculator_meta = { + "@module": "chgnet.model.dynamics", + "@callable": "CHGNetCalculator" + } + ``` + calculator_kwargs : dict or None (default) + kwargs to pass to the calculator object + + Returns + ------- + ASE .Calculator + """ + calculator = None + + if isinstance(calculator_meta, str) and calculator_meta in [ + f"{name}" for name in MLFF + ]: + calculator_name = MLFF(calculator_meta.split("MLFF.")[-1]) + calculator_kwargs = calculator_kwargs or _ase_calculator_defaults.get( + calculator_meta, {} + ) + + if calculator_name == MLFF.CHGNet: + from chgnet.model.dynamics import CHGNetCalculator + + calculator = CHGNetCalculator(**calculator_kwargs) + + elif calculator_name == MLFF.M3GNet: + import matgl + from matgl.ext.ase import PESCalculator + + potential = matgl.load_model("M3GNet-MP-2021.2.8-PES") + calculator = PESCalculator(potential, **calculator_kwargs) + + elif calculator_name == MLFF.MACE: + from mace.calculators import mace_mp + + calculator = mace_mp(**calculator_kwargs) + + elif calculator_name == MLFF.GAP: + from quippy.potential import Potential + + calculator = Potential(**calculator_kwargs) + + elif isinstance(calculator_meta, dict): + _calculator = MontyDecoder().decode(json.dumps(calculator_meta)) + calculator_kwargs = calculator_kwargs or {} + calculator = _calculator(**calculator_kwargs) + + return calculator diff --git a/tests/forcefields/flows/test_elastic.py b/tests/forcefields/flows/test_elastic.py index 1f41671fcb..dea097eb95 100644 --- a/tests/forcefields/flows/test_elastic.py +++ b/tests/forcefields/flows/test_elastic.py @@ -11,9 +11,8 @@ def test_elastic_wf_with_mace(clean_dir, si_structure, test_dir): si_prim = SpacegroupAnalyzer(si_structure).get_primitive_standard_structure() model_path = f"{test_dir}/forcefields/mace/MACE.model" common_kwds = dict( - model=model_path, + calculator_kwargs={"model": model_path, "default_dtype": "float64"}, relax_kwargs={"fmax": 0.00001}, - model_kwargs={"default_dtype": "float64"}, ) flow = ElasticMaker( diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index e6249486f8..1deb07c0d6 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -39,8 +39,9 @@ def test_chgnet_relax_maker(si_structure, relax_cell: bool): # translate one atom to ensure a small number of relaxation steps are taken si_structure.translate_sites(0, [0, 0, 0.1]) + max_step = 25 # generate job - job = CHGNetRelaxMaker(steps=25, relax_cell=relax_cell).make(si_structure) + job = CHGNetRelaxMaker(steps=max_step, relax_cell=relax_cell).make(si_structure) # run the flow or job and ensure that it finished running successfully responses = run_locally(job, ensure_success=True) @@ -48,11 +49,14 @@ def test_chgnet_relax_maker(si_structure, relax_cell: bool): # validate job outputs output1 = responses[job.uuid][1].output assert isinstance(output1, ForceFieldTaskDocument) - assert output1.output.n_steps >= 12 if relax_cell: + assert not output1.is_force_converged + assert output1.output.n_steps == max_step + 2 assert output1.output.energy == approx(-10.62461, abs=1e-2) assert output1.output.ionic_steps[-1].magmoms[0] == approx(0.00251674, rel=1e-1) else: + assert output1.is_force_converged + assert output1.output.n_steps == 13 assert output1.output.energy == approx(-10.6274, rel=1e-2) assert output1.output.ionic_steps[-1].magmoms[0] == approx(0.00303572, rel=1e-2) @@ -78,7 +82,8 @@ def test_m3gnet_relax_maker(si_structure): si_structure.translate_sites(0, [0, 0, 0.1]) # generate job - job = M3GNetRelaxMaker(steps=25).make(si_structure) + max_step = 25 + job = M3GNetRelaxMaker(steps=max_step).make(si_structure) # run the flow or job and ensure that it finished running successfully responses = run_locally(job, ensure_success=True) @@ -86,8 +91,9 @@ def test_m3gnet_relax_maker(si_structure): # validate job outputs output1 = responses[job.uuid][1].output assert isinstance(output1, ForceFieldTaskDocument) + assert not output1.is_force_converged assert output1.output.energy == approx(-10.8, abs=0.2) - assert output1.output.n_steps == 27 + assert output1.output.n_steps == max_step + 2 mace_paths = pytest.mark.parametrize( @@ -106,9 +112,9 @@ def test_mace_static_maker(si_structure, test_dir, model): # generate job # NOTE the test model is not trained on Si, so the energy is not accurate - job = MACEStaticMaker(model=model, task_document_kwargs=task_doc_kwargs).make( - si_structure - ) + job = MACEStaticMaker( + calculator_kwargs={"model": model}, task_document_kwargs=task_doc_kwargs + ).make(si_structure) # run the flow or job and ensure that it finished running successfully responses = run_locally(job, ensure_success=True) @@ -129,7 +135,7 @@ def test_mace_relax_maker(si_structure, test_dir, model, relax_cell: bool): # generate job # NOTE the test model is not trained on Si, so the energy is not accurate job = MACERelaxMaker( - model=model, + calculator_kwargs={"model": model}, steps=25, optimizer_kwargs={"optimizer": "BFGSLineSearch"}, relax_cell=relax_cell, @@ -140,6 +146,7 @@ def test_mace_relax_maker(si_structure, test_dir, model, relax_cell: bool): # validating the outputs of the job output1 = responses[job.uuid][1].output + assert output1.is_force_converged assert isinstance(output1, ForceFieldTaskDocument) if relax_cell: assert output1.output.energy == approx(-0.0526856, rel=1e-1) @@ -157,8 +164,10 @@ def test_gap_static_maker(si_structure, test_dir): # generate job # Test files have been provided by Yuanbin Liu (University of Oxford) job = GAPStaticMaker( - potential_args_str="IP GAP", - potential_param_file_name=test_dir / "forcefields" / "gap" / "gap_file.xml", + calculator_kwargs={ + "args_str": "IP GAP", + "param_filename": str(test_dir / "forcefields" / "gap" / "gap_file.xml"), + }, task_document_kwargs=task_doc_kwargs, ).make(si_structure) @@ -181,9 +190,13 @@ def test_gap_relax_maker(si_structure, test_dir, relax_cell: bool): # generate job # Test files have been provided by Yuanbin Liu (University of Oxford) + max_step = 25 job = GAPRelaxMaker( - potential_param_file_name=test_dir / "forcefields" / "gap" / "gap_file.xml", - steps=25, + calculator_kwargs={ + "args_str": "IP GAP", + "param_filename": str(test_dir / "forcefields" / "gap" / "gap_file.xml"), + }, + steps=max_step, relax_cell=relax_cell, ).make(si_structure) @@ -194,8 +207,10 @@ def test_gap_relax_maker(si_structure, test_dir, relax_cell: bool): output1 = responses[job.uuid][1].output assert isinstance(output1, ForceFieldTaskDocument) if relax_cell: + assert not output1.is_force_converged assert output1.output.energy == approx(-13.08492, rel=1e-2) - assert output1.output.n_steps == 27 + assert output1.output.n_steps == max_step + 2 else: + assert output1.is_force_converged assert output1.output.energy == approx(-10.8523, rel=1e-4) assert output1.output.n_steps == 17 diff --git a/tests/forcefields/test_utils.py b/tests/forcefields/test_utils.py index 9856b881d4..55f113b926 100644 --- a/tests/forcefields/test_utils.py +++ b/tests/forcefields/test_utils.py @@ -6,7 +6,13 @@ from numpy.testing import assert_allclose from pymatgen.io.ase import AseAtomsAdaptor -from atomate2.forcefields.utils import FrechetCellFilter, Relaxer, TrajectoryObserver +from atomate2.forcefields import MLFF +from atomate2.forcefields.utils import ( + FrechetCellFilter, + Relaxer, + TrajectoryObserver, + ase_calculator, +) def test_safe_import(): @@ -41,14 +47,14 @@ def test_trajectory_observer(si_structure, test_dir, tmp_dir): ] assert_allclose(traj.stresses[0], expected_stresses, atol=1e-8) - save_file_name = "log_file.json.gz" + save_file_name = "log_file.traj" traj.save(save_file_name) assert os.path.isfile(save_file_name) @pytest.mark.parametrize( ("optimizer", "traj_file"), - [("BFGS", None), (None, None), (BFGS, "log_file.json.gz")], + [("BFGS", None), (None, None), (BFGS, "log_file.traj")], ) def test_relaxer(si_structure, test_dir, tmp_dir, optimizer, traj_file): if FrechetCellFilter: @@ -109,11 +115,28 @@ 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"].energies[-1] == pytest.approx(expected_energy) + assert relax_output["trajectory"].frame_properties[-1]["energy"] == pytest.approx( + expected_energy + ) - assert_allclose(relax_output["trajectory"].forces[-1], expected_forces) + assert_allclose( + relax_output["trajectory"].frame_properties[-1]["forces"], expected_forces + ) - assert_allclose(relax_output["trajectory"].stresses[-1], expected_stresses) + assert_allclose( + relax_output["trajectory"].frame_properties[-1]["stress"], expected_stresses + ) if traj_file: assert os.path.isfile(traj_file) + + +def test_ext_load(): + forcefield_to_callable = { + "CHGNet": {"@module": "chgnet.model.dynamics", "@callable": "CHGNetCalculator"}, + "MACE": {"@module": "mace.calculators", "@callable": "mace_mp"}, + } + for forcefield in ["CHGNet", "MACE"]: + calc_from_decode = ase_calculator(forcefield_to_callable[forcefield]) + calc_from_preset = ase_calculator(f"{MLFF(forcefield)}") + assert isinstance(calc_from_decode, type(calc_from_preset)) From 21dd709dc125789a514bf897a3cf04e00096bc09 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Wed, 6 Mar 2024 08:56:40 -0800 Subject: [PATCH 25/48] Fix forcefield utils test and lint --- src/atomate2/forcefields/jobs.py | 5 +++-- src/atomate2/forcefields/md.py | 1 + src/atomate2/forcefields/utils.py | 1 + tests/forcefields/test_md.py | 2 +- tests/forcefields/test_utils.py | 1 + 5 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 9b3aa876df..64a48bc685 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -11,8 +11,7 @@ from atomate2.forcefields import MLFF from atomate2.forcefields.schemas import ForceFieldTaskDocument -from atomate2.forcefields.utils import Relaxer, ase_calculator, revert_default_dtype - +from atomate2.forcefields.utils import Relaxer, ase_calculator if TYPE_CHECKING: from pathlib import Path @@ -340,6 +339,7 @@ class MACERelaxMaker(ForceFieldRelaxMaker): optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) + @dataclass class MACEStaticMaker(ForceFieldStaticMaker): """ @@ -367,6 +367,7 @@ class MACEStaticMaker(ForceFieldStaticMaker): force_field_name: str = f"{MLFF.MACE}" task_document_kwargs: dict = field(default_factory=dict) + @dataclass class GAPRelaxMaker(ForceFieldRelaxMaker): """ diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index d42a80363f..1ee41d7db5 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -1,4 +1,5 @@ """Makers to perform MD with forcefields.""" + from __future__ import annotations import contextlib diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index c273313874..217f890f89 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -453,6 +453,7 @@ def ase_calculator( return calculator + @contextmanager def revert_default_dtype() -> Generator[None, None, None]: """Context manager for torch.default_dtype. diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index bcc77a87b8..570f8fea2c 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -1,4 +1,4 @@ -""" Tests for forcefield MD flows. """ +"""Tests for forcefield MD flows.""" import numpy as np import pytest diff --git a/tests/forcefields/test_utils.py b/tests/forcefields/test_utils.py index 6c69c94f92..55f113b926 100644 --- a/tests/forcefields/test_utils.py +++ b/tests/forcefields/test_utils.py @@ -125,6 +125,7 @@ def test_relaxer(si_structure, test_dir, tmp_dir, optimizer, traj_file): assert_allclose( relax_output["trajectory"].frame_properties[-1]["stress"], expected_stresses + ) if traj_file: assert os.path.isfile(traj_file) From 9ba7d298febd5c2ff954758b75fc3e0383a0c613 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Wed, 6 Mar 2024 10:01:21 -0800 Subject: [PATCH 26/48] Add revert_dtype env for running forcefield relax and md jobs / undo removal --- src/atomate2/forcefields/jobs.py | 11 +++---- src/atomate2/forcefields/md.py | 50 ++++++++++++++++++-------------- tests/forcefields/test_md.py | 4 --- 3 files changed, 34 insertions(+), 31 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 64a48bc685..ef0293a734 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -11,7 +11,7 @@ from atomate2.forcefields import MLFF from atomate2.forcefields.schemas import ForceFieldTaskDocument -from atomate2.forcefields.utils import Relaxer, ase_calculator +from atomate2.forcefields.utils import Relaxer, ase_calculator, revert_default_dtype if TYPE_CHECKING: from pathlib import Path @@ -121,10 +121,11 @@ def make( "Behavior may vary..." ) - relaxer = Relaxer( - self._calculator(), relax_cell=self.relax_cell, **self.optimizer_kwargs - ) - result = relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) + with revert_default_dtype(): + relaxer = Relaxer( + self._calculator(), relax_cell=self.relax_cell, **self.optimizer_kwargs + ) + result = relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) return ForceFieldTaskDocument.from_ase_compatible_result( self.force_field_name, diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 1ee41d7db5..692dbfe6ce 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -30,7 +30,11 @@ from atomate2.forcefields import MLFF from atomate2.forcefields.jobs import forcefield_job from atomate2.forcefields.schemas import ForcefieldResult, ForceFieldTaskDocument -from atomate2.forcefields.utils import TrajectoryObserver, ase_calculator +from atomate2.forcefields.utils import ( + TrajectoryObserver, + ase_calculator, + revert_default_dtype, +) if TYPE_CHECKING: from pathlib import Path @@ -292,34 +296,36 @@ def make( ZeroRotation(atoms) self.calculator_kwargs = self.calculator_kwargs or {} - atoms.calc = self._calculator() - with contextlib.redirect_stdout(io.StringIO()): - md_observer = TrajectoryObserver(atoms, store_md_outputs=True) + with revert_default_dtype(): + atoms.calc = self._calculator() - md_runner = md_func( - atoms=atoms, - timestep=self.timestep * units.fs, - **self.ase_md_kwargs, - ) + with contextlib.redirect_stdout(io.StringIO()): + md_observer = TrajectoryObserver(atoms, store_md_outputs=True) + + md_runner = md_func( + atoms=atoms, + timestep=self.timestep * units.fs, + **self.ase_md_kwargs, + ) - md_runner.attach(md_observer, interval=self.traj_interval) + md_runner.attach(md_observer, interval=self.traj_interval) - def _callback(dyn: MolecularDynamics = md_runner) -> None: - if self.ensemble == "nve": - return - dyn.set_temperature(temperature_K=self.tschedule[dyn.nsteps]) - if self.ensemble == "nvt": - return - dyn.set_stress(self.pschedule[dyn.nsteps] * 1e3 * units.bar) + def _callback(dyn: MolecularDynamics = md_runner) -> None: + if self.ensemble == "nve": + return + dyn.set_temperature(temperature_K=self.tschedule[dyn.nsteps]) + if self.ensemble == "nvt": + return + dyn.set_stress(self.pschedule[dyn.nsteps] * 1e3 * units.bar) - md_runner.attach(_callback, interval=1) - md_runner.run(steps=self.nsteps) + md_runner.attach(_callback, interval=1) + md_runner.run(steps=self.nsteps) - if self.traj_file is not None: - md_observer.save(filename=self.traj_file, fmt=self.traj_file_fmt) + if self.traj_file is not None: + md_observer.save(filename=self.traj_file, fmt=self.traj_file_fmt) - structure = AseAtomsAdaptor.get_structure(atoms) + structure = AseAtomsAdaptor.get_structure(atoms) self.task_document_kwargs = self.task_document_kwargs or {} diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 570f8fea2c..d1446e3551 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -2,7 +2,6 @@ import numpy as np import pytest -import torch from ase import units from ase.io import Trajectory from jobflow import run_locally @@ -13,9 +12,6 @@ _to_maker = {"CHGNet": CHGNetMDMaker, "M3GNet": M3GNetMDMaker, "MACE": MACEMDMaker} -# MACE changes the default dtype, ensure consistent dtype here -torch.set_default_dtype(torch.float32) - @pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE"]) def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): From 8a2451ef3d4f43f6c34c89a43acca6970ae63c44 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 7 Mar 2024 16:07:47 -0800 Subject: [PATCH 27/48] refactor nequip jobs, add md default for nequip --- src/atomate2/forcefields/jobs.py | 62 ++++++++++--------------------- src/atomate2/forcefields/md.py | 22 +++++++++-- src/atomate2/forcefields/utils.py | 32 +++++++--------- tests/forcefields/test_jobs.py | 10 ++++- 4 files changed, 59 insertions(+), 67 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 51f19594db..e6d7c6ea1c 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -14,8 +14,9 @@ from atomate2.forcefields.utils import Relaxer, ase_calculator, revert_default_dtype if TYPE_CHECKING: + from collections.abc import Sequence from pathlib import Path - from typing import Callable + from typing import Any, Callable from ase.calculators.calculator import Calculator from pymatgen.core.structure import Structure @@ -85,6 +86,8 @@ class ForceFieldRelaxMaker(Maker): Keyword arguments that will get passed to :obj:`Relaxer.relax`. optimizer_kwargs : dict Keyword arguments that will get passed to :obj:`Relaxer()`. + calculator_args : Sequence[Any] + Keyword arguments that will get passed to the ASE calculator. calculator_kwargs : dict Keyword arguments that will get passed to the ASE calculator. task_document_kwargs : dict @@ -97,6 +100,7 @@ class ForceFieldRelaxMaker(Maker): steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) + calculator_args: Sequence[Any] = field(default_factory=list) calculator_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) @@ -139,7 +143,9 @@ def make( def _calculator(self) -> Calculator: """ASE calculator, can be overwritten by user.""" - return ase_calculator(self.force_field_name, self.calculator_kwargs) + return ase_calculator( + self.force_field_name, *self.calculator_args, **self.calculator_kwargs + ) @dataclass @@ -302,13 +308,6 @@ class NequipRelaxMaker(ForceFieldRelaxMaker): Keyword arguments that will get passed to :obj:`Relaxer()`. task_document_kwargs : dict Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`. - model_path: str | Path - deployed model checkpoint to load with - :obj:`nequip.calculators.NequipCalculator.from_deployed_model()'`. - model_kwargs: dict[str, Any] - Further keywords (e.g. device: Union[str, torch.device], - species_to_type_name: Optional[Dict[str, str]] = None) for - :obj:`nequip.calculators.NequipCalculator()'`. """ name: str = f"{MLFF.Nequip} relax" @@ -318,20 +317,6 @@ class NequipRelaxMaker(ForceFieldRelaxMaker): relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - model_path: str | Path = "" - model_kwargs: dict = field(default_factory=dict) - - def _relax(self, structure: Structure) -> dict: - from nequip.ase import NequIPCalculator - - calculator = NequIPCalculator.from_deployed_model( - self.model_path, **self.model_kwargs - ) - relaxer = Relaxer( - calculator, relax_cell=self.relax_cell, **self.optimizer_kwargs - ) - - return relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) @dataclass @@ -347,30 +332,11 @@ class NequipStaticMaker(ForceFieldStaticMaker): The name of the force field. task_document_kwargs : dict Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`. - model_path: str | Path - deployed model checkpoint to load with - :obj:`nequip.calculators.NequipCalculator()'`. - model_kwargs: dict[str, Any] - Further keywords (e.g. device: Union[str, torch.device], - species_to_type_name: Optional[Dict[str, str]] = None) for - :obj:`nequip.calculators.NequipCalculator()'`. """ name: str = f"{MLFF.Nequip} static" force_field_name: str = f"{MLFF.Nequip}" task_document_kwargs: dict = field(default_factory=dict) - model_path: str | Path = "" - model_kwargs: dict = field(default_factory=dict) - - def _evaluate_static(self, structure: Structure) -> dict: - from nequip.ase import NequIPCalculator - - calculator = NequIPCalculator.from_deployed_model( - self.model_path, **self.model_kwargs - ) - relaxer = Relaxer(calculator, relax_cell=False) - - return relaxer.relax(structure, steps=1) @dataclass @@ -490,6 +456,12 @@ class GAPRelaxMaker(ForceFieldRelaxMaker): steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field( + default_factory=lambda: { + "args_str": "IP GAP", + "param_filename": "gap.xml", + } + ) task_document_kwargs: dict = field(default_factory=dict) @@ -511,3 +483,9 @@ class GAPStaticMaker(ForceFieldStaticMaker): name: str = f"{MLFF.GAP} static" force_field_name: str = f"{MLFF.GAP}" task_document_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field( + default_factory=lambda: { + "args_str": "IP GAP", + "param_filename": "gap.xml", + } + ) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 692dbfe6ce..f30d3e5c41 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -351,7 +351,7 @@ def _calculator(self) -> Calculator: class MACEMDMaker(ForceFieldMDMaker): """Perform an MD run with MACE.""" - name: str = "MACE MD" + name: str = f"{MLFF.MACE} MD" force_field_name: str = f"{MLFF.MACE}" calculator_kwargs: dict = field( default_factory=lambda: {"default_dtype": "float32"} @@ -362,7 +362,7 @@ class MACEMDMaker(ForceFieldMDMaker): class M3GNetMDMaker(ForceFieldMDMaker): """Perform an MD run with M3GNet.""" - name: str = "M3GNet MD" + name: str = f"{MLFF.M3GNet} MD" force_field_name: str = f"{MLFF.M3GNet}" @@ -370,7 +370,7 @@ class M3GNetMDMaker(ForceFieldMDMaker): class CHGNetMDMaker(ForceFieldMDMaker): """Perform an MD run with CHGNet.""" - name: str = "CHGNet MD" + name: str = f"{MLFF.CHGNet} MD" force_field_name: str = f"{MLFF.CHGNet}" @@ -378,5 +378,19 @@ class CHGNetMDMaker(ForceFieldMDMaker): class GAPMDMaker(ForceFieldMDMaker): """Perform an MD run with GAP.""" - name: str = "CHGNet MD" + name: str = f"{MLFF.GAP} MD" force_field_name: str = f"{MLFF.GAP}" + calculator_kwargs: dict = field( + default_factory=lambda: { + "args_str": "IP GAP", + "param_filename": "gap.xml", + } + ) + + +@dataclass +class NequipMDMaker(ForceFieldMDMaker): + """Perform an MD run with nequip.""" + + name: str = f"{MLFF.Nequip} MD" + force_field_name: str = f"{MLFF.Nequip}" diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 217f890f89..0c3c73dc4e 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -68,13 +68,6 @@ "BFGSLineSearch": BFGSLineSearch, } -_ase_calculator_defaults = { - "gap": { - "args_str": "IP GAP", - "param_filename": "gap.xml", - } -} - def _get_pymatgen_trajectory_from_observer( trajectory_observer: Any, frame_property_keys: list[str] @@ -390,7 +383,7 @@ def relax( def ase_calculator( - calculator_meta: str | dict, calculator_kwargs: dict | None = None + calculator_meta: str | dict, *args: Any, **kwargs: Any ) -> Calculator | None: """ Create an ASE calculator from a given set of metadata. @@ -407,8 +400,8 @@ def ase_calculator( "@callable": "CHGNetCalculator" } ``` - calculator_kwargs : dict or None (default) - kwargs to pass to the calculator object + args : optional args to pass to a calculator + kwargs : optional kwargs to pass to a calculator Returns ------- @@ -420,36 +413,37 @@ def ase_calculator( f"{name}" for name in MLFF ]: calculator_name = MLFF(calculator_meta.split("MLFF.")[-1]) - calculator_kwargs = calculator_kwargs or _ase_calculator_defaults.get( - calculator_meta, {} - ) if calculator_name == MLFF.CHGNet: from chgnet.model.dynamics import CHGNetCalculator - calculator = CHGNetCalculator(**calculator_kwargs) + calculator = CHGNetCalculator(**kwargs) elif calculator_name == MLFF.M3GNet: import matgl from matgl.ext.ase import PESCalculator potential = matgl.load_model("M3GNet-MP-2021.2.8-PES") - calculator = PESCalculator(potential, **calculator_kwargs) + calculator = PESCalculator(potential, **kwargs) elif calculator_name == MLFF.MACE: from mace.calculators import mace_mp - calculator = mace_mp(**calculator_kwargs) + calculator = mace_mp(**kwargs) elif calculator_name == MLFF.GAP: from quippy.potential import Potential - calculator = Potential(**calculator_kwargs) + calculator = Potential(**kwargs) + + elif calculator_name == MLFF.Nequip: + from nequip.ase import NequIPCalculator + + calculator = NequIPCalculator.from_deployed_model(*args, **kwargs) elif isinstance(calculator_meta, dict): _calculator = MontyDecoder().decode(json.dumps(calculator_meta)) - calculator_kwargs = calculator_kwargs or {} - calculator = _calculator(**calculator_kwargs) + calculator = _calculator(*args, **kwargs) return calculator diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 85abaa2d0e..638a4e67d8 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -187,13 +187,16 @@ def test_gap_static_maker(si_structure: Structure, test_dir): def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): + importorskip("nequip") task_doc_kwargs = {"ionic_step_data": ("structure", "energy")} # generate job # NOTE the test model is not trained on Si, so the energy is not accurate job = NequipStaticMaker( task_document_kwargs=task_doc_kwargs, - model_path=test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth", + calculator_args=[ + test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" + ], ).make(sr_ti_o3_structure) # run the flow or job and ensure that it finished running successfully @@ -210,6 +213,7 @@ def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): def test_nequip_relax_maker( sr_ti_o3_structure: Structure, test_dir: Path, relax_cell: bool ): + importorskip("nequip") # translate one atom to ensure a small number of relaxation steps are taken sr_ti_o3_structure.translate_sites(0, [0, 0, 0.2]) # generate job @@ -217,7 +221,9 @@ def test_nequip_relax_maker( steps=25, optimizer_kwargs={"optimizer": "BFGSLineSearch"}, relax_cell=relax_cell, - model_path=test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth", + calculator_args=[ + test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" + ], ).make(sr_ti_o3_structure) # run the flow or job and ensure that it finished running successfully From 38ad2d6a6c34c84b137b22a4845ce86ac2c28a59 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 7 Mar 2024 16:30:54 -0800 Subject: [PATCH 28/48] Add GAP and Nequip MD tests, fix arg / kwarg passing in MD --- src/atomate2/forcefields/md.py | 15 ++++++----- tests/forcefields/test_md.py | 47 ++++++++++++++++++++++++++++------ 2 files changed, 48 insertions(+), 14 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index f30d3e5c41..f5664991c9 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -38,7 +38,7 @@ if TYPE_CHECKING: from pathlib import Path - from typing import Literal + from typing import Any, Literal from ase.calculators.calculator import Calculator from pymatgen.core.structure import Structure @@ -106,7 +106,9 @@ class ForceFieldMDMaker(Maker): specifying thermostat as a string. ase_md_kwargs : dict | None = None Options except for temperature and pressure to pass into the ASE MD function - calculator_kwargs : dict | None = None + calculator_args : Sequence[Any] + args to pass to the ASE calculator class + calculator_kwargs : dict kwargs to pass to the ASE calculator class traj_file : str | Path | None = None If a str or Path, the name of the file to save the MD trajectory to. @@ -133,7 +135,8 @@ class ForceFieldMDMaker(Maker): # end_temp: float | None = 300.0 pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None - calculator_kwargs: dict | None = None + calculator_args: Sequence[Any] = field(default_factory=list) + calculator_kwargs: dict = field(default_factory=dict) traj_file: str | Path | None = None traj_file_fmt: Literal["pmg", "ase"] = "ase" traj_interval: int = 1 @@ -295,8 +298,6 @@ def make( if self.zero_angular_momentum: ZeroRotation(atoms) - self.calculator_kwargs = self.calculator_kwargs or {} - with revert_default_dtype(): atoms.calc = self._calculator() @@ -344,7 +345,9 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: def _calculator(self) -> Calculator: """ASE calculator, can be overwritten by user.""" - return ase_calculator(self.force_field_name, self.calculator_kwargs) + return ase_calculator( + self.force_field_name, *self.calculator_args, **self.calculator_kwargs + ) @dataclass diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index d1446e3551..ea5676a2ea 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -8,30 +8,61 @@ from monty.serialization import loadfn from pymatgen.analysis.structure_matcher import StructureMatcher -from atomate2.forcefields.md import CHGNetMDMaker, M3GNetMDMaker, MACEMDMaker - -_to_maker = {"CHGNet": CHGNetMDMaker, "M3GNet": M3GNetMDMaker, "MACE": MACEMDMaker} - - -@pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE"]) -def test_ml_ff_md_maker(ff_name, si_structure, clean_dir): +from atomate2.forcefields.md import ( + CHGNetMDMaker, + GAPMDMaker, + M3GNetMDMaker, + MACEMDMaker, + NequipMDMaker, +) + +_to_maker = { + "CHGNet": CHGNetMDMaker, + "M3GNet": M3GNetMDMaker, + "MACE": MACEMDMaker, + "GAP": GAPMDMaker, + "Nequip": NequipMDMaker, +} + + +@pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE", "GAP", "Nequip"]) +def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, clean_dir): nsteps = 5 ref_energies_per_atom = { "CHGNet": -5.280157089233398, "M3GNet": -5.387282371520996, "MACE": -5.311369895935059, + "GAP": -5.391255755606209, + "Nequip": -8.84670181274414, } - structure = si_structure.to_conventional() * (2, 2, 2) # ASE can slightly change tolerances on structure positions matcher = StructureMatcher() + calculator_kwargs = {} + calculator_args = [] + unit_cell_structure = si_structure.copy() + if ff_name == "GAP": + calculator_kwargs = { + "args_str": "IP GAP", + "param_filename": str(test_dir / "forcefields" / "gap" / "gap_file.xml"), + } + elif ff_name == "Nequip": + calculator_args = [ + test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" + ] + unit_cell_structure = sr_ti_o3_structure.copy() + + structure = unit_cell_structure.to_conventional() * (2, 2, 2) + job = _to_maker[ff_name]( nsteps=nsteps, traj_file="md_traj.json.gz", traj_file_fmt="pmg", task_document_kwargs={"store_trajectory": "partial"}, + calculator_args=calculator_args, + calculator_kwargs=calculator_kwargs, ).make(structure) response = run_locally(job, ensure_success=True) taskdoc = response[next(iter(response))][1].output From ac4916bd428fe3b416c9f0d12f95d95e96baf314 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 8 Mar 2024 09:05:41 -0800 Subject: [PATCH 29/48] remove calculator_args / ase_calculator args, revert phonon job change --- src/atomate2/common/jobs/phonons.py | 10 ++++------ src/atomate2/forcefields/jobs.py | 10 ++-------- src/atomate2/forcefields/md.py | 9 ++------- src/atomate2/forcefields/utils.py | 8 +++----- tests/forcefields/test_jobs.py | 12 ++++++------ tests/forcefields/test_md.py | 8 +++----- 6 files changed, 20 insertions(+), 37 deletions(-) diff --git a/src/atomate2/common/jobs/phonons.py b/src/atomate2/common/jobs/phonons.py index e70b080d56..d273f1546a 100644 --- a/src/atomate2/common/jobs/phonons.py +++ b/src/atomate2/common/jobs/phonons.py @@ -153,13 +153,11 @@ def generate_phonon_displacements( "of the symmetry of the structure and thus will be removed now.", stacklevel=1, ) - - # Note that `remove_site_property` returns None - if "magmom" in structure.site_properties: + cell = get_phonopy_structure( structure.remove_site_property(property_name="magmom") - - cell = get_phonopy_structure(structure) - + if "magmom" in structure.site_properties + else structure + ) factor = get_factor(code) # a bit of code repetition here as I currently diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index e6d7c6ea1c..c52676b4f6 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -14,9 +14,8 @@ from atomate2.forcefields.utils import Relaxer, ase_calculator, revert_default_dtype if TYPE_CHECKING: - from collections.abc import Sequence from pathlib import Path - from typing import Any, Callable + from typing import Callable from ase.calculators.calculator import Calculator from pymatgen.core.structure import Structure @@ -86,8 +85,6 @@ class ForceFieldRelaxMaker(Maker): Keyword arguments that will get passed to :obj:`Relaxer.relax`. optimizer_kwargs : dict Keyword arguments that will get passed to :obj:`Relaxer()`. - calculator_args : Sequence[Any] - Keyword arguments that will get passed to the ASE calculator. calculator_kwargs : dict Keyword arguments that will get passed to the ASE calculator. task_document_kwargs : dict @@ -100,7 +97,6 @@ class ForceFieldRelaxMaker(Maker): steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) - calculator_args: Sequence[Any] = field(default_factory=list) calculator_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) @@ -143,9 +139,7 @@ def make( def _calculator(self) -> Calculator: """ASE calculator, can be overwritten by user.""" - return ase_calculator( - self.force_field_name, *self.calculator_args, **self.calculator_kwargs - ) + return ase_calculator(self.force_field_name, **self.calculator_kwargs) @dataclass diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index f5664991c9..8d8a2da1c1 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -38,7 +38,7 @@ if TYPE_CHECKING: from pathlib import Path - from typing import Any, Literal + from typing import Literal from ase.calculators.calculator import Calculator from pymatgen.core.structure import Structure @@ -106,8 +106,6 @@ class ForceFieldMDMaker(Maker): specifying thermostat as a string. ase_md_kwargs : dict | None = None Options except for temperature and pressure to pass into the ASE MD function - calculator_args : Sequence[Any] - args to pass to the ASE calculator class calculator_kwargs : dict kwargs to pass to the ASE calculator class traj_file : str | Path | None = None @@ -135,7 +133,6 @@ class ForceFieldMDMaker(Maker): # end_temp: float | None = 300.0 pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None - calculator_args: Sequence[Any] = field(default_factory=list) calculator_kwargs: dict = field(default_factory=dict) traj_file: str | Path | None = None traj_file_fmt: Literal["pmg", "ase"] = "ase" @@ -345,9 +342,7 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: def _calculator(self) -> Calculator: """ASE calculator, can be overwritten by user.""" - return ase_calculator( - self.force_field_name, *self.calculator_args, **self.calculator_kwargs - ) + return ase_calculator(self.force_field_name, **self.calculator_kwargs) @dataclass diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 0c3c73dc4e..ef997d5631 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -382,9 +382,7 @@ def relax( ) -def ase_calculator( - calculator_meta: str | dict, *args: Any, **kwargs: Any -) -> Calculator | None: +def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | None: """ Create an ASE calculator from a given set of metadata. @@ -439,11 +437,11 @@ def ase_calculator( elif calculator_name == MLFF.Nequip: from nequip.ase import NequIPCalculator - calculator = NequIPCalculator.from_deployed_model(*args, **kwargs) + calculator = NequIPCalculator.from_deployed_model(**kwargs) elif isinstance(calculator_meta, dict): _calculator = MontyDecoder().decode(json.dumps(calculator_meta)) - calculator = _calculator(*args, **kwargs) + calculator = _calculator(**kwargs) return calculator diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 638a4e67d8..7136e834d4 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -194,9 +194,9 @@ def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): # NOTE the test model is not trained on Si, so the energy is not accurate job = NequipStaticMaker( task_document_kwargs=task_doc_kwargs, - calculator_args=[ - test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" - ], + calculator_kwargs={ + "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" + }, ).make(sr_ti_o3_structure) # run the flow or job and ensure that it finished running successfully @@ -221,9 +221,9 @@ def test_nequip_relax_maker( steps=25, optimizer_kwargs={"optimizer": "BFGSLineSearch"}, relax_cell=relax_cell, - calculator_args=[ - test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" - ], + calculator_kwargs={ + "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" + }, ).make(sr_ti_o3_structure) # run the flow or job and ensure that it finished running successfully diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index ea5676a2ea..c863c03068 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -41,7 +41,6 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle matcher = StructureMatcher() calculator_kwargs = {} - calculator_args = [] unit_cell_structure = si_structure.copy() if ff_name == "GAP": calculator_kwargs = { @@ -49,9 +48,9 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle "param_filename": str(test_dir / "forcefields" / "gap" / "gap_file.xml"), } elif ff_name == "Nequip": - calculator_args = [ - test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" - ] + calculator_kwargs = { + "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" + } unit_cell_structure = sr_ti_o3_structure.copy() structure = unit_cell_structure.to_conventional() * (2, 2, 2) @@ -61,7 +60,6 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle traj_file="md_traj.json.gz", traj_file_fmt="pmg", task_document_kwargs={"store_trajectory": "partial"}, - calculator_args=calculator_args, calculator_kwargs=calculator_kwargs, ).make(structure) response = run_locally(job, ensure_success=True) From c19b39b039a0a52acc87575fad7faebbcbd94eb1 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 8 Mar 2024 09:09:16 -0800 Subject: [PATCH 30/48] remove todo about adding magmoms to forcefield traj observer as that's now implemented --- src/atomate2/forcefields/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index ef997d5631..034d950f21 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -144,7 +144,6 @@ def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: def __call__(self) -> None: """Save the properties of an Atoms during the relaxation.""" - # TODO: maybe include magnetic moments self.energies.append(self.compute_energy()) self.forces.append(self.atoms.get_forces()) # TODO: MD needs kinetic energy parts of stress, From d95397d064c4cef8cc445869abe143c153effa5f Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Wed, 13 Mar 2024 14:26:03 -0700 Subject: [PATCH 31/48] Remove comments, add option to seed rng for MB velocities, turn on ideal gas stress contribution only when MD trajectory outputs (temp / velocity) are stored --- src/atomate2/forcefields/md.py | 47 ++++++++++--------------------- src/atomate2/forcefields/utils.py | 13 ++++----- 2 files changed, 21 insertions(+), 39 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 8d8a2da1c1..c2d750e623 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -99,11 +99,11 @@ class ForceFieldMDMaker(Maker): The pressure in kilobar. If a sequence or 1D array, the pressure schedule will be interpolated linearly between the given values. If a float, the pressure will be constant throughout the run. - thermostat : str | ASE .MolecularDynamics = "langevin" - The thermostat to use. If thermostat is an ASE .MolecularDynamics + dynamics : str | ASE .MolecularDynamics = "langevin" + The dynamical thermostat to use. If dynamics is an ASE .MolecularDynamics object, this uses the option specified explicitly by the user. See _valid_dynamics for a list of pre-defined options when - specifying thermostat as a string. + specifying dynamics as a string. ase_md_kwargs : dict | None = None Options except for temperature and pressure to pass into the ASE MD function calculator_kwargs : dict @@ -111,15 +111,24 @@ class ForceFieldMDMaker(Maker): traj_file : str | Path | None = None If a str or Path, the name of the file to save the MD trajectory to. If None, the trajectory is not written to disk + traj_file_fmt : Literal["ase","pmg"] + The format of the trajectory file to write. If "ase", writes an + ASE trajectory, if "pmg", writes a Pymatgen trajectory. traj_interval : int The step interval for saving the trajectories. + mb_velocity_seed : int | None = None + If an int, a random number seed for generating initial velocities + from a Maxwell-Boltzmann distribution. zero_linear_momentum : bool = False Whether to initialize the atomic velocities with zero linear momentum zero_angular_momentum : bool = False Whether to initialize the atomic velocities with zero angular momentum task_document_kwargs: dict Options to pass to the TaskDoc. Default choice - {"store_trajectory": "partial"} + { + "store_trajectory": "partial", + "ionic_step_data": ("energy",), + } is consistent with atomate2.vasp.md.MDMaker """ @@ -130,13 +139,13 @@ class ForceFieldMDMaker(Maker): ensemble: Literal["nve", "nvt", "npt"] = "nvt" dynamics: str | MolecularDynamics = "langevin" temperature: float | Sequence | np.ndarray | None = 300.0 - # end_temp: float | None = 300.0 pressure: float | Sequence | np.ndarray | None = None ase_md_kwargs: dict | None = None calculator_kwargs: dict = field(default_factory=dict) traj_file: str | Path | None = None traj_file_fmt: Literal["pmg", "ase"] = "ase" traj_interval: int = 1 + mb_velocity_seed: int | None = None zero_linear_momentum: bool = False zero_angular_momentum: bool = False task_document_kwargs: dict = field( @@ -202,37 +211,11 @@ def _get_ensemble_defaults(self) -> None: self.ase_md_kwargs["temperature_K"] = self.tschedule[0] self.ase_md_kwargs["externalstress"] = self.pschedule[0] * 1e3 * units.bar - # NOTE: We take in the temperature in _get_ensemble_schedule instead - # if self.ensemble in ("nvt", "npt") and all( - # self.ase_md_kwargs.get(key) is None - # for key in ("temperature_K", "temperature") - # ): - # self.ase_md_kwargs["temperature_K"] = ( - # self.end_temp if self.end_temp else self.start_temp - # ) - - # NOTE: We take care of the units when passing into the ASE MD function - # if self.ensemble == "npt" and isinstance(self.pressure, float): - # # convert from kilobar to eV/Ang**3 - # self.ase_md_kwargs["pressure_au"] = self.pressure * 1.0e-3 / bar - if self.dynamics.lower() == "langevin": - # NOTE: Unless we have a detailed documentation on the conversion of all - # parameters for all different ASE dynamics, it is not a good idea to - # convert the parameters here. It is better to let the user to consult - # the ASE documentation and set the parameters themselves. self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get( "friction", 10.0 * 1e-3 / units.fs, # Same default as in VASP: 10 ps^-1 ) - # NOTE: same as above, user can specify per atom friction but we don't - # expect to change their intention to pass into ASE MD function - # if isinstance(self.ase_md_kwargs["friction"], (list, tuple)): - # self.ase_md_kwargs["friction"] = [ - # coeff * 1.0e-3 / fs for coeff in self.ase_md_kwargs["friction"] - # ] - # else: - # self.ase_md_kwargs["friction"] *= 1.0e-3 / fs @forcefield_job def make( @@ -288,7 +271,7 @@ def make( MaxwellBoltzmannDistribution( atoms=atoms, temperature_K=self.tschedule[0], - rng=None, # TODO: we might want to use a seed for reproducibility + rng=np.random.default_rng(seed=self.mb_velocity_seed), ) if self.zero_linear_momentum: Stationary(atoms) diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 034d950f21..4efd0d53ad 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -146,10 +146,12 @@ def __call__(self) -> None: """Save the properties of an Atoms during the relaxation.""" self.energies.append(self.compute_energy()) self.forces.append(self.atoms.get_forces()) - # TODO: MD needs kinetic energy parts of stress, - # we might just need to refactor TrajectoryObserver - # Now it is safe for 0K relaxation as the atoms don't have momenta - self.stresses.append(self.atoms.get_stress(include_ideal_gas=True)) + # MD needs kinetic energy parts of stress, relaxations do not + # When _store_md_outputs is True, ideal gas contribution to + # stress is included. + self.stresses.append( + self.atoms.get_stress(include_ideal_gas=self._store_md_outputs) + ) if self._store_magmoms: try: @@ -163,9 +165,6 @@ def __call__(self) -> None: if self._store_md_outputs: self.velocities.append(self.atoms.get_velocities()) self.temperatures.append(self.atoms.get_temperature()) - # self.stresses.append(self.atoms.get_stress( - # voigt=True, include_ideal_gas=True) - # ) def compute_energy(self) -> float: """ From 8ecc1d16dd86c50466ad8ea8cef3577d69b03fb4 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 25 Mar 2024 10:19:23 -0700 Subject: [PATCH 32/48] Ensure CHGNet and M3GNet relax / static makers convert stress to eV/A0**3; add tests for MD NVE ensemble and specifying MolecularDynamics object as input --- src/atomate2/forcefields/jobs.py | 13 +++++++++ src/atomate2/forcefields/md.py | 14 +++++----- tests/forcefields/test_jobs.py | 4 +-- tests/forcefields/test_md.py | 46 ++++++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 9 deletions(-) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index c52676b4f6..00d4f7eed0 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING +from ase.units import GPa as _GPa_to_eV_per_A3 from jobflow import Maker, job from pymatgen.core.trajectory import Trajectory @@ -229,6 +230,9 @@ class CHGNetRelaxMaker(ForceFieldRelaxMaker): relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field( + default_factory=lambda: {"stress_weight": _GPa_to_eV_per_A3} + ) @dataclass @@ -247,6 +251,9 @@ class CHGNetStaticMaker(ForceFieldStaticMaker): name: str = f"{MLFF.CHGNet} static" force_field_name: str = f"{MLFF.CHGNet}" task_document_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field( + default_factory=lambda: {"stress_weight": _GPa_to_eV_per_A3} + ) @dataclass @@ -279,6 +286,9 @@ class M3GNetRelaxMaker(ForceFieldRelaxMaker): relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field( + default_factory=lambda: {"stress_weight": _GPa_to_eV_per_A3} + ) @dataclass @@ -351,6 +361,9 @@ class M3GNetStaticMaker(ForceFieldStaticMaker): name: str = f"{MLFF.M3GNet} static" force_field_name: str = f"{MLFF.M3GNet}" task_document_kwargs: dict = field(default_factory=dict) + calculator_kwargs: dict = field( + default_factory=lambda: {"stress_weight": _GPa_to_eV_per_A3} + ) @dataclass diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index c2d750e623..a6c75e1c2b 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -211,7 +211,7 @@ def _get_ensemble_defaults(self) -> None: self.ase_md_kwargs["temperature_K"] = self.tschedule[0] self.ase_md_kwargs["externalstress"] = self.pschedule[0] * 1e3 * units.bar - if self.dynamics.lower() == "langevin": + if isinstance(self.dynamics, str) and self.dynamics.lower() == "langevin": self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get( "friction", 10.0 * 1e-3 / units.fs, # Same default as in VASP: 10 ps^-1 @@ -239,12 +239,8 @@ def make( initial_velocities = structure.site_properties.get("velocities") - if isinstance(self.dynamics, MolecularDynamics): - # Allow user to explicitly run ASE Dynamics class - md_func = self.dynamics - - elif isinstance(self.dynamics, str): - # Otherwise, use known dynamics + if isinstance(self.dynamics, str): + # Use known dynamics if `self.dynamics` is a str self.dynamics = self.dynamics.lower() if self.dynamics not in _valid_dynamics[self.ensemble]: raise ValueError( @@ -257,6 +253,10 @@ def make( self.dynamics = "velocityverlet" md_func = _preset_dynamics[f"{self.ensemble}_{self.dynamics}"] + elif issubclass(self.dynamics, MolecularDynamics): + # Allow user to explicitly run ASE Dynamics class + md_func = self.dynamics + atoms = structure.to_ase_atoms() if md_func is NPT: diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 7136e834d4..f70a358702 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -94,9 +94,9 @@ def test_m3gnet_relax_maker(si_structure): # validate job outputs output1 = responses[job.uuid][1].output assert isinstance(output1, ForceFieldTaskDocument) - assert not output1.is_force_converged + assert output1.is_force_converged assert output1.output.energy == approx(-10.8, abs=0.2) - assert output1.output.n_steps == max_step + 2 + assert output1.output.n_steps == 13 mace_paths = pytest.mark.parametrize( diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index c863c03068..bdebc33f73 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -1,12 +1,16 @@ """Tests for forcefield MD flows.""" +from pathlib import Path + import numpy as np import pytest from ase import units from ase.io import Trajectory +from ase.md.verlet import VelocityVerlet from jobflow import run_locally from monty.serialization import loadfn from pymatgen.analysis.structure_matcher import StructureMatcher +from pymatgen.core import Structure from atomate2.forcefields.md import ( CHGNetMDMaker, @@ -152,6 +156,48 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): ) +def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): + # This test serves two purposes: + # 1. Test the NVE calculator + # 2. Test specifying the `dynamics` kwarg of the `MDMaker` as a str + # vs. as an ase.md.md.MolecularDyanmics object + + output = {} + for k in ["from_str", "from_dyn"]: + if k == "from_str": + dyn = "velocityverlet" + elif k == "from_dyn": + dyn = VelocityVerlet + + job = CHGNetMDMaker( + ensemble="nve", + dynamics=dyn, + nsteps=50, + traj_file=None, + ).make(si_structure) + + response = run_locally(job, ensure_success=True) + output[k] = response[next(iter(response))][1].output + + # check that energy and volume are constants + assert output["from_str"].output.energy == pytest.approx(-10.6, abs=0.1) + assert output["from_str"].output.structure.volume == pytest.approx( + output["from_str"].input.structure.volume + ) + assert all( + step.energy == pytest.approx(-10.6, abs=0.1) + for step in output["from_str"].output.ionic_steps + ) + + # ensure that output is consistent if molecular dynamics object is specified + # as str or as MolecularDynamics object + assert all( + output["from_str"].output.__getattribute__(attr) + == output["from_dyn"].output.__getattribute__(attr) + for attr in ("energy", "forces", "stress", "structure") + ) + + @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_temp_schedule(ff_name, si_structure, clean_dir): nsteps = 100 From 7b8f584b65762e15c09a4ad3c646ce0e997c90b3 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 25 Mar 2024 11:12:24 -0700 Subject: [PATCH 33/48] fix M3GNet test related to outdated cached model --- tests/forcefields/test_jobs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index f70a358702..dcdd28217b 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -96,7 +96,7 @@ def test_m3gnet_relax_maker(si_structure): assert isinstance(output1, ForceFieldTaskDocument) assert output1.is_force_converged assert output1.output.energy == approx(-10.8, abs=0.2) - assert output1.output.n_steps == 13 + assert output1.output.n_steps == 24 mace_paths = pytest.mark.parametrize( From d3aa25dae872b5627063f5bfc78ceb6aec350200 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 28 Mar 2024 17:26:04 -0700 Subject: [PATCH 34/48] Fix dependencies, amend forcefield static option parsing --- pyproject.toml | 1 + src/atomate2/forcefields/jobs.py | 39 ++++---------------------------- 2 files changed, 5 insertions(+), 35 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 63b1d51560..d641f0b6da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ defects = [ ] forcefields = [ "ase>=3.22.1", + "torch<=2.2.1", # incompatibility with dgl if newer versions are used "chgnet>=0.2.2", "mace-torch>=0.3.3", "matgl>=1.0.0", diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 00d4f7eed0..450864a923 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -148,6 +148,10 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker): """ Maker to calculate forces and stresses using any force field. + Note that while `steps = 1` by default, the user could override + this setting along with cell shape relaxation (`relax_cell = False` + by default). + Parameters ---------- name : str @@ -167,41 +171,6 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker): calculator_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) - @forcefield_job - def make( - self, structure: Structure, prev_dir: str | Path | None = None - ) -> ForceFieldTaskDocument: - """ - Perform a relaxation of a structure using a force field. - - Parameters - ---------- - structure: .Structure - pymatgen structure. - prev_dir : str or Path or None - A previous calculation directory to copy output files from. Unused, just - added to match the method signature of other makers. - """ - if self.steps < 0: - logger.warning( - "WARNING: A negative number of steps is not possible. " - "Behavior may vary..." - ) - - static_calc = Relaxer(self._calculator(), relax_cell=False) - result = static_calc.relax(structure, steps=1) - - return ForceFieldTaskDocument.from_ase_compatible_result( - forcefield_name=self.force_field_name, - result=result, - relax_cell=False, - steps=1, - relax_kwargs=None, - optimizer_kwargs=None, - **self.task_document_kwargs, - ) - - @dataclass class CHGNetRelaxMaker(ForceFieldRelaxMaker): """ From ecf9356b57b249f8caac01306616ad038dc7d2dd Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Thu, 28 Mar 2024 17:29:00 -0700 Subject: [PATCH 35/48] linting --- src/atomate2/forcefields/jobs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 450864a923..a210f5908b 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -171,6 +171,7 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker): calculator_kwargs: dict = field(default_factory=dict) task_document_kwargs: dict = field(default_factory=dict) + @dataclass class CHGNetRelaxMaker(ForceFieldRelaxMaker): """ From 733f60e43c5dd3dbf53753ecd486d2978102ef60 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 29 Mar 2024 08:18:15 -0700 Subject: [PATCH 36/48] try to fix ci tests caused by torch / dgl incompatibility --- .github/workflows/testing.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index e5ed45cbce..a30a7a58ac 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -39,8 +39,7 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - cache: pip - cache-dependency-path: pyproject.toml + - name: Install dependencies # ERROR: Cannot install atomate2 and atomate2[strict,tests]==0.0.1 because these package versions have conflicting dependencies. @@ -54,6 +53,7 @@ jobs: python -m pip install --upgrade pip # ase needed to get FrechetCellFilter used by ML force fields pip install git+https://gitlab.com/ase/ase + pip install torch<=2.2.1 # incompatibility with dgl if newer versions are used pip install .[strict,tests] pip install torch-runstats pip install --no-deps nequip From a0ead2002327147972e1f467d9fc976497c08cf8 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Fri, 29 Mar 2024 08:35:58 -0700 Subject: [PATCH 37/48] try to fix ci tests caused by torch / dgl incompatibility --- .github/workflows/testing.yml | 1 - pyproject.toml | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index a30a7a58ac..35ab100172 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -53,7 +53,6 @@ jobs: python -m pip install --upgrade pip # ase needed to get FrechetCellFilter used by ML force fields pip install git+https://gitlab.com/ase/ase - pip install torch<=2.2.1 # incompatibility with dgl if newer versions are used pip install .[strict,tests] pip install torch-runstats pip install --no-deps nequip diff --git a/pyproject.toml b/pyproject.toml index d641f0b6da..fd35919e92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ strict = [ # must use >= for ase to not uninstall main branch install in CI # last known working commit: https://gitlab.com/ase/ase@2bab58f4e "ase>=3.22.1", + "torch==2.2.1", "cclib==1.8.1", "chgnet==0.3.5", "click==8.1.7", From d3c5ad6b2f17228554d4ee22dc37aa289fae9dca Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 1 Apr 2024 09:34:24 -0700 Subject: [PATCH 38/48] Change: default time_step, units for time_step, initialization of temperature and velocity in TrajectoryObserver, snake case change --- src/atomate2/forcefields/md.py | 40 ++++++++++++++++++------------- src/atomate2/forcefields/utils.py | 7 +++--- tests/forcefields/test_md.py | 30 +++++++++++------------ 3 files changed, 43 insertions(+), 34 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index a6c75e1c2b..000df2ce10 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -85,9 +85,11 @@ class ForceFieldMDMaker(Maker): The name of the MD Maker force_field_name : str The name of the forcefield (for provenance) - timestep : float | None = 2. - The timestep of the MD run in femtoseconds - nsteps : int = 1000 + time_step : float | None = None. + The timestep of the MD run, specified in ASE's units. + If `None`, defaults to 0.5 fs if a structure contains an isotope of + hydrogen and 2 fs otherwise. + n_steps : int = 1000 The number of MD steps to run ensemble : str = "nvt" The ensemble to use. Valid ensembles are nve, nvt, or npt @@ -134,8 +136,8 @@ class ForceFieldMDMaker(Maker): name: str = "Forcefield MD" force_field_name: str = f"{MLFF.Forcefield}" - timestep: float | None = 2.0 - nsteps: int = 1000 + time_step: float | None = None + n_steps: int = 1000 ensemble: Literal["nve", "nvt", "npt"] = "nvt" dynamics: str | MolecularDynamics = "langevin" temperature: float | Sequence | np.ndarray | None = 300.0 @@ -160,15 +162,15 @@ def _get_ensemble_schedule(self) -> None: # Distable thermostat and barostat self.temperature = np.nan self.pressure = np.nan - self.tschedule = np.full(self.nsteps + 1, self.temperature) - self.pschedule = np.full(self.nsteps + 1, self.pressure) + self.tschedule = np.full(self.n_steps + 1, self.temperature) + self.pschedule = np.full(self.n_steps + 1, self.pressure) return if isinstance(self.temperature, Sequence) or ( isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1 ): self.tschedule = np.interp( - np.arange(self.nsteps + 1), + np.arange(self.n_steps + 1), np.arange(len(self.temperature)), self.temperature, ) @@ -176,25 +178,25 @@ def _get_ensemble_schedule(self) -> None: # scalars, but in principle one quantity per atom could be specified by giving # an array. This is not implemented yet here. else: - self.tschedule = np.full(self.nsteps + 1, self.temperature) + self.tschedule = np.full(self.n_steps + 1, self.temperature) if self.ensemble == "nvt": self.pressure = np.nan - self.pschedule = np.full(self.nsteps + 1, self.pressure) + self.pschedule = np.full(self.n_steps + 1, self.pressure) return if isinstance(self.pressure, Sequence) or ( isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 1 ): self.pschedule = np.interp( - np.arange(self.nsteps + 1), np.arange(len(self.pressure)), self.pressure + np.arange(self.n_steps + 1), np.arange(len(self.pressure)), self.pressure ) elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: self.pschedule = interp1d( - np.arange(self.nsteps + 1), self.pressure, kind="linear" + np.arange(self.n_steps + 1), self.pressure, kind="linear" ) else: - self.pschedule = np.full(self.nsteps + 1, self.pressure) + self.pschedule = np.full(self.n_steps + 1, self.pressure) def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" @@ -237,6 +239,12 @@ def make( self._get_ensemble_schedule() self._get_ensemble_defaults() + if self.time_step is None: + # If a structure contains an isotope of hydrogen, set default `time_step` + # to 0.5 fs, and 2 fs otherwise. + has_h_isotope = any(element.Z == 1 for element in structure.composition) + self.time_step = (0.5 if has_h_isotope else 2.) * units.fs + initial_velocities = structure.site_properties.get("velocities") if isinstance(self.dynamics, str): @@ -286,7 +294,7 @@ def make( md_runner = md_func( atoms=atoms, - timestep=self.timestep * units.fs, + timestep=self.time_step, **self.ase_md_kwargs, ) @@ -301,7 +309,7 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: dyn.set_stress(self.pschedule[dyn.nsteps] * 1e3 * units.bar) md_runner.attach(_callback, interval=1) - md_runner.run(steps=self.nsteps) + md_runner.run(steps=self.n_steps) if self.traj_file is not None: md_observer.save(filename=self.traj_file, fmt=self.traj_file_fmt) @@ -317,7 +325,7 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: trajectory=md_observer.to_pymatgen_trajectory(None), ), relax_cell=(self.ensemble == "npt"), - steps=self.nsteps, + steps=self.n_steps, relax_kwargs=None, optimizer_kwargs=None, **self.task_document_kwargs, diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 4efd0d53ad..2d9a69d781 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -138,9 +138,10 @@ def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: self.cells: list[np.ndarray] = [] self._store_md_outputs = store_md_outputs - if self._store_md_outputs: - self.velocities: list[np.ndarray] = [] - self.temperatures: list[float] = [] + # `self.{velocities,temperatures}` always initialized, + # but data is only stored / saved to trajectory for MD runs + self.velocities: list[np.ndarray] = [] + self.temperatures: list[float] = [] def __call__(self) -> None: """Save the properties of an Atoms during the relaxation.""" diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index bdebc33f73..1a498fd802 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -31,7 +31,7 @@ @pytest.mark.parametrize("ff_name", ["CHGNet", "M3GNet", "MACE", "GAP", "Nequip"]) def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, clean_dir): - nsteps = 5 + n_steps = 5 ref_energies_per_atom = { "CHGNet": -5.280157089233398, @@ -60,7 +60,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle structure = unit_cell_structure.to_conventional() * (2, 2, 2) job = _to_maker[ff_name]( - nsteps=nsteps, + n_steps=n_steps, traj_file="md_traj.json.gz", traj_file_fmt="pmg", task_document_kwargs={"store_trajectory": "partial"}, @@ -77,7 +77,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle # Check that we have the right number of MD steps # ASE logs the initial structure energy, and doesn't count this as an MD step assert matcher.fit(taskdoc.output.ionic_steps[0].structure, structure) - assert len(taskdoc.output.ionic_steps) == nsteps + 1 + assert len(taskdoc.output.ionic_steps) == n_steps + 1 # Check that the ionic steps have the expected physical properties assert all( @@ -88,7 +88,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle # Check that the trajectory has expected physical properties assert taskdoc.included_objects == ["trajectory"] - assert len(taskdoc.forcefield_objects["trajectory"]) == nsteps + 1 + assert len(taskdoc.forcefield_objects["trajectory"]) == n_steps + 1 assert all( key in step for key in ("energy", "forces", "stress", "velocities", "temperature") @@ -98,7 +98,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle @pytest.mark.parametrize("traj_file", ["trajectory.json.gz", "atoms.traj"]) def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): - nsteps = 5 + n_steps = 5 # Check that traj file written to disk is consistent with trajectory # stored to the task document @@ -112,7 +112,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): structure = si_structure.to_conventional() * (2, 2, 2) job = _to_maker[ff_name]( - nsteps=nsteps, + n_steps=n_steps, traj_file=traj_file, traj_file_fmt=traj_file_fmt, ).make(structure) @@ -121,7 +121,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): traj_from_file = traj_file_loader(traj_file) - assert len(traj_from_file) == nsteps + 1 + assert len(traj_from_file) == n_steps + 1 if traj_file_fmt == "pmg": assert all( @@ -132,7 +132,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): .get(key) ) for key in ("energy", "temperature", "forces", "velocities") - for idx in range(nsteps + 1) + for idx in range(n_steps + 1) ) elif traj_file_fmt == "ase": traj_as_dict = [ @@ -142,7 +142,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): "forces": traj_from_file[idx].get_forces(), "velocities": traj_from_file[idx].get_velocities(), } - for idx in range(1, nsteps + 1) + for idx in range(1, n_steps + 1) ] assert all( np.all( @@ -152,7 +152,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): .get(key) ) for key in ("energy", "temperature", "forces", "velocities") - for idx in range(1, nsteps + 1) + for idx in range(1, n_steps + 1) ) @@ -172,7 +172,7 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): job = CHGNetMDMaker( ensemble="nve", dynamics=dyn, - nsteps=50, + n_steps=50, traj_file=None, ).make(si_structure) @@ -200,13 +200,13 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_temp_schedule(ff_name, si_structure, clean_dir): - nsteps = 100 + n_steps = 100 temp_schedule = [300, 3000] structure = si_structure.to_conventional() * (2, 2, 2) job = _to_maker[ff_name]( - nsteps=nsteps, + n_steps=n_steps, traj_file=None, dynamics="nose-hoover", temperature=temp_schedule, @@ -225,14 +225,14 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_press_schedule(ff_name, si_structure, clean_dir): - nsteps = 100 + n_steps = 100 press_schedule = [0, 10] # kbar structure = si_structure.to_conventional() * (3, 3, 3) job = _to_maker[ff_name]( ensemble="npt", - nsteps=nsteps, + n_steps=n_steps, traj_file="md_traj.json.gz", traj_file_fmt="pmg", dynamics="nose-hoover", From 626b73bb4ddae7f7519d2753a1db081d07f2fb7a Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 1 Apr 2024 09:34:56 -0700 Subject: [PATCH 39/48] linting / pre-commit --- src/atomate2/forcefields/md.py | 6 ++++-- src/atomate2/forcefields/utils.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 000df2ce10..98643dfcc0 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -189,7 +189,9 @@ def _get_ensemble_schedule(self) -> None: isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 1 ): self.pschedule = np.interp( - np.arange(self.n_steps + 1), np.arange(len(self.pressure)), self.pressure + np.arange(self.n_steps + 1), + np.arange(len(self.pressure)), + self.pressure, ) elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: self.pschedule = interp1d( @@ -243,7 +245,7 @@ def make( # If a structure contains an isotope of hydrogen, set default `time_step` # to 0.5 fs, and 2 fs otherwise. has_h_isotope = any(element.Z == 1 for element in structure.composition) - self.time_step = (0.5 if has_h_isotope else 2.) * units.fs + self.time_step = (0.5 if has_h_isotope else 2.0) * units.fs initial_velocities = structure.site_properties.get("velocities") diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 2d9a69d781..38410e9fe7 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -139,7 +139,7 @@ def __init__(self, atoms: Atoms, store_md_outputs: bool = False) -> None: self._store_md_outputs = store_md_outputs # `self.{velocities,temperatures}` always initialized, - # but data is only stored / saved to trajectory for MD runs + # but data is only stored / saved to trajectory for MD runs self.velocities: list[np.ndarray] = [] self.temperatures: list[float] = [] From 9cee4122f67058e64dc18f1de9be369546a144e7 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 1 Apr 2024 09:39:13 -0700 Subject: [PATCH 40/48] revert default time_step units to fs --- src/atomate2/forcefields/md.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 98643dfcc0..7a21062283 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -66,9 +66,9 @@ class ForceFieldMDMaker(Maker): Perform MD with a forcefield. Note the the following units are consistent with the VASP MD implementation: - - Temperature is in Kelvin (TEBEG and TEEND) - - Time steps are in femtoseconds (POTIM) - - Pressure in kB (PSTRESS) + - `temperature` in Kelvin (TEBEG and TEEND) + - `time_step` in femtoseconds (POTIM) + - `pressure` in kB (PSTRESS) The default dynamics is Langevin NVT consistent with VASP MD, with the friction coefficient set to 10 ps^-1 (LANGEVIN_GAMMA). @@ -86,7 +86,7 @@ class ForceFieldMDMaker(Maker): force_field_name : str The name of the forcefield (for provenance) time_step : float | None = None. - The timestep of the MD run, specified in ASE's units. + The timestep of the MD run in fs. If `None`, defaults to 0.5 fs if a structure contains an isotope of hydrogen and 2 fs otherwise. n_steps : int = 1000 @@ -245,7 +245,7 @@ def make( # If a structure contains an isotope of hydrogen, set default `time_step` # to 0.5 fs, and 2 fs otherwise. has_h_isotope = any(element.Z == 1 for element in structure.composition) - self.time_step = (0.5 if has_h_isotope else 2.0) * units.fs + self.time_step = 0.5 if has_h_isotope else 2.0 initial_velocities = structure.site_properties.get("velocities") @@ -296,7 +296,7 @@ def make( md_runner = md_func( atoms=atoms, - timestep=self.time_step, + timestep=self.time_step * units.fs, **self.ase_md_kwargs, ) From da809dfd1fa31dfc60aa7c86de3c2518f48590b0 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 1 Apr 2024 10:03:15 -0700 Subject: [PATCH 41/48] add missing forcefields to version check in forcefield task document --- src/atomate2/forcefields/schemas.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 06004ec754..a912cb700c 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -268,6 +268,8 @@ def from_ase_compatible_result( MLFF.M3GNet: "matgl", MLFF.CHGNet: "chgnet", MLFF.MACE: "mace-torch", + MLFF.GAP: "quippy-ase", + MLFF.Nequip: "nequip", }.get(forcefield_name) # type: ignore[call-overload] if pkg_name: import importlib.metadata From 2fd3b6b177e41cd31c1d9b7dc5e55816a8abf749 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Mon, 1 Apr 2024 10:15:11 -0700 Subject: [PATCH 42/48] allow complex schur decomposition --- src/atomate2/forcefields/md.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 7a21062283..5b25481ba7 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -272,8 +272,8 @@ def make( if md_func is NPT: # Note that until md_func is instantiated, isinstance(md_func,NPT) is False # ASE NPT implementation requires upper triangular cell - u, _ = schur(atoms.get_cell(complete=True)) - atoms.set_cell(u, scale_atoms=True) + u, _ = schur(atoms.get_cell(complete=True), output="complex") + atoms.set_cell(u.real, scale_atoms=True) if initial_velocities: atoms.set_velocities(initial_velocities) From 3e7ced36bcdcce23b5ea52a733fb1a5877311d7c Mon Sep 17 00:00:00 2001 From: Janosh Riebesell Date: Mon, 1 Apr 2024 20:12:23 +0200 Subject: [PATCH 43/48] more snake case --- src/atomate2/common/schemas/phonons.py | 2 +- src/atomate2/forcefields/md.py | 50 ++++++++++++-------------- src/atomate2/forcefields/schemas.py | 18 +++++----- src/atomate2/forcefields/utils.py | 17 ++++----- tests/forcefields/test_md.py | 40 ++++++++++----------- 5 files changed, 58 insertions(+), 69 deletions(-) diff --git a/src/atomate2/common/schemas/phonons.py b/src/atomate2/common/schemas/phonons.py index 9c4b66af0d..ab954b5be0 100644 --- a/src/atomate2/common/schemas/phonons.py +++ b/src/atomate2/common/schemas/phonons.py @@ -128,7 +128,7 @@ class PhononJobDirs(BaseModel): None, description="Directory where optimization run was performed." ) taskdoc_run_job_dir: Optional[str] = Field( - None, description="Directory where taskdoc was generated." + None, description="Directory where task doc was generated." ) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 5b25481ba7..d3b951b907 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -29,7 +29,7 @@ from atomate2.forcefields import MLFF from atomate2.forcefields.jobs import forcefield_job -from atomate2.forcefields.schemas import ForcefieldResult, ForceFieldTaskDocument +from atomate2.forcefields.schemas import ForceFieldResult, ForceFieldTaskDocument from atomate2.forcefields.utils import ( TrajectoryObserver, ase_calculator, @@ -63,7 +63,7 @@ @dataclass class ForceFieldMDMaker(Maker): """ - Perform MD with a forcefield. + Perform MD with a force field. Note the the following units are consistent with the VASP MD implementation: - `temperature` in Kelvin (TEBEG and TEEND) @@ -127,10 +127,7 @@ class ForceFieldMDMaker(Maker): Whether to initialize the atomic velocities with zero angular momentum task_document_kwargs: dict Options to pass to the TaskDoc. Default choice - { - "store_trajectory": "partial", - "ionic_step_data": ("energy",), - } + {"store_trajectory": "partial", "ionic_step_data": ("energy",),} is consistent with atomate2.vasp.md.MDMaker """ @@ -153,23 +150,23 @@ class ForceFieldMDMaker(Maker): task_document_kwargs: dict = field( default_factory=lambda: { "store_trajectory": "partial", - "ionic_step_data": ("energy",), # energy is required in ionicsteps + "ionic_step_data": ("energy",), # energy is required in ionic steps } ) def _get_ensemble_schedule(self) -> None: if self.ensemble == "nve": - # Distable thermostat and barostat + # Disable thermostat and barostat self.temperature = np.nan self.pressure = np.nan - self.tschedule = np.full(self.n_steps + 1, self.temperature) - self.pschedule = np.full(self.n_steps + 1, self.pressure) + self.t_schedule = np.full(self.n_steps + 1, self.temperature) + self.p_schedule = np.full(self.n_steps + 1, self.pressure) return if isinstance(self.temperature, Sequence) or ( isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1 ): - self.tschedule = np.interp( + self.t_schedule = np.interp( np.arange(self.n_steps + 1), np.arange(len(self.temperature)), self.temperature, @@ -178,27 +175,27 @@ def _get_ensemble_schedule(self) -> None: # scalars, but in principle one quantity per atom could be specified by giving # an array. This is not implemented yet here. else: - self.tschedule = np.full(self.n_steps + 1, self.temperature) + self.t_schedule = np.full(self.n_steps + 1, self.temperature) if self.ensemble == "nvt": self.pressure = np.nan - self.pschedule = np.full(self.n_steps + 1, self.pressure) + self.p_schedule = np.full(self.n_steps + 1, self.pressure) return if isinstance(self.pressure, Sequence) or ( isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 1 ): - self.pschedule = np.interp( + self.p_schedule = np.interp( np.arange(self.n_steps + 1), np.arange(len(self.pressure)), self.pressure, ) elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: - self.pschedule = interp1d( + self.p_schedule = interp1d( np.arange(self.n_steps + 1), self.pressure, kind="linear" ) else: - self.pschedule = np.full(self.n_steps + 1, self.pressure) + self.p_schedule = np.full(self.n_steps + 1, self.pressure) def _get_ensemble_defaults(self) -> None: """Update ASE MD kwargs with defaults consistent with VASP MD.""" @@ -209,11 +206,11 @@ def _get_ensemble_defaults(self) -> None: self.ase_md_kwargs.pop("temperature_K", None) self.ase_md_kwargs.pop("externalstress", None) elif self.ensemble == "nvt": - self.ase_md_kwargs["temperature_K"] = self.tschedule[0] + self.ase_md_kwargs["temperature_K"] = self.t_schedule[0] self.ase_md_kwargs.pop("externalstress", None) elif self.ensemble == "npt": - self.ase_md_kwargs["temperature_K"] = self.tschedule[0] - self.ase_md_kwargs["externalstress"] = self.pschedule[0] * 1e3 * units.bar + self.ase_md_kwargs["temperature_K"] = self.t_schedule[0] + self.ase_md_kwargs["externalstress"] = self.p_schedule[0] * 1e3 * units.bar if isinstance(self.dynamics, str) and self.dynamics.lower() == "langevin": self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get( @@ -277,10 +274,10 @@ def make( if initial_velocities: atoms.set_velocities(initial_velocities) - elif not np.isnan(self.tschedule).any(): + elif not np.isnan(self.t_schedule).any(): MaxwellBoltzmannDistribution( atoms=atoms, - temperature_K=self.tschedule[0], + temperature_K=self.t_schedule[0], rng=np.random.default_rng(seed=self.mb_velocity_seed), ) if self.zero_linear_momentum: @@ -305,10 +302,10 @@ def make( def _callback(dyn: MolecularDynamics = md_runner) -> None: if self.ensemble == "nve": return - dyn.set_temperature(temperature_K=self.tschedule[dyn.nsteps]) + dyn.set_temperature(temperature_K=self.t_schedule[dyn.nsteps]) if self.ensemble == "nvt": return - dyn.set_stress(self.pschedule[dyn.nsteps] * 1e3 * units.bar) + dyn.set_stress(self.p_schedule[dyn.nsteps] * 1e3 * units.bar) md_runner.attach(_callback, interval=1) md_runner.run(steps=self.n_steps) @@ -322,7 +319,7 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: return ForceFieldTaskDocument.from_ase_compatible_result( self.force_field_name, - ForcefieldResult( + ForceFieldResult( final_structure=structure, trajectory=md_observer.to_pymatgen_trajectory(None), ), @@ -372,10 +369,7 @@ class GAPMDMaker(ForceFieldMDMaker): name: str = f"{MLFF.GAP} MD" force_field_name: str = f"{MLFF.GAP}" calculator_kwargs: dict = field( - default_factory=lambda: { - "args_str": "IP GAP", - "param_filename": "gap.xml", - } + default_factory=lambda: {"args_str": "IP GAP", "param_filename": "gap.xml"} ) diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index a912cb700c..28689f3878 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -15,7 +15,7 @@ from atomate2.forcefields import MLFF -class ForcefieldResult(dict): +class ForceFieldResult(dict): """Schema to store outputs in ForceFieldTaskDocument.""" def __init__(self, **kwargs: Any) -> None: @@ -28,8 +28,8 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) -class ForcefieldObject(ValueEnum): - """Types of forcefield data objects.""" +class ForceFieldObject(ValueEnum): + """Types of force field data objects.""" TRAJECTORY = "trajectory" @@ -127,10 +127,10 @@ class ForceFieldTaskDocument(StructureMetadata): None, description="Directory where the force field calculations are performed." ) - included_objects: Optional[list[ForcefieldObject]] = Field( - None, description="list of forcefield objects included with this task document" + included_objects: Optional[list[ForceFieldObject]] = Field( + None, description="list of force field objects included with this task document" ) - forcefield_objects: Optional[dict[ForcefieldObject, Any]] = Field( + forcefield_objects: Optional[dict[ForceFieldObject, Any]] = Field( None, description="Forcefield objects associated with this task" ) @@ -214,7 +214,7 @@ def from_ase_compatible_result( output_structure = result["final_structure"] final_energy = trajectory.frame_properties[-1]["energy"] - final_energy_per_atom = final_energy / input_structure.num_sites + final_energy_per_atom = final_energy / len(input_structure) final_forces = trajectory.frame_properties[-1]["forces"] final_stress = trajectory.frame_properties[-1]["stress"] @@ -246,12 +246,12 @@ def from_ase_compatible_result( ionic_steps.append(ionic_step) - forcefield_objects: dict[ForcefieldObject, Any] = {} + forcefield_objects: dict[ForceFieldObject, Any] = {} if store_trajectory != StoreTrajectoryOption.NO: # For VASP calculations, the PARTIAL trajectory option removes # electronic step info. There is no equivalent for forcefields, # so we just save the same info for FULL and PARTIAL options. - forcefield_objects[ForcefieldObject.TRAJECTORY] = trajectory # type: ignore[index] + forcefield_objects[ForceFieldObject.TRAJECTORY] = trajectory # type: ignore[index] output_doc = OutputDoc( structure=output_structure, diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 38410e9fe7..12cba25316 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -30,7 +30,7 @@ from pymatgen.io.ase import AseAtomsAdaptor from atomate2.forcefields import MLFF -from atomate2.forcefields.schemas import ForcefieldResult +from atomate2.forcefields.schemas import ForceFieldResult try: from ase.filters import FrechetCellFilter @@ -272,12 +272,7 @@ def as_dict(self) -> dict: traj_dict["magmoms"] = self.magmoms if self._store_md_outputs: - traj_dict.update( - { - "velocities": self.velocities, - "temperature": self.temperatures, - } - ) + traj_dict.update(velocities=self.velocities, temperature=self.temperatures) # sanitize dict for key in traj_dict: if all(isinstance(val, np.ndarray) for val in traj_dict[key]): @@ -328,7 +323,7 @@ def relax( verbose: bool = False, cell_filter: Filter = FrechetCellFilter, **kwargs, - ) -> ForcefieldResult: + ) -> ForceFieldResult: """ Relax the structure. @@ -373,10 +368,10 @@ def relax( struct = self.ase_adaptor.get_structure(atoms) traj = obs.to_pymatgen_trajectory(None) is_force_conv = all( - np.linalg.norm(traj.frame_properties[-1]["forces"][i]) < abs(fmax) - for i in range(struct.num_sites) + np.linalg.norm(traj.frame_properties[-1]["forces"][idx]) < abs(fmax) + for idx in range(len(struct)) ) - return ForcefieldResult( + return ForceFieldResult( final_structure=struct, trajectory=traj, is_force_converged=is_force_conv ) diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 1a498fd802..130af7eff7 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -20,7 +20,7 @@ NequipMDMaker, ) -_to_maker = { +name_to_maker = { "CHGNet": CHGNetMDMaker, "M3GNet": M3GNetMDMaker, "MACE": MACEMDMaker, @@ -59,7 +59,7 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle structure = unit_cell_structure.to_conventional() * (2, 2, 2) - job = _to_maker[ff_name]( + job = name_to_maker[ff_name]( n_steps=n_steps, traj_file="md_traj.json.gz", traj_file_fmt="pmg", @@ -67,32 +67,32 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle calculator_kwargs=calculator_kwargs, ).make(structure) response = run_locally(job, ensure_success=True) - taskdoc = response[next(iter(response))][1].output + task_doc = response[next(iter(response))][1].output # Check that energies are reasonably close to reference values - assert taskdoc.output.energy / len(structure) == pytest.approx( + assert task_doc.output.energy / len(structure) == pytest.approx( ref_energies_per_atom[ff_name], abs=0.1 ) # Check that we have the right number of MD steps # ASE logs the initial structure energy, and doesn't count this as an MD step - assert matcher.fit(taskdoc.output.ionic_steps[0].structure, structure) - assert len(taskdoc.output.ionic_steps) == n_steps + 1 + assert matcher.fit(task_doc.output.ionic_steps[0].structure, structure) + assert len(task_doc.output.ionic_steps) == n_steps + 1 # Check that the ionic steps have the expected physical properties assert all( key in step.model_dump() for key in ("energy", "forces", "stress", "structure") - for step in taskdoc.output.ionic_steps + for step in task_doc.output.ionic_steps ) # Check that the trajectory has expected physical properties - assert taskdoc.included_objects == ["trajectory"] - assert len(taskdoc.forcefield_objects["trajectory"]) == n_steps + 1 + assert task_doc.included_objects == ["trajectory"] + assert len(task_doc.forcefield_objects["trajectory"]) == n_steps + 1 assert all( key in step for key in ("energy", "forces", "stress", "velocities", "temperature") - for step in taskdoc.forcefield_objects["trajectory"].frame_properties + for step in task_doc.forcefield_objects["trajectory"].frame_properties ) @@ -111,13 +111,13 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): traj_file_loader = Trajectory structure = si_structure.to_conventional() * (2, 2, 2) - job = _to_maker[ff_name]( + job = name_to_maker[ff_name]( n_steps=n_steps, traj_file=traj_file, traj_file_fmt=traj_file_fmt, ).make(structure) response = run_locally(job, ensure_success=True) - taskdoc = response[next(iter(response))][1].output + task_doc = response[next(iter(response))][1].output traj_from_file = traj_file_loader(traj_file) @@ -127,7 +127,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): assert all( np.all( traj_from_file.frame_properties[idx][key] - == taskdoc.forcefield_objects["trajectory"] + == task_doc.forcefield_objects["trajectory"] .frame_properties[idx] .get(key) ) @@ -147,7 +147,7 @@ def test_traj_file(traj_file, si_structure, clean_dir, ff_name="CHGNet"): assert all( np.all( traj_as_dict[idx - 1][key] - == taskdoc.forcefield_objects["trajectory"] + == task_doc.forcefield_objects["trajectory"] .frame_properties[idx] .get(key) ) @@ -205,7 +205,7 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): structure = si_structure.to_conventional() * (2, 2, 2) - job = _to_maker[ff_name]( + job = name_to_maker[ff_name]( n_steps=n_steps, traj_file=None, dynamics="nose-hoover", @@ -213,11 +213,11 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): ase_md_kwargs=dict(ttime=50.0 * units.fs, pfactor=None), ).make(structure) response = run_locally(job, ensure_success=True) - taskdoc = response[next(iter(response))][1].output + task_doc = response[next(iter(response))][1].output temp_history = [ step["temperature"] - for step in taskdoc.forcefield_objects["trajectory"].frame_properties + for step in task_doc.forcefield_objects["trajectory"].frame_properties ] assert temp_history[-1] > temp_schedule[0] @@ -226,11 +226,11 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): @pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) def test_press_schedule(ff_name, si_structure, clean_dir): n_steps = 100 - press_schedule = [0, 10] # kbar + press_schedule = [0, 10] # kBar structure = si_structure.to_conventional() * (3, 3, 3) - job = _to_maker[ff_name]( + job = name_to_maker[ff_name]( ensemble="npt", n_steps=n_steps, traj_file="md_traj.json.gz", @@ -243,7 +243,7 @@ def test_press_schedule(ff_name, si_structure, clean_dir): ), ).make(structure) run_locally(job, ensure_success=True) - # taskdoc = response[next(iter(response))][1].output + # task_doc = response[next(iter(response))][1].output traj_from_file = loadfn("md_traj.json.gz") From a4c6bd533335efa491fc905dd575b4cc902fd651 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 1 Apr 2024 14:28:58 -0700 Subject: [PATCH 44/48] Decrease number of steps in tests for speed. Fix (?) temperature / pressure scheduling --- src/atomate2/forcefields/md.py | 26 ++++++++++++++------------ src/atomate2/forcefields/schemas.py | 24 +++++++++++++----------- src/atomate2/forcefields/utils.py | 6 +++--- tests/forcefields/test_jobs.py | 9 +++++++-- tests/forcefields/test_md.py | 26 ++++++++++++++++---------- 5 files changed, 53 insertions(+), 38 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index d3b951b907..dd28e9059b 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -29,7 +29,7 @@ from atomate2.forcefields import MLFF from atomate2.forcefields.jobs import forcefield_job -from atomate2.forcefields.schemas import ForceFieldResult, ForceFieldTaskDocument +from atomate2.forcefields.schemas import ForcefieldResult, ForceFieldTaskDocument from atomate2.forcefields.utils import ( TrajectoryObserver, ase_calculator, @@ -154,6 +154,16 @@ class ForceFieldMDMaker(Maker): } ) + @staticmethod + def _interpolate_quantity(values : Sequence | np.ndarray, n_pts: int) -> np.ndarray: + """ Utility function to interpolate temperature / pressure on a schedule. """ + n_vals = len(values) + return np.interp( + np.arange(n_pts + 1)*n_vals/(n_pts - 1), + np.arange(n_vals), + values, + ) + def _get_ensemble_schedule(self) -> None: if self.ensemble == "nve": # Disable thermostat and barostat @@ -166,11 +176,7 @@ def _get_ensemble_schedule(self) -> None: if isinstance(self.temperature, Sequence) or ( isinstance(self.temperature, np.ndarray) and self.temperature.ndim == 1 ): - self.t_schedule = np.interp( - np.arange(self.n_steps + 1), - np.arange(len(self.temperature)), - self.temperature, - ) + self.t_schedule = self._interpolate_quantity(self.temperature, self.n_steps) # NOTE: In ASE Langevin dynamics, the temperature are normally # scalars, but in principle one quantity per atom could be specified by giving # an array. This is not implemented yet here. @@ -185,11 +191,7 @@ def _get_ensemble_schedule(self) -> None: if isinstance(self.pressure, Sequence) or ( isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 1 ): - self.p_schedule = np.interp( - np.arange(self.n_steps + 1), - np.arange(len(self.pressure)), - self.pressure, - ) + self.p_schedule = self._interpolate_quantity(self.pressure, self.n_steps) elif isinstance(self.pressure, np.ndarray) and self.pressure.ndim == 4: self.p_schedule = interp1d( np.arange(self.n_steps + 1), self.pressure, kind="linear" @@ -319,7 +321,7 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None: return ForceFieldTaskDocument.from_ase_compatible_result( self.force_field_name, - ForceFieldResult( + ForcefieldResult( final_structure=structure, trajectory=md_observer.to_pymatgen_trajectory(None), ), diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 28689f3878..ace5683751 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -15,7 +15,7 @@ from atomate2.forcefields import MLFF -class ForceFieldResult(dict): +class ForcefieldResult(dict): """Schema to store outputs in ForceFieldTaskDocument.""" def __init__(self, **kwargs: Any) -> None: @@ -28,8 +28,8 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) -class ForceFieldObject(ValueEnum): - """Types of force field data objects.""" +class ForcefieldObject(ValueEnum): + """Types of forcefield data objects.""" TRAJECTORY = "trajectory" @@ -127,10 +127,10 @@ class ForceFieldTaskDocument(StructureMetadata): None, description="Directory where the force field calculations are performed." ) - included_objects: Optional[list[ForceFieldObject]] = Field( - None, description="list of force field objects included with this task document" + included_objects: Optional[list[ForcefieldObject]] = Field( + None, description="list of forcefield objects included with this task document" ) - forcefield_objects: Optional[dict[ForceFieldObject, Any]] = Field( + forcefield_objects: Optional[dict[ForcefieldObject, Any]] = Field( None, description="Forcefield objects associated with this task" ) @@ -214,7 +214,7 @@ def from_ase_compatible_result( output_structure = result["final_structure"] final_energy = trajectory.frame_properties[-1]["energy"] - final_energy_per_atom = final_energy / len(input_structure) + final_energy_per_atom = final_energy / input_structure.num_sites final_forces = trajectory.frame_properties[-1]["forces"] final_stress = trajectory.frame_properties[-1]["stress"] @@ -246,12 +246,12 @@ def from_ase_compatible_result( ionic_steps.append(ionic_step) - forcefield_objects: dict[ForceFieldObject, Any] = {} + forcefield_objects: dict[ForcefieldObject, Any] = {} if store_trajectory != StoreTrajectoryOption.NO: # For VASP calculations, the PARTIAL trajectory option removes # electronic step info. There is no equivalent for forcefields, # so we just save the same info for FULL and PARTIAL options. - forcefield_objects[ForceFieldObject.TRAJECTORY] = trajectory # type: ignore[index] + forcefield_objects[ForcefieldObject.TRAJECTORY] = trajectory # type: ignore[index] output_doc = OutputDoc( structure=output_structure, @@ -264,13 +264,15 @@ def from_ase_compatible_result( ) # map force field name to its package name - pkg_name = { + pkg_names = { MLFF.M3GNet: "matgl", MLFF.CHGNet: "chgnet", MLFF.MACE: "mace-torch", MLFF.GAP: "quippy-ase", MLFF.Nequip: "nequip", - }.get(forcefield_name) # type: ignore[call-overload] + } + pkg_names = {str(k): v for k,v in pkg_names.items()} + pkg_name = pkg_names.get(forcefield_name) if pkg_name: import importlib.metadata diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 12cba25316..4f6d1720fc 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -30,7 +30,7 @@ from pymatgen.io.ase import AseAtomsAdaptor from atomate2.forcefields import MLFF -from atomate2.forcefields.schemas import ForceFieldResult +from atomate2.forcefields.schemas import ForcefieldResult try: from ase.filters import FrechetCellFilter @@ -323,7 +323,7 @@ def relax( verbose: bool = False, cell_filter: Filter = FrechetCellFilter, **kwargs, - ) -> ForceFieldResult: + ) -> ForcefieldResult: """ Relax the structure. @@ -371,7 +371,7 @@ def relax( np.linalg.norm(traj.frame_properties[-1]["forces"][idx]) < abs(fmax) for idx in range(len(struct)) ) - return ForceFieldResult( + return ForcefieldResult( final_structure=struct, trajectory=traj, is_force_converged=is_force_conv ) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index dcdd28217b..fb837d7b6c 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -5,6 +5,8 @@ from pymatgen.core import Structure from pytest import approx, importorskip +from importlib.metadata import version as get_imported_version + from atomate2.forcefields.jobs import ( CHGNetRelaxMaker, CHGNetStaticMaker, @@ -36,6 +38,7 @@ def test_chgnet_static_maker(si_structure): assert output1.output.ionic_steps[-1].magmoms is None assert output1.output.n_steps == 1 + assert output1.forcefield_version == get_imported_version("chgnet") @pytest.mark.parametrize("relax_cell", [True, False]) def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): @@ -79,6 +82,7 @@ def test_m3gnet_static_maker(si_structure): assert output1.output.energy == approx(-10.8, abs=0.2) assert output1.output.n_steps == 1 + assert output1.forcefield_version == get_imported_version("matgl") def test_m3gnet_relax_maker(si_structure): # translate one atom to ensure a small number of relaxation steps are taken @@ -127,6 +131,7 @@ def test_mace_static_maker(si_structure: Structure, test_dir: Path, model): assert isinstance(output1, ForceFieldTaskDocument) assert output1.output.energy == approx(-0.068231, rel=1e-4) assert output1.output.n_steps == 1 + assert output1.forcefield_version == get_imported_version("mace-torch") @pytest.mark.parametrize("relax_cell", [True, False]) @@ -184,7 +189,7 @@ def test_gap_static_maker(si_structure: Structure, test_dir): assert isinstance(output1, ForceFieldTaskDocument) assert output1.output.energy == approx(-10.8523, rel=1e-4) assert output1.output.n_steps == 1 - + assert output1.forcefield_version == get_imported_version("quippy-ase") def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): importorskip("nequip") @@ -207,7 +212,7 @@ def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): assert isinstance(output1, ForceFieldTaskDocument) assert output1.output.energy == approx(-44.40017, rel=1e-4) assert output1.output.n_steps == 1 - + assert output1.forcefield_version == get_imported_version("nequip") @pytest.mark.parametrize("relax_cell", [True, False]) def test_nequip_relax_maker( diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 130af7eff7..0029317e42 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -191,16 +191,22 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): # ensure that output is consistent if molecular dynamics object is specified # as str or as MolecularDynamics object - assert all( - output["from_str"].output.__getattribute__(attr) - == output["from_dyn"].output.__getattribute__(attr) - for attr in ("energy", "forces", "stress", "structure") - ) - + for attr in ("energy","forces","stress","structure"): + vals = {k: output[k].output.__getattribute__(attr) for k in ("from_str","from_dyn",)} + if isinstance(vals["from_str"],float): + assert vals["from_str"] == pytest.approx(vals["from_dyn"]) + elif isinstance(vals["from_str"], Structure): + assert vals["from_str"] == vals["from_dyn"] + else: + assert all( + vals["from_str"][i][j] == pytest.approx(vals["from_dyn"][i][j]) + for i in range(len(vals["from_str"])) + for j in range(len(vals["from_str"][i])) + ) -@pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) +@pytest.mark.parametrize("ff_name", [ "CHGNet"]) def test_temp_schedule(ff_name, si_structure, clean_dir): - n_steps = 100 + n_steps = 50 temp_schedule = [300, 3000] structure = si_structure.to_conventional() * (2, 2, 2) @@ -223,9 +229,9 @@ def test_temp_schedule(ff_name, si_structure, clean_dir): assert temp_history[-1] > temp_schedule[0] -@pytest.mark.parametrize("ff_name", ["MACE", "CHGNet"]) +@pytest.mark.parametrize("ff_name", ["CHGNet"]) def test_press_schedule(ff_name, si_structure, clean_dir): - n_steps = 100 + n_steps = 20 press_schedule = [0, 10] # kBar structure = si_structure.to_conventional() * (3, 3, 3) From edbb2e7df61a013efec519fdb512631af61c41dc Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Mon, 1 Apr 2024 14:31:43 -0700 Subject: [PATCH 45/48] linting --- src/atomate2/forcefields/md.py | 6 +++--- src/atomate2/forcefields/schemas.py | 14 ++++++++------ tests/forcefields/test_jobs.py | 7 +++++-- tests/forcefields/test_md.py | 11 +++++++---- 4 files changed, 23 insertions(+), 15 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index dd28e9059b..9aa6a8d0d0 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -155,11 +155,11 @@ class ForceFieldMDMaker(Maker): ) @staticmethod - def _interpolate_quantity(values : Sequence | np.ndarray, n_pts: int) -> np.ndarray: - """ Utility function to interpolate temperature / pressure on a schedule. """ + def _interpolate_quantity(values: Sequence | np.ndarray, n_pts: int) -> np.ndarray: + """Interpolate temperature / pressure on a schedule.""" n_vals = len(values) return np.interp( - np.arange(n_pts + 1)*n_vals/(n_pts - 1), + np.arange(n_pts + 1) * n_vals / (n_pts - 1), np.arange(n_vals), values, ) diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index ace5683751..c4b7063176 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -265,13 +265,15 @@ def from_ase_compatible_result( # map force field name to its package name pkg_names = { - MLFF.M3GNet: "matgl", - MLFF.CHGNet: "chgnet", - MLFF.MACE: "mace-torch", - MLFF.GAP: "quippy-ase", - MLFF.Nequip: "nequip", + str(k): v + for k, v in { + MLFF.M3GNet: "matgl", + MLFF.CHGNet: "chgnet", + MLFF.MACE: "mace-torch", + MLFF.GAP: "quippy-ase", + MLFF.Nequip: "nequip", + }.items() } - pkg_names = {str(k): v for k,v in pkg_names.items()} pkg_name = pkg_names.get(forcefield_name) if pkg_name: import importlib.metadata diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index fb837d7b6c..ebd3837bf4 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -1,3 +1,4 @@ +from importlib.metadata import version as get_imported_version from pathlib import Path import pytest @@ -5,8 +6,6 @@ from pymatgen.core import Structure from pytest import approx, importorskip -from importlib.metadata import version as get_imported_version - from atomate2.forcefields.jobs import ( CHGNetRelaxMaker, CHGNetStaticMaker, @@ -40,6 +39,7 @@ def test_chgnet_static_maker(si_structure): assert output1.forcefield_version == get_imported_version("chgnet") + @pytest.mark.parametrize("relax_cell", [True, False]) def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): # translate one atom to ensure a small number of relaxation steps are taken @@ -84,6 +84,7 @@ def test_m3gnet_static_maker(si_structure): assert output1.forcefield_version == get_imported_version("matgl") + def test_m3gnet_relax_maker(si_structure): # translate one atom to ensure a small number of relaxation steps are taken si_structure.translate_sites(0, [0, 0, 0.1]) @@ -191,6 +192,7 @@ def test_gap_static_maker(si_structure: Structure, test_dir): assert output1.output.n_steps == 1 assert output1.forcefield_version == get_imported_version("quippy-ase") + def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): importorskip("nequip") task_doc_kwargs = {"ionic_step_data": ("structure", "energy")} @@ -214,6 +216,7 @@ def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): assert output1.output.n_steps == 1 assert output1.forcefield_version == get_imported_version("nequip") + @pytest.mark.parametrize("relax_cell", [True, False]) def test_nequip_relax_maker( sr_ti_o3_structure: Structure, test_dir: Path, relax_cell: bool diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 0029317e42..53e0e72d33 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -191,9 +191,11 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): # ensure that output is consistent if molecular dynamics object is specified # as str or as MolecularDynamics object - for attr in ("energy","forces","stress","structure"): - vals = {k: output[k].output.__getattribute__(attr) for k in ("from_str","from_dyn",)} - if isinstance(vals["from_str"],float): + for attr in ("energy", "forces", "stress", "structure"): + vals = { + k: output[k].output.__getattribute__(attr) for k in ("from_str", "from_dyn") + } + if isinstance(vals["from_str"], float): assert vals["from_str"] == pytest.approx(vals["from_dyn"]) elif isinstance(vals["from_str"], Structure): assert vals["from_str"] == vals["from_dyn"] @@ -204,7 +206,8 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): for j in range(len(vals["from_str"][i])) ) -@pytest.mark.parametrize("ff_name", [ "CHGNet"]) + +@pytest.mark.parametrize("ff_name", ["CHGNet"]) def test_temp_schedule(ff_name, si_structure, clean_dir): n_steps = 50 temp_schedule = [300, 3000] From 400bbc3ffb0c71acb6606c70792ef29e931f7033 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Tue, 2 Apr 2024 10:41:34 -0700 Subject: [PATCH 46/48] fix T/P schedule --- src/atomate2/forcefields/md.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 9aa6a8d0d0..655d6169d0 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -159,9 +159,9 @@ def _interpolate_quantity(values: Sequence | np.ndarray, n_pts: int) -> np.ndarr """Interpolate temperature / pressure on a schedule.""" n_vals = len(values) return np.interp( - np.arange(n_pts + 1) * n_vals / (n_pts - 1), - np.arange(n_vals), - values, + np.linspace(0, n_vals-1, n_pts+1), + np.linspace(0, n_vals-1, n_vals), + values ) def _get_ensemble_schedule(self) -> None: From 0a063739db96fb6afa17679b1c3c14a9c405cb0d Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Tue, 2 Apr 2024 11:02:40 -0700 Subject: [PATCH 47/48] linting --- src/atomate2/forcefields/md.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py index 655d6169d0..15cf1c850b 100644 --- a/src/atomate2/forcefields/md.py +++ b/src/atomate2/forcefields/md.py @@ -159,9 +159,9 @@ def _interpolate_quantity(values: Sequence | np.ndarray, n_pts: int) -> np.ndarr """Interpolate temperature / pressure on a schedule.""" n_vals = len(values) return np.interp( - np.linspace(0, n_vals-1, n_pts+1), - np.linspace(0, n_vals-1, n_vals), - values + np.linspace(0, n_vals - 1, n_pts + 1), + np.linspace(0, n_vals - 1, n_vals), + values, ) def _get_ensemble_schedule(self) -> None: From 4a88d7b188d9947c011f9f5552400cf7a3aa4dc4 Mon Sep 17 00:00:00 2001 From: esoteric-ephemera Date: Tue, 2 Apr 2024 15:08:40 -0700 Subject: [PATCH 48/48] replace dunder __getattribute__ --> getattr in forcefield test_md --- tests/forcefields/test_md.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 53e0e72d33..d8d6a1cc97 100644 --- a/tests/forcefields/test_md.py +++ b/tests/forcefields/test_md.py @@ -192,9 +192,7 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path): # ensure that output is consistent if molecular dynamics object is specified # as str or as MolecularDynamics object for attr in ("energy", "forces", "stress", "structure"): - vals = { - k: output[k].output.__getattribute__(attr) for k in ("from_str", "from_dyn") - } + vals = {k: getattr(output[k].output, attr) for k in ("from_str", "from_dyn")} if isinstance(vals["from_str"], float): assert vals["from_str"] == pytest.approx(vals["from_dyn"]) elif isinstance(vals["from_str"], Structure):