File size: 2,318 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 | """Run one Equiformer V3 energy, force, and stress prediction through 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
from onescience.utils.equiformer_v3 import (
EquiformerV3Calculator,
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)
# Keep the example runnable without requiring a separate structure file.
return bulk("Cu")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint")
parser.add_argument(
"--input",
help=(
"CIF, POSCAR, XYZ, trajectory, or another ASE-readable structure; "
"defaults to the built-in periodic Cu example"
),
)
parser.add_argument("--device", default="cuda")
parser.add_argument("--output", default="outputs/single_point.json")
args = parser.parse_args()
checkpoint = args.checkpoint or default_checkpoint()
atoms = load_structure(args.input)
atoms.calc = EquiformerV3Calculator.from_checkpoint(
checkpoint,
device=args.device,
)
result = {
"formula": atoms.get_chemical_formula(),
"natoms": len(atoms),
"input": str(Path(args.input).expanduser()) if args.input else None,
"input_source": args.input or "ASE bulk Cu default",
"checkpoint": str(Path(checkpoint).expanduser()),
"pbc": atoms.pbc.tolist(),
"cell_angstrom": atoms.cell.array.tolist(),
"energy_ev": float(atoms.get_potential_energy()),
"forces_ev_per_angstrom": atoms.get_forces().tolist(),
"stress_ev_per_angstrom_cubed_voigt": atoms.get_stress().tolist(),
}
output = write_workflow_result(result, args.output)
print("formula:", result["formula"])
print("atoms:", result["natoms"])
print("energy (eV):", result["energy_ev"])
print("result:", output)
if __name__ == "__main__":
main()
|