diff --git a/pyproject.toml b/pyproject.toml index f4d2dfa67..0f0e288d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ forcefields = [ "matgl>=1.1.3", "torchdata<=0.7.1", # quippy-ase support for py3.12 tracked in https://github.com/libAtoms/QUIP/issues/645 + "mattersim>=1.0.1rc1", "quippy-ase>=0.9.14; python_version < '3.12'", "sevenn>=0.9.3", "torchdata<=0.7.1", # TODO: remove when issue fixed @@ -123,6 +124,7 @@ strict-forcefields = [ "chgnet==0.4.0", "mace-torch>=0.3.6", "matgl==1.1.3", + "mattersim==1.0.1rc1", "quippy-ase==0.9.14; python_version < '3.12'", "sevenn==0.10.1", "torch==2.5.1", diff --git a/src/atomate2/common/jobs/qha.py b/src/atomate2/common/jobs/qha.py index 51f657cd7..9bd58808c 100644 --- a/src/atomate2/common/jobs/qha.py +++ b/src/atomate2/common/jobs/qha.py @@ -1,4 +1,4 @@ -"""Jobs for running qha calculations.""" +"""Jobs for running QHA calculations.""" from __future__ import annotations diff --git a/src/atomate2/common/schemas/qha.py b/src/atomate2/common/schemas/qha.py index bda924fbb..f3c2ea8d5 100644 --- a/src/atomate2/common/schemas/qha.py +++ b/src/atomate2/common/schemas/qha.py @@ -1,4 +1,4 @@ -"""Schemas for qha documents.""" +"""Schemas for QHA documents.""" import logging from typing import Optional, 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: Optional[Structure] = Field( None, description="Structure of Materials Project." @@ -64,7 +64,7 @@ class PhononQHADoc(StructureMetadata, extra="allow"): # type: ignore[call-arg] description="Gruneisen parameters at temperatures.Shape: (temperatures, )", ) pressure: Optional[float] = Field( - None, description="Pressure in GPA at which Gibb's energy was computed" + None, description="Pressure in GPA at which Gibbs energy was computed" ) t_max: Optional[float] = Field( None, @@ -75,17 +75,17 @@ class PhononQHADoc(StructureMetadata, extra="allow"): # type: ignore[call-arg] free_energies: Optional[list[list[float]]] = Field( None, description="List of free energies in J/mol for per formula unit. " - "Shape: (temperatuers, volumes)", + "Shape: (temperatures, volumes)", ) heat_capacities: Optional[list[list[float]]] = Field( None, description="List of heat capacities in J/K/mol per formula unit. " - "Shape: (temperatuers, volumes)", + "Shape: (temperatures, volumes)", ) entropies: Optional[list[list[float]]] = Field( None, description="List of entropies in J/(K*mol) per formula unit. " - "Shape: (temperatuers, volumes) ", + "Shape: (temperatures, volumes) ", ) formula_units: Optional[int] = Field(None, description="Formula units") @@ -108,7 +108,7 @@ def from_phonon_runs( eos_type: str = "vinet", **kwargs, ) -> Self: - """Generate qha results. + """Generate QHA results. Parameters ---------- @@ -151,35 +151,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") @@ -199,21 +192,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 diff --git a/src/atomate2/forcefields/__init__.py b/src/atomate2/forcefields/__init__.py index 59e25c1be..00c7b390d 100644 --- a/src/atomate2/forcefields/__init__.py +++ b/src/atomate2/forcefields/__init__.py @@ -20,10 +20,13 @@ class MLFF(Enum): # TODO inherit from StrEnum when 3.11+ NEP = "NEP" Nequip = "Nequip" SevenNet = "SevenNet" + Allegro = "Allegro" + OCP = "OCP" # for loading model checkpoint with fairchem.core.OCPCalculator + MatterSim = "MatterSim" @classmethod def _missing_(cls, value: Any) -> Any: - """Allow input of str(MLFF) as valid enum.""" + """Allow feeding output of str(MLFF.) back into MLFF(...).""" if isinstance(value, str): value = value.split("MLFF.")[-1] for member in cls: diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 043b7b037..c5aeefaab 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -90,11 +90,28 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N calculator = NequIPCalculator.from_deployed_model(**kwargs) + elif calculator_name == MLFF.Allegro: + from allegro.ase import AllegroCalculator + + calculator = AllegroCalculator.from_deployed_model(**kwargs) + elif calculator_name == MLFF.SevenNet: from sevenn.sevennet_calculator import SevenNetCalculator calculator = SevenNetCalculator(**{"model": "7net-0"} | kwargs) + elif calculator_name == MLFF.OCP: + # this package is 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) + + elif calculator_name == MLFF.MatterSim: + from mattersim.forcefield import MatterSimCalculator + + calculator = MatterSimCalculator(**kwargs) + elif isinstance(calculator_meta, dict): calc_cls = MontyDecoder().process_decoded(calculator_meta) calculator = calc_cls(**kwargs) diff --git a/src/atomate2/vasp/flows/qha.py b/src/atomate2/vasp/flows/qha.py index 2b9742d45..9b73aafab 100644 --- a/src/atomate2/vasp/flows/qha.py +++ b/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 ---------- diff --git a/src/atomate2/vasp/sets/eos.py b/src/atomate2/vasp/sets/eos.py index 5758a9c2e..5690640c0 100644 --- a/src/atomate2/vasp/sets/eos.py +++ b/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: diff --git a/tests/forcefields/test_utils.py b/tests/forcefields/test_utils.py index b43eb6ff0..2d91ac638 100644 --- a/tests/forcefields/test_utils.py +++ b/tests/forcefields/test_utils.py @@ -1,27 +1,48 @@ +import numpy as np import pytest from atomate2.forcefields import MLFF from atomate2.forcefields.utils import ase_calculator -@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]) -@pytest.mark.parametrize(("force_field"), ["CHGNet", "MACE"]) -def test_ext_load(force_field: str): +@pytest.mark.parametrize("mlff", ["CHGNet", "MACE", MLFF.MatterSim, MLFF.SevenNet]) +def test_ext_load(mlff: str): + from ase.build import bulk + decode_dict = { "CHGNet": {"@module": "chgnet.model.dynamics", "@callable": "CHGNetCalculator"}, "MACE": {"@module": "mace.calculators", "@callable": "mace_mp"}, - }[force_field] + MLFF.MatterSim: { + "@module": "mattersim.forcefield", + "@callable": "MatterSimCalculator", + }, + MLFF.SevenNet: { + "@module": "sevenn.sevennet_calculator", + "@callable": "SevenNetCalculator", + }, + }[mlff] calc_from_decode = ase_calculator(decode_dict) - calc_from_preset = ase_calculator(str(MLFF(force_field))) + calc_from_preset = ase_calculator(str(MLFF(mlff))) assert type(calc_from_decode) is type(calc_from_preset) assert calc_from_decode.name == calc_from_preset.name assert calc_from_decode.parameters == calc_from_preset.parameters == {} + atoms = bulk("Si", "diamond", a=5.43) + + 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"):