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