fix import sort and add OMAT-0

ACEsuit · Feb 4, 2025 · 76f5790 · 76f5790
1 parent bb2162c
commit 76f5790
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/mace/calculators/foundations_models.py b/mace/calculators/foundations_models.py
@@ -40,6 +40,7 @@ def download_mace_mp_checkpoint(model: Union[str, Path] = None) -> str:
         "large-0b2": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b2/mace-large-density-agnesi-stress.model",
         "medium-0b3": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b3/mace-mp-0b3-medium.model",
         "medium-mpa-0": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model",
+        "medium-omat-0": "https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-medium.model",
     }
 
     checkpoint_url = (
@@ -57,6 +58,7 @@ def download_mace_mp_checkpoint(model: Union[str, Path] = None) -> str:
             "large-0b2",
             "medium-0b3",
             "medium-mpa-0",
+            "medium-omat-0",
         )
         else model
     )
@@ -65,6 +67,10 @@ def download_mace_mp_checkpoint(model: Union[str, Path] = None) -> str:
         print(
             "Using medium MPA-0 model as default MACE-MP model, to use previous (before 3.10) default model please specify 'medium' as model argument"
         )
+    if checkpoint_url == urls["medium-omat-0"]:
+        print(
+            "Using medium OMAT-0 model under Academic Software License (ASL) license, see https://github.com/gabor1/ASL \n To use this model you accept the terms of the license."
+        )
 
     cache_dir = os.path.expanduser("~/.cache/mace")
     checkpoint_url_name = "".join(
@@ -138,6 +144,7 @@ def mace_mp(
             "medium-0b2",
             "medium-0b3",
             "large-0b2",
+            "medium-omat-0",
         ) or str(model).startswith("https:"):
             model_path = download_mace_mp_checkpoint(model)
             print(f"Using Materials Project MACE for MACECalculator with {model_path}")

diff --git a/mace/calculators/mace.py b/mace/calculators/mace.py
@@ -4,9 +4,10 @@
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
 
+import logging
+
 # pylint: disable=wrong-import-position
 import os
-import logging
 from glob import glob
 from pathlib import Path
 from typing import List, Union

diff --git a/mace/cli/create_lammps_model.py b/mace/cli/create_lammps_model.py
@@ -1,6 +1,6 @@
 # pylint: disable=wrong-import-position
-import os
 import argparse
+import os
 
 os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"