add task `stability` and md run script

mlip_arena/tasks/stability/run.py

@@ -1,3 +1,190 @@
+from __future__ import annotations
 
-from mlip_arena.tasks.utils import _valid_dynamics, _preset_dynamics
+import datetime
+from datetime import datetime
+from pathlib import Path
+from typing import Literal, Sequence
 
+import numpy as np
+import torch
+from ase import Atoms, units
+from ase.calculators.mixing import SumCalculator
+from ase.io import read
+from ase.io.trajectory import Trajectory
+from ase.md.md import MolecularDynamics
+from ase.md.npt import NPT
+from ase.md.velocitydistribution import (
+    MaxwellBoltzmannDistribution,
+    Stationary,
+    ZeroRotation,
+)
+from scipy.linalg import schur
+from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+from tqdm.auto import tqdm
+
+from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
+from mlip_arena.tasks.utils import (
+    _get_ensemble_defaults,
+    _get_ensemble_schedule,
+    _preset_dynamics,
+    _valid_dynamics,
+)
+
+
+def md(
+    atoms: Atoms,
+    calculator_name: str | EXTMLIPEnum,
+    calculator_kwargs: dict | None,
+    dispersion: str | None = None,
+    dispersion_kwargs: dict | None = None,
+    device: str | None = None,
+    ensemble: Literal["nve", "nvt", "npt"] = "nvt",
+    dynamics: str | MolecularDynamics = "langevin",
+    time_step: float | None = None,
+    total_time: float = 1000,
+    temperature: float | Sequence | np.ndarray | None = 300.0,
+    pressure: float | Sequence | np.ndarray | None = None,
+    ase_md_kwargs: dict | None = None,
+    mb_velocity_seed: int | None = None,
+    zero_linear_momentum: bool = True,
+    zero_angular_momentum: bool = True,
+    traj_file: str | Path | None = None,
+    traj_interval: int = 1,
+    # ttime: float = 25 * units.fs,
+    # pfactor: float = (75 * units.fs) ** 1 * units.GPa,
+    # mask: np.ndarray | list[int] | None = None,
+    # traceless: float = 1.0,
+    restart: bool = True,
+    # interval: int = 500,
+    # device: str | None = None,
+    # dtype: str = "float64",
+):
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    print(f"Using device: {device}")
+
+    calculator_kwargs = calculator_kwargs or {}
+
+    if isinstance(calculator_name, EXTMLIPEnum) and calculator_name in EXTMLIPEnum:
+        calc = external_ase_calculator(calculator_name, **calculator_kwargs)
+    elif calculator_name in MLIPMap:
+        calc = MLIPMap[calculator_name](**calculator_kwargs)
+
+    print(f"Using calculator: {calc}")
+
+    dispersion_kwargs = dispersion_kwargs or {}
+
+    dispersion_kwargs.update({"device": device})
+
+    if dispersion is not None:
+        disp_calc = TorchDFTD3Calculator(
+            **dispersion_kwargs,
+        )
+        calc = SumCalculator([calc, disp_calc])
+
+        print(f"Using dispersion: {dispersion}")
+
+    atoms.calc = calc
+
+    if time_step is None:
+        # If a structure contains an isotope of hydrogen, set default `time_step`
+        # to 0.5 fs, and 2 fs otherwise.
+        has_h_isotope = "H" in atoms.get_chemical_symbols()
+        time_step = 0.5 if has_h_isotope else 2.0
+
+    n_steps = int(total_time / time_step)
+
+    t_schedule, p_schedule = _get_ensemble_schedule(
+        ensemble=ensemble,
+        n_steps=n_steps,
+        temperature=temperature,
+        pressure=pressure,
+    )
+
+    ase_md_kwargs = _get_ensemble_defaults(
+        ensemble=ensemble,
+        dynamics=dynamics,
+        t_schedule=t_schedule,
+        p_schedule=p_schedule,
+        ase_md_kwargs=ase_md_kwargs,
+    )
+
+    if isinstance(dynamics, str):
+        # Use known dynamics if `self.dynamics` is a str
+        dynamics = dynamics.lower()
+        if dynamics not in _valid_dynamics[ensemble]:
+            raise ValueError(
+                f"{dynamics} thermostat not available for {ensemble}."
+                f"Available {ensemble} thermostats are:"
+                " ".join(_valid_dynamics[ensemble])
+            )
+
+        if ensemble == "nve" and dynamics is None:
+            dynamics = "velocityverlet"
+        md_class = _preset_dynamics[f"{ensemble}_{dynamics}"]
+    elif issubclass(dynamics, MolecularDynamics):
+        md_class = dynamics
+
+    if md_class is NPT:
+        #  Note that until md_func is instantiated, isinstance(md_func,NPT) is False
+        # ASE NPT implementation requires upper triangular cell
+        u, _ = schur(atoms.get_cell(complete=True), output="complex")
+        atoms.set_cell(u.real, scale_atoms=True)
+
+    last_step = 0
+
+    if traj_file is not None:
+        traj_file = Path(traj_file)
+
+        if restart and traj_file.exists():
+            traj = read(traj_file, index=":")
+            last_step = len(traj)
+            n_steps -= len(traj)
+            last_atoms = traj[-1]
+            traj = Trajectory(traj_file, "a", atoms)
+            atoms.set_positions(last_atoms.get_positions())
+            atoms.set_momenta(last_atoms.get_momenta())
+        else:
+            traj = Trajectory(traj_file, "w", atoms)
+
+            if not np.isnan(t_schedule).any():
+                MaxwellBoltzmannDistribution(
+                    atoms=atoms,
+                    temperature_K=t_schedule[last_step],
+                    rng=np.random.default_rng(seed=mb_velocity_seed),
+                )
+
+            if zero_linear_momentum:
+                Stationary(atoms)
+            if zero_angular_momentum:
+                ZeroRotation(atoms)
+
+    md_runner = md_class(
+        atoms=atoms,
+        timestep=time_step * units.fs,
+        **ase_md_kwargs,
+    )
+
+    if traj_file is not None:
+        md_runner.attach(traj.write, interval=traj_interval)
+
+    with tqdm(total=n_steps) as pbar:
+
+        def _callback(dyn: MolecularDynamics = md_runner) -> None:
+            if ensemble == "nve":
+                return
+            dyn.set_temperature(temperature_K=t_schedule[last_step + dyn.nsteps])
+            if ensemble == "nvt":
+                return
+            dyn.set_stress(p_schedule[last_step + dyn.nsteps] * 1e3 * units.bar)
+            pbar.update()
+
+        md_runner.attach(_callback, interval=1)
+
+        start_time = datetime.now()
+        md_runner.run(steps=n_steps)
+        end_time = datetime.now()
+
+        traj.close()
+
+    return {"md_runtime": end_time - start_time}