Added NEP MLIP relax, static, and MD makers (materialsproject#893)

* add NEP FF relax. static and MD makers * add tests for NEP relax, static and md makers * add model license and source * Update tests/forcefields/test_jobs.py Co-authored-by: Christina Ertural <[email protected]> --------- Co-authored-by: Christina Ertural <[email protected]>
hrushikesh-s · Nov 16, 2024 · fa67259 · fa67259
1 parent 6580110
commit fa67259
Show file tree

Hide file tree

Showing 11 changed files with 78,607 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -51,6 +51,7 @@ defects = [
 ]
 forcefields = [
     "ase>=3.22.1",
+    "calorine<=2.2.1",
     "chgnet>=0.2.2",
     "mace-torch>=0.3.3",
     "matgl>=1.1.1",
@@ -88,6 +89,7 @@ strict = [
     # must use >= for ase to not uninstall main branch install in CI
     # last known working commit: https://gitlab.com/ase/ase@2bab58f4e
     "ase>=3.22.1",
+    "calorine==2.2.1",
     "cclib==1.8.1",
     "chgnet==0.3.8",
     "click==8.1.7",

diff --git a/src/atomate2/forcefields/__init__.py b/src/atomate2/forcefields/__init__.py
@@ -11,4 +11,5 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     M3GNet = "M3GNet"
     CHGNet = "CHGNet"
     Forcefield = "Forcefield"  # default placeholder option
+    NEP = "NEP"
     Nequip = "Nequip"
diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py
@@ -298,6 +298,79 @@ class M3GNetRelaxMaker(ForceFieldRelaxMaker):
     )
 
 
+@dataclass
+class NEPRelaxMaker(ForceFieldRelaxMaker):
+    """
+    Base Maker to calculate forces and stresses using a NEP potential.
+
+    Parameters
+    ----------
+    name : str
+        The job name.
+    force_field_name : str
+        The name of the force field.
+    relax_cell : bool = True
+        Whether to allow the cell shape/volume to change during relaxation.
+    fix_symmetry : bool = False
+        Whether to fix the symmetry during relaxation.
+        Refines the symmetry of the initial structure.
+    symprec : float = 1e-2
+        Tolerance for symmetry finding in case of fix_symmetry.
+    steps : int
+        Maximum number of ionic steps allowed during relaxation.
+    relax_kwargs : dict
+        Keyword arguments that will get passed to :obj:`Relaxer.relax`.
+    optimizer_kwargs : dict
+        Keyword arguments that will get passed to :obj:`Relaxer()`.
+    calculator_kwargs : dict
+        Keyword arguments that will get passed to the ASE calculator.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+    """
+
+    name: str = f"{MLFF.NEP} relax"
+    force_field_name: str = f"{MLFF.NEP}"
+    relax_cell: bool = True
+    fix_symmetry: bool = False
+    symprec: float = 1e-2
+    steps: int = 500
+    relax_kwargs: dict = field(default_factory=dict)
+    optimizer_kwargs: dict = field(default_factory=dict)
+    calculator_kwargs: dict = field(
+        default_factory=lambda: {
+            "model_filename": "nep.txt",
+        }
+    )
+    task_document_kwargs: dict = field(default_factory=dict)
+
+
+@dataclass
+class NEPStaticMaker(ForceFieldStaticMaker):
+    """
+    Base Maker to calculate forces and stresses using a NEP potential.
+
+    Parameters
+    ----------
+    name : str
+        The job name.
+    force_field_name : str
+        The name of the force field.
+    calculator_kwargs : dict
+        Keyword arguments that will get passed to the ASE calculator.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+    """
+
+    name: str = f"{MLFF.NEP} static"
+    force_field_name: str = f"{MLFF.NEP}"
+    task_document_kwargs: dict = field(default_factory=dict)
+    calculator_kwargs: dict = field(
+        default_factory=lambda: {
+            "model_filename": "nep.txt",
+        }
+    )
+
+
 @dataclass
 class NequipRelaxMaker(ForceFieldRelaxMaker):
     """

diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py
@@ -333,6 +333,17 @@ def calculator(self) -> Calculator:
         return ase_calculator(self.force_field_name, **self.calculator_kwargs)
 
 
+@dataclass
+class NEPMDMaker(ForceFieldMDMaker):
+    """Perform an MD run with NEP."""
+
+    name: str = f"{MLFF.NEP} MD"
+    force_field_name: str = f"{MLFF.NEP}"
+    calculator_kwargs: dict = field(
+        default_factory=lambda: {"model_filename": "nep.txt"}
+    )
+
+
 @dataclass
 class MACEMDMaker(ForceFieldMDMaker):
     """Perform an MD run with MACE."""

diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py
@@ -435,6 +435,11 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
 
             calculator = Potential(**kwargs)
 
+        elif calculator_name == MLFF.NEP:
+            from calorine.calculators import CPUNEP
+
+            calculator = CPUNEP(**kwargs)
+
         elif calculator_name == MLFF.Nequip:
             from nequip.ase import NequIPCalculator
 

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -97,6 +97,11 @@ def si_structure(test_dir):
     return Structure.from_file(test_dir / "structures" / "Si.cif")
 
 
+@pytest.fixture()
+def al2_au_structure(test_dir):
+    return Structure.from_file(test_dir / "structures" / "Al2Au.cif")
+
+
 @pytest.fixture()
 def sr_ti_o3_structure(test_dir):
     return Structure.from_file(test_dir / "structures" / "SrTiO3.cif")

diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py
@@ -16,6 +16,8 @@
     M3GNetStaticMaker,
     MACERelaxMaker,
     MACEStaticMaker,
+    NEPRelaxMaker,
+    NEPStaticMaker,
     NequipRelaxMaker,
     NequipStaticMaker,
 )
@@ -295,6 +297,79 @@ def test_gap_static_maker(si_structure: Structure, test_dir):
     assert output1.forcefield_version == get_imported_version("quippy-ase")
 
 
+def test_nep_static_maker(al2_au_structure: Structure, test_dir: Path):
+    task_doc_kwargs = {"ionic_step_data": ("structure", "energy")}
+
+    # NOTE: The test NEP model is specifically trained on 16 elemental metals
+    # thus a new Al2Au structure is added.
+    # The NEP model used for the tests is licensed under a
+    # [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode)
+    # and downloaded from https://doi.org/10.5281/zenodo.10081677
+    # Also cite the original work if you use this specific model : https://arxiv.org/abs/2311.04732
+    job = NEPStaticMaker(
+        task_document_kwargs=task_doc_kwargs,
+        calculator_kwargs={
+            "model_filename": test_dir / "forcefields" / "nep" / "nep.txt"
+        },
+    ).make(al2_au_structure)
+
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(job, ensure_success=True)
+
+    # validation the outputs of the job
+    output1 = responses[job.uuid][1].output
+    assert isinstance(output1, ForceFieldTaskDocument)
+    assert output1.output.energy == approx(-47.65972, rel=1e-4)
+    assert output1.output.n_steps == 1
+
+
+@pytest.mark.parametrize(
+    ("relax_cell", "fix_symmetry"),
+    [(True, False), (False, True)],
+)
+def test_nep_relax_maker(
+    al2_au_structure: Structure,
+    test_dir: Path,
+    relax_cell: bool,
+    fix_symmetry: bool,
+):
+    # NOTE: The test NEP model is specifically trained on 16 elemental metals
+    # thus a new Al2Au structure is added.
+    # The NEP model used for the tests is licensed under a
+    # [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode)
+    # and downloaded from https://doi.org/10.5281/zenodo.10081677
+    # Also cite the original work if you use this specific model : https://arxiv.org/abs/2311.04732
+
+    # generate job
+    job = NEPRelaxMaker(
+        steps=25,
+        optimizer_kwargs={"optimizer": "BFGSLineSearch"},
+        relax_cell=relax_cell,
+        fix_symmetry=fix_symmetry,
+        calculator_kwargs={
+            "model_filename": test_dir / "forcefields" / "nep" / "nep.txt"
+        },
+    ).make(al2_au_structure)
+
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(job, ensure_success=True)
+
+    # validate the outputs of the job
+    output1 = responses[job.uuid][1].output
+    assert isinstance(output1, ForceFieldTaskDocument)
+    if relax_cell:
+        assert output1.output.energy == approx(-47.6727, rel=1e-3)
+        assert output1.output.n_steps == 3
+    else:
+        assert output1.output.energy == approx(-47.659721, rel=1e-4)
+        assert output1.output.n_steps == 2
+
+    # fix_symmetry makes no difference for this structure relaxer combo
+    # just testing that passing fix_symmetry doesn't break
+    final_spg_num = output1.output.structure.get_space_group_info()[1]
+    assert final_spg_num == 225
+
+
 def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path):
     importorskip("nequip")
     task_doc_kwargs = {"ionic_step_data": ("structure", "energy")}

diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py
@@ -17,6 +17,7 @@
     GAPMDMaker,
     M3GNetMDMaker,
     MACEMDMaker,
+    NEPMDMaker,
     NequipMDMaker,
 )
 
@@ -25,22 +26,26 @@
     "M3GNet": M3GNetMDMaker,
     "MACE": MACEMDMaker,
     "GAP": GAPMDMaker,
+    "NEP": NEPMDMaker,
     "Nequip": NequipMDMaker,
 }
 
 
 @pytest.mark.parametrize(
     "ff_name",
-    ["CHGNet", "M3GNet", "MACE", "GAP", "Nequip"],
+    ["CHGNet", "M3GNet", "MACE", "GAP", "NEP", "Nequip"],
 )
-def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, clean_dir):
+def test_ml_ff_md_maker(
+    ff_name, si_structure, sr_ti_o3_structure, al2_au_structure, test_dir, clean_dir
+):
     n_steps = 5
 
     ref_energies_per_atom = {
         "CHGNet": -5.280157089233398,
         "M3GNet": -5.387282371520996,
         "MACE": -5.311369895935059,
         "GAP": -5.391255755606209,
+        "NEP": -3.966232215741286,
         "Nequip": -8.84670181274414,
     }
 
@@ -54,6 +59,17 @@ def test_ml_ff_md_maker(ff_name, si_structure, sr_ti_o3_structure, test_dir, cle
             "args_str": "IP GAP",
             "param_filename": str(test_dir / "forcefields" / "gap" / "gap_file.xml"),
         }
+    elif ff_name == "NEP":
+        # NOTE: The test NEP model is specifically trained on 16 elemental metals
+        # thus a new Al2Au structure is added.
+        # The NEP model used for the tests is licensed under a
+        # [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode)
+        # and downloaded from https://doi.org/10.5281/zenodo.10081677
+        # Also cite the original work if you use this specific model : https://arxiv.org/abs/2311.04732
+        calculator_kwargs = {
+            "model_filename": test_dir / "forcefields" / "nep" / "nep.txt"
+        }
+        unit_cell_structure = al2_au_structure.copy()
     elif ff_name == "Nequip":
         calculator_kwargs = {
             "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth"

diff --git a/tests/test_data/forcefields/nep/README.md b/tests/test_data/forcefields/nep/README.md
@@ -0,0 +1,3 @@
+The NEP model used for the tests is licensed under a [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode) and
+downloaded from https://doi.org/10.5281/zenodo.10081677. Also cite the original work if you use this
+specific model : https://arxiv.org/abs/2311.04732