| import argparse | |
| from ase.build import bulk | |
| from onescience.utils.mattersim import relax | |
| parser = argparse.ArgumentParser(description="MatterSim structure relaxation") | |
| parser.add_argument( | |
| "--checkpoint", | |
| default="../weight/mattersim-v1.0.0-1M.pth", | |
| help="Path to MatterSim checkpoint", | |
| ) | |
| parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") | |
| args = parser.parse_args() | |
| converged, atoms = relax( | |
| bulk("Si", "diamond", a=5.43), | |
| checkpoint=args.checkpoint, | |
| filter="FrechetCellFilter", | |
| fmax=0.01, | |
| steps=20, | |
| device=args.device, | |
| ) | |
| print(f"Converged: {converged}") | |
| print(f"Energy: {atoms.get_potential_energy():.6f} eV") | |