File size: 2,746 Bytes
fecdc11 | 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 | """Run a short NVT molecular-dynamics trajectory with eSEN."""
from __future__ import annotations
import argparse
import os
os.environ.setdefault(
"ONESCIENCE_ESEN_JD_PATH",
os.path.join(os.path.dirname(__file__), "weight", "Jd.pt"),
)
import numpy as np
from ase import units
from ase.build import bulk
from ase.io import read
from ase.io.trajectory import Trajectory
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary
from onescience.utils.esen import eSENCalculator
def default_checkpoint() -> str:
return os.path.join(os.path.dirname(__file__), "weight", "esen_30m_mptrj.pt")
def load_structure(path: str | None, repeat: tuple[int, int, int] | None):
if path:
atoms = read(path)
else:
atoms = bulk("Si", "diamond", a=5.43).repeat((2, 2, 2))
return atoms.repeat(repeat) if repeat else atoms
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", default=default_checkpoint())
parser.add_argument("--input", help="CIF, POSCAR, XYZ, or another ASE-readable structure")
parser.add_argument(
"--repeat",
type=int,
nargs=3,
metavar=("NX", "NY", "NZ"),
help="repeat the input structure along its three cell vectors",
)
parser.add_argument("--steps", type=int, default=10)
parser.add_argument("--temperature", type=float, default=300.0)
parser.add_argument("--timestep", type=float, default=1.0, help="time step in fs")
parser.add_argument(
"--friction", type=float, default=0.01, help="Langevin friction in 1/fs"
)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--output", default="md.traj")
parser.add_argument("--device", default="cuda")
args = parser.parse_args()
repeat = tuple(args.repeat) if args.repeat else None
atoms = load_structure(args.input, repeat)
atoms.calc = eSENCalculator.from_checkpoint(args.checkpoint, device=args.device)
rng = np.random.default_rng(args.seed)
MaxwellBoltzmannDistribution(atoms, temperature_K=args.temperature, rng=rng)
Stationary(atoms)
dynamics = Langevin(
atoms,
timestep=args.timestep * units.fs,
temperature_K=args.temperature,
friction=args.friction / units.fs,
)
trajectory = Trajectory(args.output, "w", atoms)
dynamics.attach(trajectory.write, interval=1)
dynamics.run(args.steps)
trajectory.close()
print("formula:", atoms.get_chemical_formula())
print("atoms:", len(atoms))
print("steps:", dynamics.nsteps)
print("energy (eV):", atoms.get_potential_energy())
if __name__ == "__main__":
main()
|