Merge pull request #489 from gpetretto/mdflow

Multi step MD flow

src/atomate2/vasp/flows/core.py

@@ -16,7 +16,10 @@
     RelaxMaker,
     StaticMaker,
 )
-from atomate2.vasp.sets.core import HSEBSSetGenerator, NonSCFSetGenerator
+from atomate2.vasp.sets.core import (
+    HSEBSSetGenerator,
+    NonSCFSetGenerator,
+)
 
 if TYPE_CHECKING:
     from pathlib import Path

src/atomate2/vasp/flows/md.py

@@ -0,0 +1,166 @@
+"""Flows for running molecular dynamics simulations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import Flow, Maker, OutputReference
+
+from atomate2.vasp.jobs.md import MDMaker, md_output
+from atomate2.vasp.sets.core import MDSetGenerator
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from pymatgen.core import Structure
+
+    from atomate2.vasp.jobs.base import BaseVaspMaker
+
+
+@dataclass
+class MultiMDMaker(Maker):
+    """
+    Maker to perform an MD run split in several steps.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    md_makers : .BaseVaspMaker
+        Maker to use to generate the first relaxation.
+    """
+
+    name: str = "multi md"
+    md_makers: list[BaseVaspMaker] = field(default_factory=lambda: [MDMaker()])
+
+    def make(
+        self,
+        structure: Structure,
+        prev_dir: str | Path | None = None,
+        prev_traj_ids: list[str] | None = None,
+    ) -> Flow:
+        """
+        Create a flow with several chained MD runs.
+
+        Parameters
+        ----------
+        structure : .Structure
+            A pymatgen structure object.
+        prev_dir : str or Path or None
+            A previous VASP calculation directory to copy output files from.
+        prev_traj_ids: a list of ids of job identifying previous steps of the
+            MD trajectory.
+
+        Returns
+        -------
+        Flow
+            A flow containing n_runs MD calculations.
+        """
+        md_job = None
+        md_jobs = []
+        for i, maker in enumerate(self.md_makers, 1):
+            if md_job is None:
+                md_structure = structure
+                md_prev_dir = prev_dir
+            else:
+                md_structure = md_job.output.structure
+                md_prev_dir = md_job.output.dir_name
+            md_job = maker.make(md_structure, prev_dir=md_prev_dir)
+            md_job.name += f" {i}"
+            md_jobs.append(md_job)
+
+        output_job = md_output(
+            structure=md_jobs[-1].output.structure,
+            vasp_dir=md_jobs[-1].output.dir_name,
+            traj_ids=[j.uuid for j in md_jobs],
+            prev_traj_ids=prev_traj_ids,
+        )
+        output_job.name = "molecular dynamics output"
+
+        md_jobs.append(output_job)
+
+        return Flow(md_jobs, output_job.output, name=self.name)
+
+    def restart_from_uuid(self, md_ref: str | OutputReference) -> Flow:
+        """
+        Create a flow from the output reference of another MultiMDMaker.
+
+        The last output will be used as the starting point and the reference to
+        all the previous steps will be included in the final document.
+
+        Parameters
+        ----------
+        md_ref: str or OutputReference
+            The reference to the output of another MultiMDMaker
+
+        Returns
+        -------
+            A flow containing n_runs MD calculations.
+        """
+        if isinstance(md_ref, str):
+            md_ref = OutputReference(md_ref)
+
+        return self.make(
+            structure=md_ref.structure,
+            prev_dir=md_ref.vasp_dir,
+            prev_traj_ids=md_ref.full_traj_ids,
+        )
+
+    @classmethod
+    def from_parameters(
+        cls,
+        nsteps: int,
+        time_step: float,
+        n_runs: int,
+        ensemble: str,
+        start_temp: float,
+        end_temp: float | None = None,
+        **kwargs,
+    ) -> MultiMDMaker:
+        """
+        Create an instance of the Maker based on the standard parameters.
+
+        Set values in the Flow maker, the Job Maker and the VaspInputGenerator,
+        using them to create the final instance of the Maker.
+
+        Parameters
+        ----------
+        nsteps: int
+            Number of time steps for simulations. The VASP `NSW` parameter.
+        time_step: float
+            The time step (in femtosecond) for the simulation. The VASP
+            `POTIM` parameter.
+        n_runs : int
+            Number of MD runs in the flow.
+        ensemble: str
+            Molecular dynamics ensemble to run. Options include `nvt`, `nve`, and `npt`.
+        start_temp: float
+            Starting temperature. The VASP `TEBEG` parameter.
+        end_temp: float or None
+            Final temperature. The VASP `TEEND` parameter. If None the same
+            as start_temp.
+        kwargs:
+            Other parameters passed
+
+        Returns
+        -------
+            A MultiMDMaker
+        """
+        if end_temp is None:
+            end_temp = start_temp
+        md_makers = []
+        start_temp_i = start_temp
+        increment = (end_temp - start_temp) / n_runs
+        for _ in range(n_runs):
+            end_temp_i = start_temp_i + increment
+            generator = MDSetGenerator(
+                nsteps=nsteps,
+                time_step=time_step,
+                ensemble=ensemble,
+                start_temp=start_temp_i,
+                end_temp=end_temp_i,
+            )
+            md_makers.append(MDMaker(input_set_generator=generator))
+            start_temp_i = end_temp_i
+        return cls(md_makers=md_makers, **kwargs)

src/atomate2/vasp/jobs/core.py

@@ -6,15 +6,6 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Literal
 
-from custodian.vasp.handlers import (
-    FrozenJobErrorHandler,
-    IncorrectSmearingHandler,
-    LargeSigmaHandler,
-    MeshSymmetryErrorHandler,
-    PositiveEnergyErrorHandler,
-    StdErrHandler,
-    VaspErrorHandler,
-)
 from pymatgen.alchemy.materials import TransformedStructure
 from pymatgen.alchemy.transmuters import StandardTransmuter
 
@@ -25,7 +16,6 @@
     HSERelaxSetGenerator,
     HSEStaticSetGenerator,
     HSETightRelaxSetGenerator,
-    MDSetGenerator,
     NonSCFSetGenerator,
     RelaxSetGenerator,
     StaticSetGenerator,
@@ -511,59 +501,3 @@ def make(
         self.write_additional_data.setdefault("transformations:json", tjson)
 
         return super().make.original(self, structure, prev_dir)
-
-
-@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": True}
-    )

src/atomate2/vasp/jobs/md.py

@@ -0,0 +1,128 @@
+"""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__)
+
+
+@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}
+    )
+
+
+@job(output_schema=MultiMDOutput)
+def md_output(
+    structure: Structure,
+    vasp_dir: str | Path,
+    traj_ids: list[str],
+    prev_traj_ids: list[str] | None,
+):
+    """
+    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)