diff --git a/src/atomate2/vasp/files.py b/src/atomate2/vasp/files.py index d24cefd5c7..217e53f85f 100644 --- a/src/atomate2/vasp/files.py +++ b/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: diff --git a/src/atomate2/vasp/flows/md.py b/src/atomate2/vasp/flows/md.py index a6109a2abe..72a0008888 100644 --- a/src/atomate2/vasp/flows/md.py +++ b/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) diff --git a/src/atomate2/vasp/jobs/md.py b/src/atomate2/vasp/jobs/md.py index aa3aae99e6..12a56b0de4 100644 --- a/src/atomate2/vasp/jobs/md.py +++ b/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, diff --git a/src/atomate2/vasp/sets/core.py b/src/atomate2/vasp/sets/core.py index 618ded50e5..25d0de06c6 100644 --- a/src/atomate2/vasp/sets/core.py +++ b/src/atomate2/vasp/sets/core.py @@ -707,6 +707,76 @@ def _get_ensemble_defaults(structure: Structure, ensemble: str) -> dict[str, Any raise ValueError(f"Expect {ensemble=} to be one of {supported}") from err +@dataclass +class MLMDSetGenerator(MDSetGenerator): + """ + Class to generate VASP molecular dynamics input sets using MLFF feature of Vasp. + + Parameters + ---------- + ml_mode + Operation mode for MLFF. Can be "train", "select", "refit" or "run" + **kwargs + Other keyword arguments that will be passed to :obj:`MDSetGenerator`. + """ + + ml_mode: str = "train" + + def __post_init__(self) -> None: + """Ensure validity of inputs.""" + super().__post_init__() + + supported_ml_modes = ("train", "select", "refit", "run") + if self.ml_mode not in supported_ml_modes: + raise ValueError( + f"Supported values for ml_mode are: {', '.join(supported_ml_modes)}" + ) + + def get_incar_updates( + self, + structure: Structure, + prev_incar: dict = None, + bandgap: float = None, + vasprun: Vasprun = None, + outcar: Outcar = None, + ) -> dict: + """ + Get updates to the INCAR for a molecular dynamics job. + + Parameters + ---------- + structure + A structure. + prev_incar + An incar from a previous calculation. + bandgap + The band gap. + vasprun + A vasprun from a previous calculation. + outcar + An outcar from a previous calculation. + + Returns + ------- + dict + A dictionary of updates to apply. + """ + updates = super().get_incar_updates( + structure=structure, + prev_incar=prev_incar, + bandgap=bandgap, + vasprun=vasprun, + outcar=outcar, + ) + updates.update( + { + "ML_LMLFF": True, + "ML_MODE": self.ml_mode, + } + ) + return updates + + @dataclass class LobsterTightStaticSetGenerator(LobsterSet): """ diff --git a/tests/vasp/sets/test_core.py b/tests/vasp/sets/test_core.py new file mode 100644 index 0000000000..6a43b56f8d --- /dev/null +++ b/tests/vasp/sets/test_core.py @@ -0,0 +1,25 @@ +import pytest +from pymatgen.core.structure import Structure + +from atomate2.vasp.sets.core import MLMDSetGenerator + + +def test_mlmd_generator(test_dir) -> None: + struct_gan = Structure.from_file(test_dir / "structures" / "GaN.cif") + + gen = MLMDSetGenerator() + incar_updates = gen.get_incar_updates(structure=struct_gan) + assert "ML_LMLFF" in incar_updates + assert incar_updates["ML_LMLFF"] is True + assert "ML_MODE" in incar_updates + assert incar_updates["ML_MODE"] == "train" + + gen = MLMDSetGenerator(ml_mode="run") + incar_updates = gen.get_incar_updates(structure=struct_gan) + assert "ML_MODE" in incar_updates + assert incar_updates["ML_MODE"] == "run" + + with pytest.raises( + ValueError, match=r"Supported values for ml_mode are: train, select, refit, run" + ): + MLMDSetGenerator(ml_mode="zztop")