add grace calculator, clarify gpu/cpu fixes #408

README.md

@@ -92,6 +92,7 @@ Current and planned features include:
   - DPA3
   - Orb
   - MatterSim
+  - GRACE
 - [x] Single point calculations
 - [x] Geometry optimisation
 - [x] Molecular Dynamics

conftest.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+import sys
+
 import pytest
 
 
@@ -34,3 +36,13 @@ def pytest_collection_modifyitems(config, items):
     for item in items:
         if "extra_mlips" in item.keywords:
             item.add_marker(skip_extra_mlips)
+
+
+@pytest.fixture(autouse=True)
+def capture_wrap():
+    """Block closure of stderr and stdout."""
+    # Prevent `ValueError: I/O operation on closed file` during testing
+    # See discussion in https://github.com/stfc/janus-core/pull/426
+    sys.stderr.close = lambda *args, **kwargs: None
+    sys.stdout.close = lambda *args, **kwargs: None
+    yield

docs/source/getting_started/getting_started.rst

@@ -60,6 +60,7 @@ Currently supported extras are:
 - ``dpa3``: `DPA3 <https://github.com/deepmodeling/deepmd-kit/tree/dpa3-alpha>`_
 - ``orb``: `Orb <https://github.com/orbital-materials/orb-models>`_
 - ``mattersim``: `MatterSim <https://github.com/microsoft/mattersim>`_
+- ``grace``: `GRACE <https://github.com/ICAMS/grace-tensorpotential>`_
 
 .. note::
 

janus_core/helpers/janus_types.py

@@ -128,6 +128,7 @@ class Correlation(TypedDict, total=True):
     "dpa3",
     "orb",
     "mattersim",
+    "grace",
 ]
 Devices = Literal["cpu", "cuda", "mps", "xpu"]
 Ensembles = Literal["nph", "npt", "nve", "nvt", "nvt-nh", "nvt-csvr", "npt-mtk"]

janus_core/helpers/mlip_calculators.py

@@ -305,6 +305,19 @@ def choose_calculator(
             potential=potential, load_path=model_path, device=device, **kwargs
         )
 
+    elif arch == "grace":
+        from tensorpotential.calculator import grace_fm
+
+        __version__ = "0.5.1"
+
+        # Default model
+        model_path = model_path if model_path else "GRACE-2L-OMAT"
+
+        if isinstance(model_path, Path):
+            model_path = str(model_path)
+
+        calculator = grace_fm(model_path, **kwargs)
+
     else:
         raise ValueError(
             f"Unrecognized {arch=}. Suported architectures "

pyproject.toml

@@ -46,6 +46,9 @@ chgnet = [
 dpa3 = [
     "deepmd-kit",
 ]
+grace = [
+    "tensorpotential == 0.5.1",
+]
 mace = [
     "mace-torch==0.3.12",
     "torch-dftd==0.4.0",
@@ -66,6 +69,7 @@ visualise = [
 all = [
     "janus-core[chgnet]",
     "janus-core[dpa3]",
+    "janus-core[grace]",
     "janus-core[mace]",
     "janus-core[nequip]",
     "janus-core[orb]",
@@ -208,6 +212,8 @@ default-groups = [
 constraint-dependencies = [
     "dgl==2.1",
     "torch<2.6",
+    "tensorflow>=2.16.1",
+    "tensorflow-io-gcs-filesystem<=0.31.0; sys_platform == 'win32'",
 ]
 conflicts = [
     [

tests/test_mlip_calculators.py

@@ -95,6 +95,8 @@
         ("m3gnet", "cpu", {"model_path": M3GNET_MODEL_PATH}),
         ("m3gnet", "cpu", {"potential": M3GNET_DIR_PATH}),
         ("m3gnet", "cpu", {"potential": M3GNET_POTENTIAL}),
+        ("grace", "cpu", {}),
+        ("grace", "cpu", {"model_path": "GRACE-1L-MP-r6"}),
     ],
 )
 def test_mlips(arch, device, kwargs):
@@ -127,6 +129,7 @@ def test_invalid_arch():
         ("nequip", "/invalid/path"),
         ("orb", "/invalid/path"),
         ("sevenn", "/invalid/path"),
+        ("grace", "/invalid/path"),
         ("alignn", "invalid/path"),
         ("m3gnet", "/invalid/path"),
     ],
@@ -159,6 +162,11 @@ def test_invalid_model_path(arch, model_path):
         {"arch": "orb", "model_path": ORB_MODEL, "path": ORB_MODEL},
         {"arch": "sevennet", "model_path": SEVENNET_PATH, "path": SEVENNET_PATH},
         {"arch": "sevennet", "model_path": SEVENNET_PATH, "model": SEVENNET_PATH},
+        {
+            "arch": "grace",
+            "model_path": "GRACE-1L-MP-r6",
+            "model": "GRACE-1L-MP-r6",
+        },
         {
             "arch": "alignn",
             "model_path": ALIGNN_PATH / "best_model.pt",

tests/test_single_point.py

@@ -97,6 +97,7 @@ def test_potential_energy(struct, expected, properties, prop_key, calc_kwargs, i
             "NaCl.cif",
             {"model_path": "SevenNet-0_11July2024"},
         ),
+        ("grace", "cpu", -27.081155042373453, "NaCl.cif", {}),
         (
             "alignn",
             "cpu",

tests/utils.py

@@ -131,6 +131,8 @@ def skip_extras(arch: str):
             pytest.importorskip("chgnet")
         case "dpa3":
             pytest.importorskip("deepmd")
+        case "grace":
+            pytest.importorskip("tensorpotential")
         case "mace" | "mace_mp" | "mace_off":
             pytest.importorskip("mace")
         case "mattersim":