Add UMA (#574)

ElliottKasoar · oerc0122 · web-flow · commit 5f5ed84edfd7 · 2025-07-24T17:50:20.000+01:00
* Add UMA dependency

* Add UMA calculator

* Add tests for UMA

* Add UMA tests to workflows

* Fix CI installations

* Update CI install and delete cached MLIPs

* Change default UMA model

* Swap tested UMA models

* Reduce memory used by UMA tests

* Add invalid path UMA test and reorder

* Skip tests if unauthorised for model

* ADD UMA to README list

* Update docs for UMA/fairchem

---------

Co-authored-by: Jacob Wilkins &lt;46597752+oerc0122@users.noreply.github.com&gt;
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -38,7 +38,7 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest --cov janus_core --cov-append .
+        run: uv run --no-sync pytest --cov janus_core --cov-append .
 
       - name: Install updated e3nn dependencies
         run: |
@@ -51,7 +51,7 @@ jobs:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: uv run pytest tests/test_{mlip_calculators,single_point}.py
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py
 
       - name: Install dgl dependencies
         run: |
@@ -62,7 +62,25 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest tests/test_{mlip_calculators,single_point,eos}.py
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point,eos}.py
+
+      - name: Create space in cache
+        run: |
+          rm -rf ~/.cache/*
+          uv cache clean
+
+      - name: Install UMA
+        run: |
+          uv sync --extra uma
+          uv pip install --reinstall pynvml
+          uv pip install fairchem-core[torch-extras] --no-build-isolation
+
+      - name: Run test suite for UMA
+        env:
+          # show timings of tests
+          PYTEST_ADDOPTS: "--durations=0"
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py
 
       - name: Report coverage to Coveralls
         uses: coverallsapp/github-action@v2
diff --git a/.github/workflows/mac.yml b/.github/workflows/mac.yml
@@ -38,7 +38,7 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest
+        run: uv run --no-sync pytest
 
       - name: Install updated e3nn dependencies
         run: |
@@ -51,7 +51,7 @@ jobs:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: uv run pytest tests/test_{mlip_calculators,single_point}.py
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py
 
       - name: Install dgl dependencies
         run: |
@@ -62,4 +62,17 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest tests/test_{mlip_calculators,single_point,eos}.py
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point,eos}.py
+
+      - name: Install uma
+        run: |
+          uv sync --extra uma
+          uv pip install --reinstall pynvml
+          uv pip install fairchem-core[torch-extras] --no-build-isolation
+
+      - name: Run test suite for UMA
+        env:
+          # show timings of tests
+          PYTEST_ADDOPTS: "--durations=0"
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py
diff --git a/.github/workflows/windows.yml b/.github/workflows/windows.yml
@@ -27,7 +27,7 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run pytest
+        run: uv run --no-sync pytest
 
       - name: Install updated e3nn dependencies
         run: |
@@ -40,4 +40,4 @@ jobs:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: uv run pytest tests/test_mlip_calculators.py tests/test_single_point.py
+        run: uv run --no-sync pytest tests/test_mlip_calculators.py tests/test_single_point.py
diff --git a/README.md b/README.md
@@ -92,6 +92,7 @@ Current and planned features include:
   - GRACE
   - EquiformerV2
   - eSEN
+  - UMA
   - PET-MAD
 - [x] Single point calculations
 - [x] Geometry optimisation
diff --git a/docs/source/user_guide/get_started.rst b/docs/source/user_guide/get_started.rst
@@ -59,7 +59,8 @@ Currently supported MLIP ``extras`` are:
 - ``orb``: `Orb <https://github.com/orbital-materials/orb-models>`_
 - ``mattersim``: `MatterSim <https://github.com/microsoft/mattersim>`_
 - ``grace``: `GRACE <https://github.com/ICAMS/grace-tensorpotential>`_
-- ``fairchem``: `eqV2 DeNS/eSEN <https://github.com/FAIR-Chem/fairchem/tree/main/src/fairchem/core>`_
+- ``fairchem``: `eqV2 DeNS/eSEN <https://github.com/facebookresearch/fairchem/tree/fairchem_core-1.10.0/src/fairchem/core>`_
+- ``uma``: `UMA <https://github.com/FAIR-Chem/fairchem/tree/main/src/fairchem/core>`_
 - ``pet-mad``: `PET-MAD <https://github.com/lab-cosmo/pet-mad>`_
 
 .. note::
diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
@@ -132,6 +132,7 @@ class Correlation(TypedDict, total=True):
     "esen",
     "equiformer",
     "pet_mad",
+    "uma",
 ]
 Devices = Literal["cpu", "cuda", "mps", "xpu"]
 Ensembles = Literal["nph", "npt", "nve", "nvt", "nvt-nh", "nvt-csvr", "npt-mtk"]
diff --git a/janus_core/helpers/mlip_calculators.py b/janus_core/helpers/mlip_calculators.py
@@ -53,6 +53,7 @@ def _set_model(
         "path",
         "model_name",
         "checkpoint_path",
+        "predict_unit",
     }
     present = kwargs.keys() & model_kwargs
 
@@ -382,6 +383,31 @@ def choose_calculator(
             if model is None:
                 model = LATEST_VERSION
 
+        case "uma":
+            from fairchem.core import FAIRChemCalculator, __version__, pretrained_mlip
+            from fairchem.core.units.mlip_unit import MLIPPredictUnit
+
+            match model:
+                case MLIPPredictUnit():
+                    predict_unit = model
+                    model = "loaded_Module"
+                case Path() | str():
+                    predict_unit = pretrained_mlip.get_predict_unit(
+                        model_name=model, device=device
+                    )
+                case None:
+                    model = "uma-m-1p1"
+                    predict_unit = pretrained_mlip.get_predict_unit(
+                        model_name=model, device=device
+                    )
+
+            kwargs.setdefault("task_name", "omat")
+
+            calculator = FAIRChemCalculator(
+                predict_unit=predict_unit,
+                **kwargs,
+            )
+
         case _:
             raise ValueError(
                 f"Unrecognized {arch=}. Suported architectures "
diff --git a/pyproject.toml b/pyproject.toml
@@ -28,7 +28,7 @@ readme = "README.md"
 dependencies = [
     "ase<4.0,>=3.25",
     "codecarbon<3.0.0,>=2.8.4",
-    "numpy<2.0.0,>=1.26.4",
+    "numpy<3.0.0,>=1.26.4",
     "phonopy<3.0.0,>=2.23.1",
     "pymatgen>=2025.1.24",
     "pyyaml<7.0.0,>=6.0.1",
@@ -93,6 +93,10 @@ mattersim = [
     "mattersim == 1.1.2; sys_platform != 'win32'",
 ]
 
+uma = [
+    "fairchem-core == 2.3.0",
+]
+
 # MLIPs with dgl dependency
 alignn = [
     "alignn == 2024.5.27; sys_platform != 'win32'",
@@ -281,4 +285,36 @@ conflicts = [
       { extra = "fairchem" },
       { extra = "all" },
     ],
+    [
+      { extra = "uma" },
+      { extra = "alignn" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "fairchem" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "grace" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "mace" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "mattersim" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "m3gnet" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "sevennet" },
+    ],
+    [
+      { extra = "uma" },
+      { extra = "all" },
+    ],
 ]
diff --git a/tests/test_mlip_calculators.py b/tests/test_mlip_calculators.py
@@ -31,6 +31,7 @@
 EQUIFORMER_LABEL = "EquiformerV2-83M-S2EF-OC20-2M"
 ESEN_LABEL = "eSEN-30M-MP"
 
+
 try:
     from fairchem.core.models.model_registry import model_name_to_local_file
 
@@ -59,6 +60,23 @@
 
 SEVENNET_PATH = MODEL_PATH / "sevennet_0.pth"
 
+UMA_LABEL = "uma-s-1"
+
+try:
+    from fairchem.core import pretrained_mlip
+    from huggingface_hub.errors import GatedRepoError
+
+    try:
+        UMA_PREDICT_UNIT = pretrained_mlip.get_predict_unit(
+            model_name=UMA_LABEL, device="cpu"
+        )
+    except GatedRepoError:
+        UMA_PREDICT_UNIT = None
+
+except ImportError:
+    UMA_PREDICT_UNIT = None
+
+
 ALIGNN_PATH = MODEL_PATH / "v5.27.2024"
 M3GNET_DIR_PATH = MODEL_PATH / "M3GNet-MP-2021.2.8-DIRECT-PES"
 M3GNET_MODEL_PATH = M3GNET_DIR_PATH / "model.pt"
@@ -78,17 +96,11 @@
 @pytest.mark.parametrize(
     "arch, device, kwargs",
     [
-        ("mace", "cpu", {"model": MACE_MP_PATH}),
-        ("mace", "cpu", {"model_paths": MACE_MP_PATH}),
-        ("mace", "cpu", {"model_path": MACE_MP_PATH}),
-        ("mace_off", "cpu", {}),
-        ("mace_off", "cpu", {"model": "small"}),
-        ("mace_off", "cpu", {"model_path": MACE_OFF_PATH}),
-        ("mace_off", "cpu", {"model": MACE_OFF_PATH}),
-        ("mace_mp", "cpu", {}),
-        ("mace_mp", "cpu", {"model": "small"}),
-        ("mace_mp", "cpu", {"model_path": MACE_MP_PATH}),
-        ("mace_mp", "cpu", {"model": MACE_MP_PATH}),
+        ("alignn", "cpu", {}),
+        ("alignn", "cpu", {"model_path": ALIGNN_PATH}),
+        ("alignn", "cpu", {"model_path": ALIGNN_PATH / "best_model.pt"}),
+        ("alignn", "cpu", {"model": "alignnff_wt10"}),
+        ("alignn", "cpu", {"path": ALIGNN_PATH}),
         ("chgnet", "cpu", {}),
         ("chgnet", "cpu", {"model": "0.2.0"}),
         ("chgnet", "cpu", {"model_path": CHGNET_PATH}),
@@ -107,8 +119,26 @@
         ("esen", "cpu", {"model_name": ESEN_LABEL}),
         ("esen", "cpu", {"model_name": ESEN_PATH}),
         ("esen", "cpu", {"checkpoint_path": ESEN_PATH}),
+        ("grace", "cpu", {}),
+        ("grace", "cpu", {"model_path": "GRACE-1L-MP-r6"}),
+        ("mace", "cpu", {"model": MACE_MP_PATH}),
+        ("mace", "cpu", {"model_paths": MACE_MP_PATH}),
+        ("mace", "cpu", {"model_path": MACE_MP_PATH}),
+        ("mace_mp", "cpu", {}),
+        ("mace_mp", "cpu", {"model": "small"}),
+        ("mace_mp", "cpu", {"model_path": MACE_MP_PATH}),
+        ("mace_mp", "cpu", {"model": MACE_MP_PATH}),
+        ("mace_off", "cpu", {}),
+        ("mace_off", "cpu", {"model": "small"}),
+        ("mace_off", "cpu", {"model_path": MACE_OFF_PATH}),
+        ("mace_off", "cpu", {"model": MACE_OFF_PATH}),
         ("mattersim", "cpu", {}),
         ("mattersim", "cpu", {"model_path": "mattersim-v1.0.0-1m"}),
+        ("m3gnet", "cpu", {}),
+        ("m3gnet", "cpu", {"model_path": M3GNET_DIR_PATH}),
+        ("m3gnet", "cpu", {"model_path": M3GNET_MODEL_PATH}),
+        ("m3gnet", "cpu", {"potential": M3GNET_DIR_PATH}),
+        ("m3gnet", "cpu", {"potential": M3GNET_POTENTIAL}),
         ("nequip", "cpu", {"model_path": NEQUIP_PATH}),
         ("nequip", "cpu", {"model": NEQUIP_PATH}),
         ("orb", "cpu", {}),
@@ -121,18 +151,11 @@
         ("sevennet", "cpu", {"model_path": SEVENNET_PATH}),
         ("sevennet", "cpu", {}),
         ("sevennet", "cpu", {"model": "sevennet-0"}),
-        ("alignn", "cpu", {}),
-        ("alignn", "cpu", {"model_path": ALIGNN_PATH}),
-        ("alignn", "cpu", {"model_path": ALIGNN_PATH / "best_model.pt"}),
-        ("alignn", "cpu", {"model": "alignnff_wt10"}),
-        ("alignn", "cpu", {"path": ALIGNN_PATH}),
-        ("m3gnet", "cpu", {}),
-        ("m3gnet", "cpu", {"model_path": M3GNET_DIR_PATH}),
-        ("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"}),
+        ("uma", "cpu", {"model": UMA_LABEL}),
+        ("uma", "cpu", {"model_path": UMA_LABEL}),
+        ("uma", "cpu", {"model_name": UMA_LABEL}),
+        ("uma", "cpu", {"model": UMA_PREDICT_UNIT}),
+        ("uma", "cpu", {"predict_unit": UMA_PREDICT_UNIT}),
     ],
 )
 def test_mlips(arch, device, kwargs):
@@ -164,20 +187,22 @@ def test_invalid_arch():
 @pytest.mark.parametrize(
     "arch, model",
     [
-        ("mace", "/invalid/path"),
-        ("mace_off", "/invalid/path"),
-        ("mace_mp", "/invalid/path"),
+        ("alignn", "invalid/path"),
         ("chgnet", "/invalid/path"),
         ("dpa3", "/invalid/path"),
-        ("fairchem", "/invalid/path"),
+        ("equiformer", "/invalid/path"),
+        ("esen", "/invalid/path"),
         ("grace", "/invalid/path"),
+        ("mace", "/invalid/path"),
+        ("mace_mp", "/invalid/path"),
+        ("mace_off", "/invalid/path"),
         ("mattersim", "/invalid/path"),
+        ("m3gnet", "/invalid/path"),
         ("nequip", "/invalid/path"),
         ("orb", "/invalid/path"),
         ("pet_mad", "/invalid/path"),
         ("sevenn", "/invalid/path"),
-        ("alignn", "invalid/path"),
-        ("m3gnet", "/invalid/path"),
+        ("uma", "/invalid/path"),
     ],
 )
 def test_invalid_model(arch, model):
@@ -190,10 +215,6 @@ def test_invalid_model(arch, model):
 @pytest.mark.parametrize(
     "kwargs",
     [
-        {"arch": "mace", "model": MACE_MP_PATH, "model_paths": MACE_MP_PATH},
-        {"arch": "mace", "model": MACE_MP_PATH, "model_paths": MACE_MP_PATH},
-        {"arch": "mace", "model_path": MACE_MP_PATH, "model": MACE_MP_PATH},
-        {"arch": "mace", "model": MACE_MP_PATH, "potential": MACE_MP_PATH},
         {
             "arch": "alignn",
             "model_path": ALIGNN_PATH / "best_model.pt",
@@ -207,11 +228,6 @@ def test_invalid_model(arch, model):
         {"arch": "chgnet", "model": CHGNET_PATH, "path": CHGNET_PATH},
         {"arch": "dpa3", "model_path": DPA3_PATH, "model": DPA3_PATH},
         {"arch": "dpa3", "model": DPA3_PATH, "path": DPA3_PATH},
-        {
-            "arch": "esen",
-            "model_path": ESEN_LABEL,
-            "model": ESEN_LABEL,
-        },
         {
             "arch": "equiformer",
             "model_path": EQUIFORMER_LABEL,
@@ -222,6 +238,11 @@ def test_invalid_model(arch, model):
             "model_path": "GRACE-1L-MP-r6",
             "model": "GRACE-1L-MP-r6",
         },
+        {"arch": "esen", "model_path": ESEN_LABEL, "model": ESEN_LABEL},
+        {"arch": "mace", "model": MACE_MP_PATH, "model_paths": MACE_MP_PATH},
+        {"arch": "mace", "model": MACE_MP_PATH, "model_paths": MACE_MP_PATH},
+        {"arch": "mace", "model_path": MACE_MP_PATH, "model": MACE_MP_PATH},
+        {"arch": "mace", "model": MACE_MP_PATH, "potential": MACE_MP_PATH},
         {
             "arch": "mattersim",
             "model": "mattersim-v1.0.0-1m",
@@ -243,6 +264,7 @@ def test_invalid_model(arch, model):
         },
         {"arch": "sevennet", "model_path": SEVENNET_PATH, "model": SEVENNET_PATH},
         {"arch": "sevennet", "model": SEVENNET_PATH, "path": SEVENNET_PATH},
+        {"arch": "uma", "model_path": UMA_LABEL, "model": UMA_LABEL},
     ],
 )
 def test_model_model_paths(kwargs):
diff --git a/tests/test_single_point.py b/tests/test_single_point.py
diff --git a/tests/utils.py b/tests/utils.py

Original file line number	Diff line number	Diff line change
`@@ -132,6 +132,7 @@ class Correlation(TypedDict, total=True):`
`132`	`132`	`"esen",`
`133`	`133`	`"equiformer",`
`134`	`134`	`"pet_mad",`
	`135`	`+ "uma",`
`135`	`136`	`]`
`136`	`137`	`Devices = Literal["cpu", "cuda", "mps", "xpu"]`
`137`	`138`	`Ensembles = Literal["nph", "npt", "nve", "nvt", "nvt-nh", "nvt-csvr", "npt-mtk"]`