Implement MD workflows

src/quacc/recipes/emt/md.py

@@ -0,0 +1,123 @@
+"""
+Molecular Dynamics recipes for EMT.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ase.calculators.emt import EMT
+from ase.md.velocitydistribution import (
+ MaxwellBoltzmannDistribution,
+ Stationary,
+ ZeroRotation,
+)
+
+from quacc import job
+from quacc.runners.ase import Runner
+from quacc.schemas.ase import summarize_md_run
+from quacc.utils.dicts import recursive_dict_merge
+
+if TYPE_CHECKING:
+ from typing import Any
+
+ from ase.atoms import Atoms
+
+ from quacc.schemas._aliases.ase import DynSchema
+
+
+@job
+def md_job(
+ atoms: Atoms,
+ maxwell_boltzmann_params: dict[str, Any] | None = None,
+ md_params: dict[str, Any] | None = None,
+ **calc_kwargs,
+) -> DynSchema:
+ """
+ Carry out a Molecular Dynamics calculation.
+
+ !!! Note "Units"
+
+ Quacc does not follow ASE standards for Molecular Dynamics units.
+
+ Quacc ALWAYS uses the following units:
+
+ - Time: femtoseconds (fs)
+ - Pressure: GPa
+ - Temperature: Kelvin (K)
+ - Compressibility: 1/GPa
+
+ !!! Note "Keywords"
+
+ Additionally, Quacc uses the keywords `fix_com` and `fix_rot` instead of
+ `fixcm` and `fixrot` to fix the center of mass and rotation, respectively.
+
+ As of March 2024, ASE still accept deprecated keywords such as:
+
+ - temperature instead of temperature_K
+ - pressure instead of pressure_au
+ - compressibility instead of compressibility_au
+ - dt instead of timestep
+
+ All of these keywords will be accepted by Quacc, but it is recommended to use the
+ new keywords to avoid confusion. Everything will always be converted to the Quacc
+ units mentioned above.
+
+ Shared keywords between various dynamics type such as `timestep` and `steps` should
+ be specified in the `md_params` dictionary. Keywords specific to the dynamics type
+ should be specified in a dictionary `dynamics_kwargs` inside `md_params`.
+
+ The dynamics type can be specified in `md_params` as `dynamics`.
+
+ Parameters
+ ----------
+ atoms
+ Atoms object
+ maxwell_boltzmann_params
+ Dictionary of custom kwargs for the initial temperature distribution. Set a value
+ to `quacc.Remove` to remove a pre-existing key entirely. For a list of available
+ keys, refer to [ase.md.velocitydistribution.MaxwellBoltzmannDistribution][].
+ Quacc has two additional parameters `fix_com` and `fix_rot` to fix the center of mass
+ and rotation from the initial velocity distribution. Defaults are
+ `{"temperature": None, "fix_com": True, "fix_rot": True}`. If `temperature` is set to
+ `None`, the initial temperature will not be set.
+ md_params
+ Dictionary of custom kwargs for the optimization process. Set a value
+ to `quacc.Remove` to remove a pre-existing key entirely. For a list of available
+ keys, refer to [quacc.runners.ase.Runner.run_md][].
+ **calc_kwargs
+ Custom kwargs for the EMT calculator. Set a value to
+ `quacc.Remove` to remove a pre-existing key entirely. For a list of available
+ keys, refer to the `ase.calculators.emt.EMT` calculator.
+
+ Returns
+ -------
+ DynSchema
+ Dictionary of results, specified in [quacc.schemas.ase.summarize_md_run][].
+ See the type-hint for the data structure.
+ """
+
+ maxwell_boltzmann_defaults = {"fix_com": True, "fix_rot": True, "temperature": None}
+ md_params = md_params or {}
+
+ maxwell_boltzmann_params = recursive_dict_merge(
+ maxwell_boltzmann_defaults, maxwell_boltzmann_params
+ )
+
+ initial_temperature = maxwell_boltzmann_params.pop("temperature", None)
+ fix_com = maxwell_boltzmann_params.pop("fix_com", False)
+ fix_rot = maxwell_boltzmann_params.pop("fix_rot", False)
+
+ if initial_temperature:
+ MaxwellBoltzmannDistribution(
+ atoms, temperature_K=initial_temperature, **maxwell_boltzmann_params
+ )
+ if fix_com:
+ Stationary(atoms)
+ if fix_rot:
+ ZeroRotation(atoms)
+
+ calc = EMT(**calc_kwargs)
+ dyn = Runner(atoms, calc).run_md(**md_params)
+
+ return summarize_md_run(dyn, additional_fields={"name": "EMT MD"})

src/quacc/runners/ase.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import logging
 import sys
 from importlib.util import find_spec
 from shutil import copy, copytree
@@ -11,6 +12,8 @@
 from ase.calculators import calculator
 from ase.filters import FrechetCellFilter
 from ase.io import Trajectory, read
+from ase.md.md import MolecularDynamics
+from ase.md.verlet import VelocityVerlet
 from ase.optimize import BFGS
 from ase.optimize.sciopt import SciPyOptimizer
 from ase.vibrations import Vibrations
@@ -22,6 +25,9 @@
 from quacc.runners._base import BaseRunner
 from quacc.runners.prep import terminate
 from quacc.utils.dicts import recursive_dict_merge
+from quacc.utils.units import convert_md_units
+
+LOGGER = logging.getLogger(__name__)
 
 has_sella = bool(find_spec("sella"))
 
@@ -32,7 +38,7 @@
 
  from ase.atoms import Atoms
  from ase.calculators.calculator import Calculator
- from ase.optimize.optimize import Optimizer
+ from ase.optimize.optimize import Dynamics
 
  from quacc.utils.files import Filenames, SourceDirectory
 
@@ -43,7 +49,7 @@ class OptParams(TypedDict, total=False):
 
  fmax: float
  max_steps: int
- optimizer: Optimizer # default = BFGS
+ optimizer: Dynamics = BFGS # default = BFGS
  optimizer_kwargs: OptimizerKwargs | None
  store_intermediate_results: bool
  fn_hook: Callable | None
@@ -167,12 +173,12 @@ def run_opt(
  relax_cell: bool = False,
  fmax: float | None = 0.01,
  max_steps: int = 1000,
- optimizer: Optimizer = BFGS,
+ optimizer: Dynamics = BFGS,
  optimizer_kwargs: OptimizerKwargs | None = None,
  store_intermediate_results: bool = False,
  fn_hook: Callable | None = None,
  run_kwargs: dict[str, Any] | None = None,
- ) -> Optimizer:
+ ) -> Dynamics:
  """
  This is a wrapper around the optimizers in ASE.
 
@@ -223,8 +229,12 @@ def run_opt(
  raise ValueError(msg)
 
  # Handle optimizer kwargs
- if issubclass(optimizer, SciPyOptimizer) or optimizer.__name__ == "IRC":
+ if (
+ issubclass(optimizer, (SciPyOptimizer, MolecularDynamics))
+ or optimizer.__name__ == "IRC"
+ ):
  # https://gitlab.com/ase/ase/-/issues/1476
+ # https://gitlab.com/ase/ase/-/merge_requests/3310
  optimizer_kwargs.pop("restart", None)
  if optimizer.__name__ == "Sella":
  self._set_sella_kwargs(optimizer_kwargs)
@@ -239,15 +249,17 @@ def run_opt(
  self.atoms = FrechetCellFilter(self.atoms)
 
  # Run optimization
+ full_run_kwargs = {"fmax": fmax, "steps": max_steps, **run_kwargs}
+ if issubclass(optimizer, MolecularDynamics):
+ full_run_kwargs.pop("fmax")
  try:
  with traj, optimizer(self.atoms, **optimizer_kwargs) as dyn:
- if issubclass(optimizer, SciPyOptimizer):
+ if issubclass(optimizer, (SciPyOptimizer, MolecularDynamics)):
  # https://gitlab.coms/ase/ase/-/issues/1475
- dyn.run(fmax=fmax, steps=max_steps, **run_kwargs)
+ # https://gitlab.com/ase/ase/-/issues/1497
+ dyn.run(**full_run_kwargs)
  else:
- for i, _ in enumerate(
- dyn.irun(fmax=fmax, steps=max_steps, **run_kwargs)
- ):
+ for i, _ in enumerate(dyn.irun(**full_run_kwargs)):
  if store_intermediate_results:
  self._copy_intermediate_files(
  i,
@@ -307,6 +319,50 @@ def run_vib(self, vib_kwargs: VibKwargs | None = None) -> Vibrations:
 
  return vib
 
+ def run_md(
+ self,
+ timestep: float = 1.0,
+ steps: int = 500,
+ dynamics: MolecularDynamics = VelocityVerlet,
+ dynamics_kwargs: dict[str, Any] | None = None,
+ ) -> MolecularDynamics:
+ """
+ Run an ASE-based MD in a scratch directory and copy the results back to
+ the original directory. This can be useful if file I/O is slow in the working
+ directory, so long as file transfer speeds are reasonable.
+
+ Parameters
+ ----------
+ timestep
+ The time step in femtoseconds.
+ steps
+ Maximum number of steps to run
+ dynamics
+ MolecularDynamics class to use, from `ase.md.md.MolecularDynamics`.
+ dynamics_kwargs
+ Dictionary of kwargs for the dynamics. Takes all valid kwargs for ASE
+ MolecularDynamics classes.
+
+ Returns
+ -------
+ Dymamics
+ The ASE MolecularDynamics object.
+ """
+
+ # Set defaults
+ dynamics_kwargs = dynamics_kwargs or {}
+ dynamics_kwargs["timestep"] = timestep
+ dynamics_kwargs["logfile"] = "-" if SETTINGS.DEBUG else self.tmpdir / "md.log"
+ dynamics_kwargs = self._fix_deprecated_md_params(dynamics_kwargs)
+ dynamics_kwargs = convert_md_units(dynamics_kwargs)
+
+ return self.run_opt(
+ fmax=None,
+ max_steps=steps,
+ optimizer=dynamics,
+ optimizer_kwargs=dynamics_kwargs,
+ )
+
  def _copy_intermediate_files(
  self, step_number: int, files_to_ignore: list[Path] | None = None
  ) -> None:
@@ -362,3 +418,80 @@ def _set_sella_kwargs(self, optimizer_kwargs: dict[str, Any]) -> None:
 
  if not self.atoms.pbc.any() and "internal" not in optimizer_kwargs:
  optimizer_kwargs["internal"] = True
+
+ @staticmethod
+ def _fix_deprecated_md_params(dynamics_kwargs: dict[str, Any]) -> None:
+ """
+ Helper function to handle deprecated MD parameters.
+
+ Parameters
+ ----------
+ dynamics_kwargs
+ Dictionary of keyword arguments for the molecular dynamics calculation.
+
+ Returns
+ -------
+ None
+ """
+
+ if "trajectory" in dynamics_kwargs:
+ msg = "Quacc does not support setting the `trajectory` kwarg."
+ raise ValueError(msg)
+
+ if "temperature" in dynamics_kwargs or "temp" in dynamics_kwargs:
+ LOGGER.warning(
+ "In quacc `temperature`, `temp` and `temperature_K` are"
+ " equivalent and will be interpreted in Kelvin."
+ )
+ dynamics_kwargs["temperature_K"] = dynamics_kwargs.pop(
+ "temperature", None
+ ) or dynamics_kwargs.pop("temp", None)
+
+ if "pressure" in dynamics_kwargs:
+ LOGGER.warning(
+ "In quacc `pressure` and `pressure_au` are equivalent and will be"
+ " interpreted in GPA."
+ )
+ dynamics_kwargs["pressure_au"] = dynamics_kwargs.pop("pressure")
+
+ if "pressure_au" in dynamics_kwargs:
+ LOGGER.warning(
+ "In quacc `pressure` and `pressure_au` are equivalent and will be"
+ " interpreted in GPA."
+ )
+
+ if "compressibility" in dynamics_kwargs:
+ LOGGER.warning(
+ "In quacc `compressibility` and `compressibility_au` are equivalent"
+ " and will be interpreted in 1/GPa."
+ )
+ dynamics_kwargs["compressibility_au"] = dynamics_kwargs.pop(
+ "compressibility"
+ )
+
+ if "compressibility_au" in dynamics_kwargs:
+ LOGGER.warning(
+ "In quacc `compressibility` and `compressibility_au` are equivalent"
+ " and will be interpreted in 1/GPa."
+ )
+
+ if "dt" in dynamics_kwargs:
+ LOGGER.warning(
+ "The `dt` kwarg is ASE deprecated and will"
+ " be interpreted as `timestep` in Quacc."
+ )
+ dynamics_kwargs["timestep"] = dynamics_kwargs.pop("dt")
+
+ if "fixcm" in dynamics_kwargs:
+ LOGGER.warning("`fixcm` is interpreted as `fix_com` in Quacc.")
+
+ if "fixrot" in dynamics_kwargs:
+ LOGGER.warning("`fixrot` is interpreted as `fix_rot` in Quacc.")
+
+ if "fix_com" in dynamics_kwargs:
+ dynamics_kwargs["fixcm"] = dynamics_kwargs.pop("fix_com")
+
+ if "fix_rot" in dynamics_kwargs:
+ dynamics_kwargs["fixrot"] = dynamics_kwargs.pop("fix_rot")
+
+ return dynamics_kwargs

src/quacc/schemas/_aliases/ase.py

@@ -23,6 +23,18 @@ class ParametersOpt(TypedDict):
  """Dictionary of parameters from Optimizer.todict()"""
 
 
+class ParametersDyn(TypedDict):
+ """Dictionary of parameters from Dynamics.todict()"""
+
+
+class TrajectoryLog(TypedDict):
+ """Dictionary of parameters related to the MD trajectory"""
+
+ kinetic_energy: float
+ temperature: float
+ time: float
+
+
 class RunSchema(AtomsSchema):
  """Schema for [quacc.schemas.ase.summarize_run][]"""
 
@@ -45,6 +57,16 @@ class OptSchema(RunSchema):
  trajectory_results: list[Results]
 
 
+class DynSchema(RunSchema):
+ """Schema for [quacc.schemas.ase.summarize_md_run][]"""
+
+ parameters_md: ParametersDyn # from Dynamics.todict()
+ nsteps: int
+ trajectory: list[Atoms]
+ trajectory_log: TrajectoryLog
+ trajectory_results: list[Results]
+
+
 class ParametersVib(TypedDict):
  delta: float
  direction: str