Merge pull request #25 from microsoft/hanyang/refactor-potential

Refactor: Update the method to load checkpoints

README.md

@@ -40,7 +40,7 @@ from mattersim.datasets.utils.build import build_dataloader
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 print(f"Running MatterSim on {device}")
 
-potential = Potential.load(load_path="pretrained_models/mattersim-v1.0.0-1M.pth", device=device)
+potential = Potential.load()
 si = bulk("Si", "diamond", a=5.43)
 dataloader = build_dataloader([si], only_inference=True)
 
@@ -53,7 +53,7 @@ print(predictions)
 We kindly request that users of MatterSim version 1.0.0 cite our preprint available on arXiv:
 ```
 @article{yang2024mattersim,
-      title={MatterSim: A Deep Learning Atomistic Model Across Elements, Temperatures and Pressures}, 
+      title={MatterSim: A Deep Learning Atomistic Model Across Elements, Temperatures and Pressures},
       author={Han Yang and Chenxi Hu and Yichi Zhou and Xixian Liu and Yu Shi and Jielan Li and Guanzhi Li and Zekun Chen and Shuizhou Chen and Claudio Zeni and Matthew Horton and Robert Pinsler and Andrew Fowler and Daniel Zügner and Tian Xie and Jake Smith and Lixin Sun and Qian Wang and Lingyu Kong and Chang Liu and Hongxia Hao and Ziheng Lu},
       year={2024},
       eprint={2405.04967},

src/mattersim/applications/relax.py

@@ -7,6 +7,7 @@
 from ase.filters import ExpCellFilter, FrechetCellFilter
 from ase.optimize import BFGS, FIRE
 from ase.optimize.optimize import Optimizer
+from ase.units import GPa
 
 
 class Relaxer(object):
@@ -53,7 +54,7 @@ def relax(
         steps: int = 500,
         fmax: float = 0.01,
         params_filter: dict = {},
-        **kwargs
+        **kwargs,
     ) -> Atoms:
         """
         Relax the atoms object.
@@ -115,7 +116,7 @@ def relax_structures(
         constrain_symmetry: bool = False,
         fix_axis: Union[bool, Iterable[bool]] = False,
         pressure_in_GPa: Union[float, None] = None,
-        **kwargs
+        **kwargs,
     ) -> Union[Tuple[bool, Atoms], Tuple[List[bool], List[Atoms]]]:
         """
         Args:
@@ -138,11 +139,15 @@ def relax_structures(
             pass
         elif filter is None and pressure_in_GPa is not None:
             filter = "ExpCellFilter"
-            params_filter["scalar_pressure"] = pressure_in_GPa / 160.21766208
+            params_filter["scalar_pressure"] = (
+                pressure_in_GPa * GPa
+            )  # GPa = 1 / 160.21766208
         elif filter is not None and pressure_in_GPa is None:
             params_filter["scalar_pressure"] = 0.0
         else:
-            params_filter["scalar_pressure"] = pressure_in_GPa / 160.21766208
+            params_filter["scalar_pressure"] = (
+                pressure_in_GPa * GPa
+            )  # GPa = / 160.21766208
 
         relaxer = Relaxer(
             optimizer=optimizer,