File size: 1,028 Bytes
24d7cbe | 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 | import argparse
import torch
from ase.build import bulk
from ase.units import GPa
from onescience.models.mattersim import load_calculator
def main() -> None:
parser = argparse.ArgumentParser(description="MatterSim single-point inference")
parser.add_argument("--checkpoint")
parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
args = parser.parse_args()
device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
atoms = bulk("Si", "diamond", a=5.43)
atoms.calc = load_calculator(checkpoint=args.checkpoint, device=device)
energy = atoms.get_potential_energy()
print(f"Energy (eV) = {energy}")
print(f"Energy per atom (eV/atom) = {energy / len(atoms)}")
print(f"Forces of first atom (eV/A) = {atoms.get_forces()[0]}")
stress = atoms.get_stress(voigt=False)
print(f"Stress[0][0] (eV/A^3) = {stress[0][0]}")
print(f"Stress[0][0] (GPa) = {stress[0][0] / GPa}")
if __name__ == "__main__":
main()
|