File size: 3,210 Bytes
096347b | 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 | """Run finite-displacement phonons with Equiformer V3 and ASE."""
from __future__ import annotations
import argparse
import os
from pathlib import Path
os.environ.setdefault(
"ONESCIENCE_EQUIFORMER_V3_JD_PATH",
str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
)
from ase.build import bulk
from ase.io import read, write
from onescience.utils.equiformer_v3 import (
EquiformerV3Calculator,
relax_structure,
run_phonon_workflow,
write_workflow_result,
)
def default_checkpoint() -> str:
return str(
Path(__file__).resolve().parent
/ "weight"
/ "omat24-mptrj-salex_gradient.pt"
)
def load_structure(path: str | None):
if path:
return read(path)
return bulk("Cu")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint")
parser.add_argument(
"--input",
help="periodic structure readable by ASE; defaults to built-in bulk Cu",
)
parser.add_argument(
"--relax",
action="store_true",
help="relax the cell and atomic positions before finite displacements",
)
parser.add_argument("--relax-fmax", type=float, default=0.01)
parser.add_argument("--relax-steps", type=int, default=200)
parser.add_argument("--supercell", type=int, nargs=3, default=(3, 3, 3))
parser.add_argument("--delta", type=float, default=0.01)
parser.add_argument("--bandpath", help="ASE band path, for example GXWKGL")
parser.add_argument("--band-points", type=int, default=100)
parser.add_argument("--dos-kpts", type=int, nargs=3, default=(10, 10, 10))
parser.add_argument("--dos-points", type=int, default=400)
parser.add_argument("--dos-width", type=float, default=0.001, help="eV")
parser.add_argument("--device", default="cuda")
parser.add_argument("--output-dir", default="outputs/phonons")
args = parser.parse_args()
checkpoint = args.checkpoint or default_checkpoint()
calculator = EquiformerV3Calculator.from_checkpoint(
checkpoint, device=args.device
)
output_dir = Path(args.output_dir)
structure = load_structure(args.input)
relaxation = None
if args.relax:
structure, relaxation = relax_structure(
structure,
calculator,
relax_cell=True,
fmax=args.relax_fmax,
steps=args.relax_steps,
)
output_dir.mkdir(parents=True, exist_ok=True)
write(output_dir / "relaxed.cif", structure)
result = run_phonon_workflow(
structure,
calculator,
output_dir,
supercell=args.supercell,
delta=args.delta,
bandpath=args.bandpath,
band_points=args.band_points,
dos_kpts=args.dos_kpts,
dos_points=args.dos_points,
dos_width_ev=args.dos_width,
)
result["checkpoint"] = str(checkpoint)
result["relaxation"] = relaxation
output = write_workflow_result(result, output_dir / "phonons.json")
print("formula:", result["formula"])
print("minimum phonon energy (eV):", result["minimum_band_energy_ev"])
print("result:", output)
if __name__ == "__main__":
main()
|