Adding MD input set to FHI-aims (#3896)

* Added MD input set and generator for aims

* Added MD input set and generator for aims

* Conditionally add velocities to site properties

* pre-commit auto-fixes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/pymatgen/io/aims/inputs.py

@@ -60,7 +60,7 @@ def from_str(cls, contents: str) -> Self:
         ]
 
         species, coords, is_frac, lattice_vectors = [], [], [], []
-        charges_dct, moments_dct = {}, {}
+        charges_dct, moments_dct, velocities_dct = {}, {}, {}
 
         for line in content_lines:
             inp = line.split()
@@ -74,6 +74,8 @@ def from_str(cls, contents: str) -> Self:
                 moments_dct[len(coords) - 1] = float(inp[1])
             if inp[0] == "initial_charge":
                 charges_dct[len(coords) - 1] = float(inp[1])
+            if inp[0] == "velocity":
+                velocities_dct[len(coords) - 1] = [float(x) for x in inp[1:]]
 
         charge = np.zeros(len(coords))
         for key, val in charges_dct.items():
@@ -83,6 +85,10 @@ def from_str(cls, contents: str) -> Self:
         for key, val in moments_dct.items():
             magmom[key] = val
 
+        velocity: list[None | list[float]] = [None for _ in coords]
+        for key, v in velocities_dct.items():
+            velocity[key] = v
+
         if len(lattice_vectors) == 3:
             lattice = Lattice(lattice_vectors)
             for cc in range(len(coords)):
@@ -96,6 +102,9 @@ def from_str(cls, contents: str) -> Self:
             raise ValueError("Incorrect number of lattice vectors passed.")
 
         site_props = {"magmom": magmom, "charge": charge}
+        if velocities_dct:
+            site_props["velocity"] = velocity
+
         if lattice is None:
             structure = Molecule(species, coords, np.sum(charge), site_properties=site_props)
         else:
@@ -137,13 +146,17 @@ def from_structure(cls, structure: Structure | Molecule) -> Self:
 
         charges = structure.site_properties.get("charge", np.zeros(len(structure.species)))
         magmoms = structure.site_properties.get("magmom", np.zeros(len(structure.species)))
-        for species, coord, charge, magmom in zip(structure.species, structure.cart_coords, charges, magmoms):
+        velocities = structure.site_properties.get("velocity", [None for _ in structure.species])
+        for species, coord, charge, magmom, v in zip(
+            structure.species, structure.cart_coords, charges, magmoms, velocities
+        ):
             content_lines.append(f"atom {coord[0]: .12e} {coord[1]: .12e} {coord[2]: .12e} {species}")
             if charge != 0:
                 content_lines.append(f"     initial_charge {charge:.12e}")
             if magmom != 0:
                 content_lines.append(f"     initial_moment {magmom:.12e}")
-
+            if v is not None and any(v_i != 0.0 for v_i in v):
+                content_lines.append(f"     velocity   {'  '.join([f'{v_i:.12e}' for v_i in v])}")
         return cls(_content="\n".join(content_lines), _structure=structure)
 
     @property

src/pymatgen/io/aims/sets/core.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from pymatgen.core import Structure
@@ -14,6 +14,13 @@
     from pymatgen.core import Molecule
 
 
+_valid_dynamics: dict[str, tuple[str, ...]] = {
+    "nve": ("", "4th_order", "damped"),
+    "nvt": ("andersen", "berendsen", "parrinello", "nose-hoover"),
+    "gle": ("thermostat",),
+}
+
+
 @dataclass
 class StaticSetGenerator(AimsInputGenerator):
     """Common class for ground-state generators.
@@ -97,3 +104,95 @@ def get_parameter_updates(self, structure: Structure | Molecule, prev_parameters
             dict: The updated for the parameters for the output section of FHI-aims
         """
         return {"use_pimd_wrapper": (self.host, self.port)}
+
+
+@dataclass
+class MDSetGenerator(AimsInputGenerator):
+    """
+    A class for generating FHI-aims input sets for molecular dynamics calculations.
+
+    Parameters
+    ----------
+    ensemble
+        Molecular dynamics ensemble to run. Options include `nvt`, `nve`, and `gle`.
+        Default: `nve`
+    ensemble_specs
+        A dictionary containing the specifications of the molecular dynamics ensemble.
+        Valid keys are `type` (the ensemble type, valid types are defined in
+        `_valid_dynamics` dict), and `parameter` - the control parameter for the thermostat
+        (not used for `nve` and `nve_4th_order`).
+    temp
+        Thermostat temperature. Default: None
+    time
+        Simulation time (in picoseconds). Negative value stands for indefinite run.
+        Default: 5 ps
+    time_step
+        The time step (in picoseconds) for the simulation. default: 1 fs
+    **kwargs
+        Other keyword arguments that will be passed to :obj:`AimsInputGenerator`.
+    """
+
+    calc_type: str = "md"
+    ensemble: str = "nve"
+    ensemble_specs: dict[str, Any] = field(default_factory=dict)
+    temp: float | None = None
+    time: float = 5.0
+    time_step: float = 0.001
+    init_velocities: bool = True
+
+    def get_parameter_updates(self, structure: Structure | Molecule, prev_parameters: dict[str, Any]) -> dict:
+        """Get the parameter updates for the calculation.
+
+        Parameters
+        ----------
+        structure (Structure or Molecule):
+            The structure to calculate the bands for
+        prev_parameters (Dict[str, Any]):
+            The previous parameters
+
+        Returns
+        -------
+        dict
+            A dictionary of updates to apply.
+        """
+        updates: dict[str, Any] = {"MD_run": [self.time], "MD_time_step": self.time_step}
+
+        # check for ensemble type validity
+        default_ensemble_types = {"nve": "", "nvt": "parrinello", "gle": "thermostat"}
+        if self.ensemble not in _valid_dynamics:
+            raise ValueError(f"Ensemble {self.ensemble} not valid")
+        ensemble_type = self.ensemble_specs.get("type", default_ensemble_types[self.ensemble])
+        if ensemble_type not in _valid_dynamics[self.ensemble]:
+            raise ValueError(
+                f"Type {ensemble_type} is not valid for {self.ensemble} ensemble. "
+                f"Valid types are: {' ,'.join(_valid_dynamics[self.ensemble])}"
+            )
+        ensemble_name = f"{self.ensemble.upper()}_{ensemble_type}" if ensemble_type else self.ensemble.upper()
+        updates["MD_run"].append(ensemble_name)
+
+        # add temperature
+        if self.ensemble == "nve":
+            if not self.init_velocities and "velocity" not in structure.site_properties:
+                raise ValueError("Velocities must be initialized for NVE ensemble")
+        else:
+            if self.temp is None:
+                raise ValueError(f"Temperature must be set for {ensemble_name} ensemble")
+            updates["MD_run"].append(self.temp)
+
+        # check for ensemble control parameter
+        ensemble_parameter = self.ensemble_specs.get("parameter", None)
+        if ensemble_name not in ("NVE", "NVE_4th_order"):
+            if ensemble_parameter is None:
+                raise ValueError(f"Ensemble {ensemble_name} parameter is not defined")
+            updates["MD_run"].append(ensemble_parameter)
+
+        # ...and put everything in the string
+        updates["MD_run"] = " ".join(map(str, updates["MD_run"]))
+
+        # initialize velocities
+        if self.init_velocities:
+            if self.temp is None:
+                raise ValueError("Temperature must be set for velocity initialisation")
+            updates["MD_MB_init"] = self.temp
+
+        return updates

tests/io/aims/aims_input_generator_ref/md-si/parameters.json

@@ -0,0 +1 @@
+{"xc": "pbe", "relativistic": "atomic_zora scalar", "MD_run": "10.0 NVT_parrinello 300 0.4", "MD_time_step": 0.002, "MD_MB_init": 300, "species_dir": "/home/andrey/workspace/pymatgen/tests/io/aims/species_directory/light", "k_grid": [2, 2, 2]}

tests/io/aims/test_sets/test_md_generator.py

@@ -0,0 +1,48 @@
+"""Tests the MD input set generator"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from pymatgen.io.aims.sets.core import MDSetGenerator
+from pymatgen.util.testing.aims import Si, compare_files
+
+module_dir = Path(__file__).resolve().parents[1]
+species_dir = module_dir / "species_directory"
+ref_path = (module_dir / "aims_input_generator_ref").resolve()
+
+
+def test_si_md(tmp_path):
+    # default behaviour
+    parameters = {
+        "species_dir": str(species_dir / "light"),
+        "k_grid": [2, 2, 2],
+    }
+    test_name = "md-si"
+    # First do the exceptions
+    with pytest.raises(ValueError, match="Ensemble something not valid"):
+        MDSetGenerator(ensemble="something").get_input_set(Si)
+    with pytest.raises(ValueError, match="Type parrinello is not valid for nve ensemble"):
+        MDSetGenerator(ensemble="nve", ensemble_specs={"type": "parrinello"}).get_input_set(Si)
+    with pytest.raises(ValueError, match="Velocities must be initialized"):
+        MDSetGenerator(ensemble="nve", init_velocities=False).get_input_set(Si)
+    with pytest.raises(ValueError, match="Temperature must be set"):
+        MDSetGenerator(ensemble="nve").get_input_set(Si)
+    with pytest.raises(ValueError, match="Temperature must be set"):
+        MDSetGenerator(ensemble="nvt").get_input_set(Si)
+    with pytest.raises(ValueError, match="parameter is not defined"):
+        MDSetGenerator(ensemble="nve", ensemble_specs={"type": "damped"}, temp=300).get_input_set(Si)
+    # then do the actual input set
+    generator = MDSetGenerator(
+        ensemble="nvt",
+        ensemble_specs={"type": "parrinello", "parameter": 0.4},
+        temp=300,
+        time=10.0,
+        time_step=0.002,
+        user_params=parameters,
+    )
+    input_set = generator.get_input_set(Si)
+    input_set.write_input(tmp_path / test_name)
+
+    return compare_files(test_name, tmp_path, ref_path)