"""Module defining molecular dynamics jobs."""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from custodian.vasp.handlers import (
FrozenJobErrorHandler,
IncorrectSmearingHandler,
LargeSigmaHandler,
MeshSymmetryErrorHandler,
PositiveEnergyErrorHandler,
StdErrHandler,
VaspErrorHandler,
)
from emmet.core.vasp.calculation import StoreTrajectoryOption
from jobflow import Response, job
from atomate2.vasp.jobs.base import BaseVaspMaker
from atomate2.vasp.schemas.md import MultiMDOutput
from atomate2.vasp.sets.core import MDSetGenerator
if TYPE_CHECKING:
from pathlib import Path
from pymatgen.core import Structure
from atomate2.vasp.sets.base import VaspInputGenerator
logger = logging.getLogger(__name__)
[docs]
@dataclass
class MDMaker(BaseVaspMaker):
"""
Maker to create VASP molecular dynamics jobs.
Parameters
----------
name : str
The job name.
input_set_generator : .VaspInputSetGenerator
A generator used to make the input set.
write_input_set_kwargs : dict
Keyword arguments that will get passed to :obj:`.write_vasp_input_set`.
copy_vasp_kwargs : dict
Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`.
run_vasp_kwargs : dict
Keyword arguments that will get passed to :obj:`.run_vasp`.
task_document_kwargs : dict
Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`.
stop_children_kwargs : dict
Keyword arguments that will get passed to :obj:`.should_stop_children`.
write_additional_data : dict
Additional data to write to the current directory. Given as a dict of
{filename: data}. Note that if using FireWorks, dictionary keys cannot contain
the "." character which is typically used to denote file extensions. To avoid
this, use the ":" character, which will automatically be converted to ".". E.g.
``{"my_file:txt": "contents of the file"}``.
"""
name: str = "molecular dynamics"
input_set_generator: VaspInputGenerator = field(default_factory=MDSetGenerator)
# Explicitly pass the handlers to not use the default ones. Some default handlers
# such as PotimErrorHandler do not apply to MD runs.
run_vasp_kwargs: dict = field(
default_factory=lambda: {
"handlers": (
VaspErrorHandler(),
MeshSymmetryErrorHandler(),
PositiveEnergyErrorHandler(),
FrozenJobErrorHandler(),
StdErrHandler(),
LargeSigmaHandler(),
IncorrectSmearingHandler(),
)
}
)
# Store ionic steps info in a pymatgen Trajectory object instead of in the output
# document.
task_document_kwargs: dict = field(
default_factory=lambda: {"store_trajectory": StoreTrajectoryOption.PARTIAL}
)
[docs]
@job(output_schema=MultiMDOutput)
def md_output(
structure: Structure,
vasp_dir: str | Path,
traj_ids: list[str],
prev_traj_ids: list[str] | None,
) -> Response:
"""
Collect output references of a multistep MD flow.
Parameters
----------
structure: .Structure
The final structure to be stored.
vasp_dir: str or Path
The path to the folder containing the last calculation of a MultiMDMaker.
traj_ids: list of str
List of the uuids of the jobs that will compose the trajectory.
prev_traj_ids: list of str
List of the uuids of the jobs coming from previous flow that will be
added to the trajectory.
Returns
-------
The output dictionary.
"""
full_traj_ids = list(traj_ids)
if prev_traj_ids:
full_traj_ids = prev_traj_ids + full_traj_ids
output = MultiMDOutput.from_structure(
structure=structure,
meta_structure=structure,
vasp_dir=str(vasp_dir),
traj_ids=traj_ids,
full_traj_ids=full_traj_ids,
)
return Response(output=output)