Merge pull request #11 from materialsvirtuallab/mace-support

Add MACE to recognized `UNIVERSAL_CALCULATORS`
materialsvirtuallab · Nov 17, 2023 · 63e91c4 · 63e91c4
2 parents 3d90071 + 52eb145
commit 63e91c4
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 56 deletions.
diff --git a/matcalc/elasticity.py b/matcalc/elasticity.py
@@ -49,10 +49,10 @@ def calc(self, structure: Structure) -> dict:
             structure: Pymatgen structure.
 
         Returns: {
-            elastic_tensor: Elastic tensor as a pymatgen ElasticTensor object,
-            shear_modulus_vrh: Voigt-Reuss-Hill shear modulus based on elastic tensor,
-            bulk_modulus_vrh: Voigt-Reuss-Hill bulk modulus based on elastic tensor,
-            youngs_modulus: Young's modulus based on elastic tensor,
+            elastic_tensor: Elastic tensor as a pymatgen ElasticTensor object (in eV/A^3),
+            shear_modulus_vrh: Voigt-Reuss-Hill shear modulus based on elastic tensor (in eV/A^3),
+            bulk_modulus_vrh: Voigt-Reuss-Hill bulk modulus based on elastic tensor (in eV/A^3),
+            youngs_modulus: Young's modulus based on elastic tensor (in eV/A^3),
         }
         """
         if self.relax_structure:

diff --git a/matcalc/phonon.py b/matcalc/phonon.py
@@ -36,11 +36,12 @@ def __init__(
         Args:
             calculator: ASE Calculator to use.
             fmax: Max forces. This criterion is more stringent than for simple relaxation.
-            atom_disp: Atomic displacement
+                Defaults to 0.1 (in eV/Angstrom)
+            atom_disp: Atomic displacement (in Angstrom).
             supercell_matrix: Supercell matrix to use. Defaults to 2x2x2 supercell.
-            t_step: Temperature step.
-            t_max: Max temperature.
-            t_min: Min temperature.
+            t_step: Temperature step (in Kelvin).
+            t_max: Max temperature (in Kelvin).
+            t_min: Min temperature (in Kelvin).
             relax_structure: Whether to first relax the structure. Set to False if structures
                 provided are pre-relaxed with the same calculator.
         """

diff --git a/matcalc/util.py b/matcalc/util.py
@@ -11,6 +11,7 @@
     "M3GNet-MP-2021.2.8-PES",
     "M3GNet-MP-2021.2.8-DIRECT-PES",
     "CHGNet",
+    "MACE",
 )
 
 
@@ -51,4 +52,9 @@ def get_universal_calculator(name: str | Calculator, **kwargs) -> Calculator:
 
         return CHGNetCalculator(**kwargs)
 
+    if name.lower() == "mace":
+        from mace.calculators import mace_mp
+
+        return mace_mp(**kwargs)
+
     raise ValueError(f"Unrecognized {name=}, must be one of {UNIVERSAL_CALCULATORS}")
diff --git a/pyproject.toml b/pyproject.toml
@@ -7,46 +7,44 @@ name = "matcalc"
 authors = [
 
     { name = "Eliott Liu", email = "[email protected]" },
-    { name = "Ji Qi", email = "[email protected]" },
-    { name = "Tsz Wai Ko", email = "[email protected]" },
     { name = "Janosh Riebesell", email = "[email protected]" },
+    { name = "Ji Qi", email = "[email protected]" },
     { name = "Shyue Ping Ong", email = "[email protected]" },
+    { name = "Tsz Wai Ko", email = "[email protected]" },
 ]
 description = "Calculators for materials properties from the potential energy surface."
 readme = "README.md"
 requires-python = ">=3.9"
 keywords = [
-    "materials",
-    "interatomic potential",
-    "force field",
-    "science",
-    "property prediction",
     "AI",
-    "machine learning",
-    "graph",
     "deep learning",
+    "force field",
+    "graph",
+    "interatomic potential",
+    "machine learning",
+    "materials",
+    "property prediction",
+    "science",
 ]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: BSD License",
     "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3",
     "Topic :: Scientific/Engineering :: Chemistry",
     "Topic :: Scientific/Engineering :: Information Analysis",
     "Topic :: Scientific/Engineering :: Physics",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = [
-    "ase", "pymatgen", "joblib", "phonopy"
-]
+dependencies = ["ase", "joblib", "phonopy", "pymatgen"]
 version = "0.0.3"
 
 [project.optional-dependencies]
-models = ["matgl", "chgnet"]
+models = ["chgnet", "mace@git+https://github.com/ACEsuit/mace", "matgl>=0.9.0"]
 
 [tool.setuptools]
 packages = ["matcalc"]
@@ -60,42 +58,42 @@ target-version = ["py39"]
 target-version = "py39"
 line-length = 120
 select = [
-    "B", # flake8-bugbear
-    "C4", # flake8-comprehensions
-    "D", # pydocstyle
-    "E", # pycodestyle error
-    "EXE", # flake8-executable
-    "F", # pyflakes
-    "I", # isort
-    "ICN", # flake8-import-conventions
-    "ISC", # flake8-implicit-str-concat
-    "PD", # pandas-vet
+    "B",    # flake8-bugbear
+    "C4",   # flake8-comprehensions
+    "D",    # pydocstyle
+    "E",    # pycodestyle error
+    "EXE",  # flake8-executable
+    "F",    # pyflakes
+    "I",    # isort
+    "ICN",  # flake8-import-conventions
+    "ISC",  # flake8-implicit-str-concat
+    "PD",   # pandas-vet
     "PERF", # perflint
-    "PIE", # flake8-pie
-    "PL", # pylint
-    "PT", # flake8-pytest-style
-    "PYI", # flakes8-pyi
-    "Q", # flake8-quotes
-    "RET", # flake8-return
-    "RSE", # flake8-raise
-    "RUF", # Ruff-specific rules
-    "SIM", # flake8-simplify
+    "PIE",  # flake8-pie
+    "PL",   # pylint
+    "PT",   # flake8-pytest-style
+    "PYI",  # flakes8-pyi
+    "Q",    # flake8-quotes
+    "RET",  # flake8-return
+    "RSE",  # flake8-raise
+    "RUF",  # Ruff-specific rules
+    "SIM",  # flake8-simplify
     "SLOT", # flake8-slots
-    "TCH", # flake8-type-checking
-    "TID", # tidy imports
-    "TID", # flake8-tidy-imports
-    "UP", # pyupgrade
-    "W", # pycodestyle warning
-    "YTT", # flake8-2020
+    "TCH",  # flake8-type-checking
+    "TID",  # tidy imports
+    "TID",  # flake8-tidy-imports
+    "UP",   # pyupgrade
+    "W",    # pycodestyle warning
+    "YTT",  # flake8-2020
 ]
 ignore = [
-    "B019", # functools.lru_cache on methods can lead to memory leaks
-    "D105", # Missing docstring in magic method
-    "D205", # 1 blank line required between summary line and description
-    "D212", # Multi-line docstring summary should start at the first line
-    "PLR", # pylint refactor
+    "B019",    # functools.lru_cache on methods can lead to memory leaks
+    "D105",    # Missing docstring in magic method
+    "D205",    # 1 blank line required between summary line and description
+    "D212",    # Multi-line docstring summary should start at the first line
+    "PLR",     # pylint refactor
     "PLW0603", # Using the global statement to update variables is discouraged
-    "SIM105", # Use contextlib.suppress(OSError) instead of try-except-pass
+    "SIM105",  # Use contextlib.suppress(OSError) instead of try-except-pass
 ]
 pydocstyle.convention = "google"
 isort.required-imports = ["from __future__ import annotations"]

diff --git a/requirements-ci.txt b/requirements-ci.txt
@@ -5,5 +5,5 @@ coveralls
 mypy
 ruff
 black
-matgl==0.8.1
+matgl==0.9.0
 chgnet==0.2.0
diff --git a/tests/test_util.py b/tests/test_util.py
@@ -5,7 +5,10 @@
 
 
 def test_get_universal_calculator():
-    for name in UNIVERSAL_CALCULATORS:
+    # skip MACE until https://github.com/ACEsuit/mace/pull/230 is merged
+    # maybe even permanently, since the checkpoint downloaded from figshare is 130MB+
+    # and hence would slow down test a lot
+    for name in {*UNIVERSAL_CALCULATORS} - {"MACE"}:
         calc = get_universal_calculator(name)
         assert isinstance(calc, Calculator)
         same_calc = get_universal_calculator(calc)  # test ASE Calculator classes are returned as-is