Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
esoteric-ephemera merged 15 commits into from
Jan 26, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
78ba4f0
fix typos + minor qha.py refactor + fix mypy "VaspInputGenerator" sup…
janosh Dec 7, 2024
e5945b4
add Allegro, OCP and MatterSim to ase_calculator
janosh Dec 7, 2024
e07e855
add MatterSim and SevenNet to test_ext_load
janosh Dec 7, 2024
461c528
test energy and forces produced by MLFFs in test_ext_load
janosh Dec 7, 2024
49f391b
pyproject use mattersim>=1.0.0rc9
janosh Dec 7, 2024
5aade53
bump mattersim to v1.0.1rc1
janosh Dec 23, 2024
471c033
resolve merge conflicts
esoteric-ephemera Jan 26, 2026
92a3f08
remove mattersim from tests to allow ci to run
esoteric-ephemera Jan 26, 2026
ad0b469
skip mattersim in md test
esoteric-ephemera Jan 26, 2026
5969dcb
camel case
esoteric-ephemera Jan 26, 2026
ed44b7d
skip mlmd tests more transparently
esoteric-ephemera Jan 26, 2026
1b55d5c
ensure pytest xdist gets same tests in order
esoteric-ephemera Jan 26, 2026
bda310a
skip some mlffs in ph wf
esoteric-ephemera Jan 26, 2026
8e48ceb
fix mlff check
esoteric-ephemera Jan 26, 2026
a1b50a7
dgl probs
esoteric-ephemera Jan 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]
Expand Down
2 changes: 1 addition & 1 deletion src/atomate2/common/jobs/qha.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Jobs for running qha calculations."""
"""Jobs for running QHA calculations."""

from __future__ import annotations

Expand Down
46 changes: 17 additions & 29 deletions src/atomate2/common/schemas/qha.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Schemas for qha documents."""
"""Schemas for QHA documents."""

import logging
from typing import Union
Expand All @@ -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."
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -106,7 +106,7 @@ def from_phonon_runs(
eos_type: str = "vinet",
**kwargs,
) -> Self:
"""Generate qha results.
"""Generate QHA results.

Parameters
----------
Expand Down Expand Up @@ -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")
Expand All @@ -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

Expand Down
2 changes: 2 additions & 0 deletions src/atomate2/forcefields/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
20 changes: 20 additions & 0 deletions src/atomate2/forcefields/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/atomate2/vasp/flows/qha.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand Down
16 changes: 8 additions & 8 deletions src/atomate2/vasp/sets/eos.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion tests/forcefields/test_jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

try:
import dgl
except ImportError:
except Exception: # noqa: BLE001
dgl = None


Expand Down
12 changes: 7 additions & 5 deletions tests/forcefields/test_md.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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

Expand Down
14 changes: 13 additions & 1 deletion tests/forcefields/test_phonon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
36 changes: 27 additions & 9 deletions tests/forcefields/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from typing import TYPE_CHECKING

import numpy as np
import pytest

from atomate2.forcefields import MLFF
Expand All @@ -13,24 +14,41 @@
from pymatgen.core import Structure


@pytest.mark.parametrize(("force_field"), [mlff.value for mlff in MLFF])
def test_mlff(force_field: str):
mlff = MLFF(force_field)
@pytest.mark.parametrize("mlff", MLFF)
def test_mlff(mlff: MLFF):
assert mlff == MLFF(str(mlff)) == MLFF(str(mlff).split(".")[-1])


def test_ext_load(test_dir):
calc_from_decode = ase_calculator(
{"@module": "mace.calculators", "@callable": "mace_mp"}
)
calc_from_preset = ase_calculator(str(MLFF.MACE_MP_0))
calc_from_enum = ase_calculator(MLFF.MACE_MP_0)
@pytest.mark.parametrize("mlff", ["MACE", MLFF.SevenNet])
def test_ext_load(mlff: str | MLFF, test_dir, si_structure: Structure):
decode_dict = {
"MACE": {"@module": "mace.calculators", "@callable": "mace_mp"},
MLFF.SevenNet: {
"@module": "sevenn.sevennet_calculator",
"@callable": "SevenNetCalculator",
},
}[mlff]
formatted_mlff = MLFF(mlff)
calc_from_decode = ase_calculator(decode_dict)
calc_from_preset = ase_calculator(str(formatted_mlff))
calc_from_enum = ase_calculator(formatted_mlff)

for other in (calc_from_preset, calc_from_enum):
assert type(calc_from_decode) is type(other)
assert calc_from_decode.name == other.name
assert calc_from_decode.parameters == other.parameters == {}

atoms = si_structure.to_ase_atoms()

atoms.calc = calc_from_preset
energy = atoms.get_potential_energy()
forces = atoms.get_forces()

assert isinstance(energy, float | np.floating)
assert energy < 0
assert forces.shape == (2, 3)
assert abs(forces.sum()) < 1e-6, f"unexpectedly large net {forces=}"


def test_raises_error():
with pytest.raises(ValueError, match="Could not create"):
Expand Down
Loading