Merge pull request #1079 from materialsproject/more-mlffs

Add MatterSim, Allegro and OCP models (`fairchem-core`) to `ase_calculators`

Author: esoteric-ephemera

pyproject.toml

@@ -61,6 +61,7 @@ forcefields = [
     "matgl>=1.2.1",
     "torchdata<=0.7.1", # TODO: remove when issue fixed
     "quippy-ase>=0.9.14",
+    "mattersim>=1.0.1",
     "sevenn>=0.9.3",
     "deepmd-kit>=2.1.4",
 ]

src/atomate2/common/jobs/qha.py

@@ -1,4 +1,4 @@
-"""Jobs for running qha calculations."""
+"""Jobs for running QHA calculations."""
 
 from __future__ import annotations
 

src/atomate2/common/schemas/qha.py

@@ -1,4 +1,4 @@
-"""Schemas for qha documents."""
+"""Schemas for QHA documents."""
 
 import logging
 from typing import Union
@@ -15,7 +15,7 @@
 
 
 class PhononQHADoc(StructureMetadata, extra="allow"):  # type: ignore[call-arg]
-    """Collection of all data produced by the qha workflow."""
+    """Collection of all data produced by the QHA workflow."""
 
     structure: Structure | None = Field(
         None, description="Structure of Materials Project."
@@ -62,7 +62,7 @@ class PhononQHADoc(StructureMetadata, extra="allow"):  # type: ignore[call-arg]
         description="Gruneisen parameters at temperatures.Shape: (temperatures,)",
     )
     pressure: float | None = Field(
-        None, description="Pressure in GPA at which Gibb's energy was computed"
+        None, description="Pressure in GPa at which the Gibbs energy was computed."
     )
     t_max: float | None = Field(
         None,
@@ -106,7 +106,7 @@ def from_phonon_runs(
         eos_type: str = "vinet",
         **kwargs,
     ) -> Self:
-        """Generate qha results.
+        """Generate QHA results.
 
         Parameters
         ----------
@@ -149,35 +149,28 @@ def from_phonon_runs(
 
         # create some plots here
         # add kwargs to change the names and file types
+        fig_ext = kwargs.get("plot_type", "pdf")
         qha.plot_helmholtz_volume().savefig(
-            f"{kwargs.get('helmholtz_volume_filename', 'helmholtz_volume')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
+            f"{kwargs.get('helmholtz_volume_filename', 'helmholtz_volume')}.{fig_ext}"
         )
         qha.plot_volume_temperature().savefig(
-            f"{kwargs.get('volume_temperature_plot', 'volume_temperature')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
+            f"{kwargs.get('volume_temperature_plot', 'volume_temperature')}.{fig_ext}"
         )
         qha.plot_thermal_expansion().savefig(
-            f"{kwargs.get('thermal_expansion_plot', 'thermal_expansion')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
+            f"{kwargs.get('thermal_expansion_plot', 'thermal_expansion')}.{fig_ext}"
         )
         qha.plot_gibbs_temperature().savefig(
-            f"{kwargs.get('gibbs_temperature_plot', 'gibbs_temperature')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
+            f"{kwargs.get('gibbs_temperature_plot', 'gibbs_temperature')}.{fig_ext}"
         )
         qha.plot_bulk_modulus_temperature().savefig(
-            f"{kwargs.get('bulk_modulus_plot', 'bulk_modulus_temperature')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
+            f"{kwargs.get('bulk_modulus_plot', 'bulk_modulus_temperature')}.{fig_ext}"
         )
         qha.plot_heat_capacity_P_numerical().savefig(
-            f"{kwargs.get('heat_capacity_plot', 'heat_capacity_P_numerical')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
+            f"{kwargs.get('heat_capacity_plot', 'heat_capacity_P_numerical')}.{fig_ext}"
         )
         # qha.plot_heat_capacity_P_polyfit().savefig("heat_capacity_P_polyfit.eps")
-        qha.plot_gruneisen_temperature().savefig(
-            f"{kwargs.get('gruneisen_temperature_plot', 'gruneisen_temperature')}"
-            f".{kwargs.get('plot_type', 'pdf')}"
-        )
+        ge_temp_plot = kwargs.get("gruneisen_temperature_plot", "gruneisen_temperature")
+        qha.plot_gruneisen_temperature().savefig(f"{ge_temp_plot}.{fig_ext}")
 
         qha.write_helmholtz_volume(
             filename=kwargs.get("helmholtz_volume_datafile", "helmholtz_volume.dat")
@@ -197,21 +190,16 @@ def from_phonon_runs(
         qha.write_gibbs_temperature(
             filename=kwargs.get("gibbs_temperature_datafile", "gibbs_temperature.dat")
         )
-        qha.write_gruneisen_temperature(
-            filename=kwargs.get(
-                "gruneisen_temperature_datafile", "gruneisen_temperature.dat"
-            )
+        ge_temp_file = kwargs.get(
+            "gruneisen_temperature_datafile", "gruneisen_temperature.dat"
         )
+        qha.write_gruneisen_temperature(filename=ge_temp_file)
         qha.write_heat_capacity_P_numerical(
             filename=kwargs.get(
                 "heat_capacity_datafile", "heat_capacity_P_numerical.dat"
             )
         )
-        qha.write_gruneisen_temperature(
-            filename=kwargs.get(
-                "gruneisen_temperature_datafile", "gruneisen_temperature.dat"
-            )
-        )
+        qha.write_gruneisen_temperature(filename=ge_temp_file)
 
         # write files as well - might be easier for plotting
 

src/atomate2/forcefields/__init__.py

@@ -2,3 +2,5 @@
 
 # ensure that this is still importable for legacy jobs
 from atomate2.forcefields.utils import MLFF, _get_formatted_ff_name
+
+__all__ = ["MLFF"]

src/atomate2/forcefields/utils.py

@@ -42,6 +42,9 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     MATPES_R2SCAN = "MatPES-r2SCAN"
     MATPES_PBE = "MatPES-PBE"
     DeepMD = "DeepMD"
+    Allegro = "Allegro"
+    OCP = "OCP"  # for loading model checkpoint with fairchem.core.OCPCalculator
+    MatterSim = "MatterSim"
 
     @classmethod
     def _missing_(cls, value: Any) -> Any:
@@ -353,6 +356,23 @@ def ase_calculator(
 
                 calculator = DP(**kwargs)
 
+            case MLFF.Allegro:
+                from allegro.ase import AllegroCalculator
+
+                calculator = AllegroCalculator.from_deployed_model(**kwargs)
+
+            case MLFF.OCP:
+                # Not available on PyPI, needs to be installed from source
+                # see https://github.com/FAIR-Chem/fairchem?tab=readme-ov-file#installation
+                from fairchem.core import OCPCalculator
+
+                calculator = OCPCalculator(**kwargs)
+
+            case MLFF.MatterSim:
+                from mattersim.forcefield import MatterSimCalculator
+
+                calculator = MatterSimCalculator(**kwargs)
+
     elif isinstance(calculator_meta, dict):
         calc_cls = _load_calc_cls(calculator_meta)
         calculator = calc_cls(**kwargs)

src/atomate2/vasp/flows/qha.py

@@ -20,7 +20,7 @@ class QhaMaker(CommonQhaMaker):
     First relax a structure using relax_maker.
     Then perform a series of deformations on the relaxed structure, and
     then compute harmonic phonons for each deformed structure.
-    Finally, compute Gibb's free energy.
+    Finally, compute Gibbs free energy.
 
     Parameters
     ----------

src/atomate2/vasp/sets/eos.py

@@ -25,7 +25,7 @@ class EosSetGenerator(VaspInputGenerator):
     force_gamma: bool = True
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -60,7 +60,7 @@ class MPLegacyEosRelaxSetGenerator(VaspInputGenerator):
     config_dict: dict = field(default_factory=lambda: MPRelaxSet.CONFIG)
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -103,7 +103,7 @@ class MPLegacyEosStaticSetGenerator(EosSetGenerator):
     config_dict: dict = field(default_factory=lambda: MPRelaxSet.CONFIG)
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -138,7 +138,7 @@ class MPGGAEosRelaxSetGenerator(VaspInputGenerator):
     config_dict: dict = field(default_factory=lambda: MPScanRelaxSet.CONFIG)
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -173,7 +173,7 @@ class MPGGAEosStaticSetGenerator(EosSetGenerator):
     config_dict: dict = field(default_factory=lambda: MPScanRelaxSet.CONFIG)
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -207,7 +207,7 @@ class MPMetaGGAEosStaticSetGenerator(VaspInputGenerator):
     config_dict: dict = field(default_factory=lambda: MPScanRelaxSet.CONFIG)
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -250,7 +250,7 @@ class MPMetaGGAEosRelaxSetGenerator(VaspInputGenerator):
     bandgap_tol: float = 1e-4
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:
@@ -295,7 +295,7 @@ class MPMetaGGAEosPreRelaxSetGenerator(VaspInputGenerator):
     bandgap_tol: float = 1e-4
     auto_ismear: bool = False
     auto_kspacing: bool = False
-    inherit_incar: bool = False
+    inherit_incar: bool | list[str] = False
 
     @property
     def incar_updates(self) -> dict:

tests/forcefields/test_jobs.py

@@ -16,7 +16,7 @@
 
 try:
     import dgl
-except ImportError:
+except Exception:  # noqa: BLE001
     dgl = None
 
 

tests/forcefields/test_md.py

@@ -42,7 +42,13 @@ def test_maker_initialization():
             ) == ForceFieldMDMaker(force_field_name=mlff)
 
 
-@pytest.mark.parametrize("ff_name, use_emmet_models", product(MLFF, [True, False]))
+_mlffs_for_test = set(MLFF).difference(
+    map(MLFF, ("Forcefield", "MatterSim", "Allegro", "OCP", "M3GNet", "MACE"))
+)
+_md_test_params = sorted(product(_mlffs_for_test, [True, False]), key=lambda x: str(x))
+
+
+@pytest.mark.parametrize("ff_name, use_emmet_models", _md_test_params)
 def test_ml_ff_md_maker(
     ff_name,
     use_emmet_models,
@@ -53,14 +59,10 @@ def test_ml_ff_md_maker(
     clean_dir,
     get_deepmd_pretrained_model_path,
 ):
-    if ff_name in map(MLFF, ("Forcefield", "MACE")):
-        return  # nothing to test here, MLFF.Forcefield is just a generic placeholder
     if ff_name == MLFF.GAP and sys.version_info >= (3, 12):
         pytest.skip(
             "GAP model not compatible with Python 3.12, waiting on https://github.com/libAtoms/QUIP/issues/645"
         )
-    if ff_name == MLFF.M3GNet:
-        pytest.skip("M3GNet requires DGL which is PyTorch 2.4 incompatible")
 
     n_steps = 5
 

tests/forcefields/test_phonon.py

@@ -68,7 +68,19 @@ def test_phonon_maker_initialization_with_all_mlff(
     # DGL which is PyTorch 2.4 incompatible, raises
     # "FileNotFoundError: Cannot find DGL C++ libgraphbolt_pytorch_2.4.1.so"
     skip_mlff = set(
-        map(MLFF, ["Forcefield", "GAP", "M3GNet", "MATPES_R2SCAN", "MATPES_PBE"])
+        map(
+            MLFF,
+            [
+                "Forcefield",
+                "GAP",
+                "M3GNet",
+                "MATPES_R2SCAN",
+                "MATPES_PBE",
+                "Allegro",
+                "OCP",
+                "MatterSim",
+            ],
+        )
     )
     for mlff in set(MLFF).difference(skip_mlff):
         calc_kwargs = {