Multi step MD flow

src/atomate2/vasp/flows/core.py

@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING
 
 from emmet.core.vasp.calculation import VaspObject
-from jobflow import Flow, Maker
+from jobflow import Flow, Maker, OutputReference
 
 from atomate2.vasp.jobs.core import (
     HSEBSMaker,
@@ -16,6 +16,7 @@
     RelaxMaker,
     StaticMaker,
 )
+from atomate2.vasp.jobs.md import MDMaker, md_output
 from atomate2.vasp.sets.core import HSEBSSetGenerator, NonSCFSetGenerator
 
 if TYPE_CHECKING:
@@ -506,3 +507,94 @@ def make(self, structure: Structure, prev_vasp_dir: str | Path | None = None):
             static_job.output.structure, prev_vasp_dir=static_job.output.dir_name
         )
         return Flow([static_job, bs_job], bs_job.output, name=self.name)
+
+
+@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_maker : .BaseVaspMaker
+        Maker to use to generate the first relaxation.
+    """
+
+    name: str = "multi md"
+    md_maker: BaseVaspMaker = field(default_factory=MDMaker)
+    n_runs: int = 5
+
+    def make(
+        self,
+        structure: Structure,
+        prev_vasp_dir: str | Path | None = None,
+        prev_traj_ids: list[str] | None = None,
+    ):
+        """
+        Create a flow with several chained MD runs.
+
+        Parameters
+        ----------
+        structure : .Structure
+            A pymatgen structure object.
+        prev_vasp_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 in range(1, self.n_runs + 1):
+            if md_job is None:
+                md_structure = structure
+                md_prev_vasp_dir = prev_vasp_dir
+            else:
+                md_structure = md_job.output.structure
+                md_prev_vasp_dir = md_job.output.dir_name
+            md_job = self.md_maker.make(md_structure, prev_vasp_dir=md_prev_vasp_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):
+        """
+        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_vasp_dir=md_ref.vasp_dir,
+            prev_traj_ids=md_ref.full_traj_ids,
+        )

src/atomate2/vasp/jobs/core.py

@@ -6,15 +6,6 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
-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,
@@ -53,7 +43,6 @@
     "TightRelaxMaker",
     "HSETightRelaxMaker",
     "TransmuterMaker",
-    "MDMaker",
 ]
 
 
@@ -524,59 +513,3 @@ def make(
         self.write_additional_data.setdefault("transformations:json", tjson)
 
         return super().make.original(self, structure, prev_vasp_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,130 @@
+"""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 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__)
+
+
+__all__ = ["MDMaker"]
+
+
+@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}
+    )
+
+
+@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)

src/atomate2/vasp/schemas/md.py

@@ -0,0 +1,18 @@
+"""Schemas for MD documents."""
+
+from typing import List
+
+from emmet.core.structure import StructureMetadata
+from pydantic import Field
+from pymatgen.core import Structure
+
+
+class MultiMDOutput(StructureMetadata):
+    """Output of a MultiMD Flow."""
+
+    structure: Structure = Field("Final structure of the last step of the flow")
+    vasp_dir: str = Field("Path to the last vasp folder of the flow")
+    traj_ids: List[str] = Field("List of uuids of the MD calculations in the flow")
+    full_traj_ids: List[str] = Field(
+        "List of uuids of the MD calculations in the flow and in previous linked flows"
+    )

src/atomate2/vasp/sets/base.py

@@ -422,10 +422,19 @@ def get_input_set(  # type: ignore
             bandgap=bandgap,
             ispin=ispin,
         )
+        site_properties = structure.site_properties
+        poscar = Poscar(
+            structure,
+            velocities=site_properties.get("velocities"),
+            predictor_corrector=site_properties.get("predictor_corrector"),
+            predictor_corrector_preamble=structure.properties.get(
+                "predictor_corrector_preamble"
+            ),
+        )
         return VaspInputSet(
             incar=incar,
             kpoints=kpoints,
-            poscar=Poscar(structure),
+            poscar=poscar,
             potcar=self._get_potcar(structure, potcar_spec=potcar_spec),
         )