WIP: Vasp ml

src/atomate2/vasp/files.py

@@ -27,7 +27,7 @@
 def copy_vasp_outputs(
     src_dir: Path | str,
     src_host: str | None = None,
-    additional_vasp_files: Sequence[str] = (),
+    additional_vasp_files: Sequence[str] | dict[str, str] = (),
     contcar_to_poscar: bool = True,
     force_overwrite: bool | str = False,
     file_client: FileClient | None = None,
@@ -49,8 +49,12 @@ def copy_vasp_outputs(
         either "username@remote_host" or just "remote_host" in which case the username
         will be inferred from the current user. If ``None``, the local filesystem will
         be used as the source.
-    additional_vasp_files : list of str
-        Additional files to copy, e.g. ["CHGCAR", "WAVECAR"].
+    additional_vasp_files : list of str or dict
+        Additional files to copy, e.g. ["CHGCAR", "WAVECAR"]. If provided as a dict,
+        the key is the original file name in the src_dir, and the value is the name
+        of the file to which it should be renamed in the current directory, e.g.
+        {"ML_ABN": "ML_AB"} will copy the ML_ABN to the current directory and then
+        rename it to ML_AB.
     contcar_to_poscar : bool
         Move CONTCAR to POSCAR (original POSCAR is not copied).
     force_overwrite : bool or str
@@ -71,7 +75,11 @@ def copy_vasp_outputs(
     directory_listing = file_client.listdir(src_dir, host=src_host)
 
     # find required files
-    files = ("INCAR", "OUTCAR", "CONTCAR", "vasprun.xml", *additional_vasp_files)
+    if isinstance(additional_vasp_files, dict):
+        additional_files = list(additional_vasp_files.keys())
+    else:
+        additional_files = list(additional_vasp_files)
+    files = ("INCAR", "OUTCAR", "CONTCAR", "vasprun.xml", *additional_files)
     required_files = [get_zfile(directory_listing, r + relax_ext) for r in files]
 
     # find optional files; do not fail if KPOINTS is missing, this might be KSPACING
@@ -114,6 +122,9 @@ def copy_vasp_outputs(
     if contcar_to_poscar:
         rename_files({"CONTCAR": "POSCAR"}, file_client=file_client)
 
+    if isinstance(additional_vasp_files, dict):
+        rename_files(additional_vasp_files, file_client=file_client)
+
     logger.info("Finished copying inputs")
 
 
@@ -188,7 +199,9 @@ def write_vasp_input_set(
     """
     prev_dir = "." if from_prev else None
     vis = input_set_generator.get_input_set(
-        structure, prev_dir=prev_dir, potcar_spec=potcar_spec
+        structure,
+        prev_dir=prev_dir,
+        potcar_spec=potcar_spec,
     )
 
     if apply_incar_updates:

src/atomate2/vasp/flows/md.py

@@ -7,7 +7,7 @@
 
 from jobflow import Flow, Maker, OutputReference
 
-from atomate2.vasp.jobs.md import MDMaker, md_output
+from atomate2.vasp.jobs.md import MDMaker, MLMDMaker, md_output
 from atomate2.vasp.sets.core import MDSetGenerator
 
 if TYPE_CHECKING:
@@ -106,6 +106,32 @@ def restart_from_uuid(self, md_ref: str | OutputReference) -> Flow:
             prev_traj_ids=md_ref.full_traj_ids,
         )
 
+    @staticmethod
+    def _split_md(
+        nsteps: int, n_runs: int, start_temp: float, end_temp: float | None = None
+    ) -> list:
+        """Get a balanced set of N runs for the required number of steps."""
+        if end_temp is None:
+            end_temp = start_temp
+        # Split steps into balanced groups
+        nsteps_run = nsteps // n_runs
+        remaining = nsteps - n_runs * nsteps_run
+        nsteps_runs = [nsteps_run] * n_runs
+        for ii in range(remaining):
+            nsteps_runs[ii] += 1
+
+        # Adapt start and end temperatures to the number of steps in each run
+        delta_temp = end_temp - start_temp
+        start_temp_runs = []
+        end_temp_runs = []
+        prevrun_end_temp = start_temp
+        for irun in range(n_runs):
+            start_temp_runs.append(prevrun_end_temp)
+            prevrun_end_temp += delta_temp / nsteps * nsteps_runs[irun]
+            end_temp_runs.append(prevrun_end_temp)
+
+        return list(zip(nsteps_runs, start_temp_runs, end_temp_runs, strict=False))
+
     @classmethod
     def from_parameters(
         cls,
@@ -145,20 +171,139 @@ def from_parameters(
         -------
             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
+        for nsteps_run, start_temp_run, end_temp_run in cls._split_md(
+            nsteps=nsteps, n_runs=n_runs, start_temp=start_temp, end_temp=end_temp
+        ):
             generator = MDSetGenerator(
-                nsteps=nsteps,
+                nsteps=nsteps_run,
                 time_step=time_step,
                 ensemble=ensemble,
-                start_temp=start_temp_i,
-                end_temp=end_temp_i,
+                start_temp=start_temp_run,
+                end_temp=end_temp_run,
             )
             md_makers.append(MDMaker(input_set_generator=generator))
-            start_temp_i = end_temp_i
+        return cls(md_makers=md_makers, **kwargs)
+
+    @classmethod
+    def onthefly_mlff(
+        cls,
+        nsteps: int,
+        time_step: float,
+        n_runs: int,
+        ensemble: str,
+        start_temp: float,
+        end_temp: float | None = None,
+        **kwargs,
+    ) -> MultiMDMaker:
+        """
+        Create an on-the-fly MLFF-based MultiMDMaker 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
+        """
+        md_makers = []
+        for nsteps_run, start_temp_run, end_temp_run in cls._split_md(
+            nsteps=nsteps, n_runs=n_runs, start_temp=start_temp, end_temp=end_temp
+        ):
+            md_makers.append(
+                MLMDMaker.train(
+                    generator_kwargs={
+                        "nsteps": nsteps_run,
+                        "time_step": time_step,
+                        "ensemble": ensemble,
+                        "start_temp": start_temp_run,
+                        "end_temp": end_temp_run,
+                    },
+                )
+            )
+        return cls(md_makers=md_makers, **kwargs)
+
+    @classmethod
+    def production_run_mlff(
+        cls,
+        nsteps: int,
+        time_step: float,
+        ensemble: str,
+        start_temp: float,
+        end_temp: float | None = None,
+        n_runs: int = 1,
+        refit: bool = True,
+        **kwargs,
+    ) -> MultiMDMaker:
+        """
+        Create an on-the-fly MLFF-based MultiMDMaker 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.
+
+        .. Note::
+            This can only work with a previous directory where training of .
+
+        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.
+        refit: bool
+            Whether to refit the ML force field based on existing ML_AB file with
+            reference configurations.
+        kwargs:
+            Other parameters passed
+
+        Returns
+        -------
+            A MultiMDMaker
+        """
+        md_makers = []
+        if refit:
+            md_makers.append(MLMDMaker.refit())
+        for nsteps_run, start_temp_run, end_temp_run in cls._split_md(
+            nsteps=nsteps, n_runs=n_runs, start_temp=start_temp, end_temp=end_temp
+        ):
+            md_makers.append(
+                MLMDMaker.run(
+                    generator_kwargs={
+                        "nsteps": nsteps_run,
+                        "time_step": time_step,
+                        "ensemble": ensemble,
+                        "start_temp": start_temp_run,
+                        "end_temp": end_temp_run,
+                    },
+                )
+            )
         return cls(md_makers=md_makers, **kwargs)

src/atomate2/vasp/jobs/md.py

@@ -20,7 +20,7 @@
 
 from atomate2.vasp.jobs.base import BaseVaspMaker
 from atomate2.vasp.schemas.md import MultiMDOutput
-from atomate2.vasp.sets.core import MDSetGenerator
+from atomate2.vasp.sets.core import MDSetGenerator, MLMDSetGenerator
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -85,10 +85,63 @@ class MDMaker(BaseVaspMaker):
     # 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}
+        default_factory=lambda: {
+            "store_trajectory": StoreTrajectoryOption.PARTIAL,
+            "vasprun_kwargs": {"parse_dos": False, "parse_eigen": False},
+        }
+    )
+    write_input_set_kwargs: dict = field(
+        default_factory=lambda: {"get_previous_bandgap": False}
     )
 
 
+@dataclass
+class MLMDMaker(MDMaker):
+    """Maker to create VASP molecular dynamics jobs using MLFF feature."""
+
+    name: str = "MLFF molecular dynamics"
+
+    input_set_generator: VaspInputGenerator = field(default_factory=MLMDSetGenerator)
+
+    @classmethod
+    def train(cls, generator_kwargs: dict | None = None, **kwargs) -> MLMDMaker:
+        """Train."""
+        generator_kwargs = generator_kwargs or {}
+        return cls(
+            name="MLFF MD train",
+            input_set_generator=MLMDSetGenerator(ml_mode="train", **generator_kwargs),
+            copy_vasp_kwargs={"additional_vasp_files": {"ML_ABN": "ML_AB"}},
+            **kwargs,
+        )
+
+    @classmethod
+    def select(cls, generator_kwargs: dict | None = None, **kwargs) -> MLMDMaker:
+        """Select."""
+        raise NotImplementedError
+
+    @classmethod
+    def refit(cls, generator_kwargs: dict | None = None, **kwargs) -> MLMDMaker:
+        """Refit."""
+        generator_kwargs = generator_kwargs or {}
+        return cls(
+            name="MLFF refit",
+            input_set_generator=MLMDSetGenerator(ml_mode="refit", **generator_kwargs),
+            copy_vasp_kwargs={"additional_vasp_files": {"ML_ABN": "ML_AB"}},
+            **kwargs,
+        )
+
+    @classmethod
+    def run(cls, generator_kwargs: dict | None = None, **kwargs) -> MLMDMaker:
+        """Run."""
+        generator_kwargs = generator_kwargs or {}
+        return cls(
+            name="MLFF MD run",
+            input_set_generator=MLMDSetGenerator(ml_mode="run", **generator_kwargs),
+            copy_vasp_kwargs={"additional_vasp_files": {"ML_FFN": "ML_FF"}},
+            **kwargs,
+        )
+
+
 @job(output_schema=MultiMDOutput)
 def md_output(
     structure: Structure,