Spaces:
Running
Running
File size: 6,056 Bytes
0ffedd3 9d1a2a5 0ffedd3 9d1a2a5 0ffedd3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 |
from __future__ import annotations
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}
|