| """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() |
|
|
