Merge pull request #380 from matthewkuner/add_m3gnet_support

Add `m3gnet` support to Atomate2

pyproject.toml

@@ -43,7 +43,7 @@ mp = ["mp-api>=0.27.5"]
 phonons = ["phonopy>=1.10.8", "seekpath"]
 lobster = ["lobsterpy"]
 defects = ["dscribe>=1.2.0", "pymatgen-analysis-defects>=2022.11.30"]
-chgnet = ["chgnet==0.1.3"]
+forcefields = ["chgnet==0.1.3", "matgl==0.5.6"]
 docs = [
     "FireWorks==2.0.3",
     "autodoc_pydantic==1.8.0",
@@ -67,6 +67,7 @@ strict = [
     "emmet-core==0.55.5",
     "jobflow==0.1.11",
     "lobsterpy==0.2.9",
+    "matgl==0.5.6",
     "monty==2023.5.8",
     "mp-api==0.33.3",
     "numpy",

src/atomate2/forcefields/flows.py

@@ -7,7 +7,7 @@
 
 from jobflow import Flow, Maker
 
-from atomate2.forcefields.jobs import CHGNetRelaxMaker
+from atomate2.forcefields.jobs import CHGNetRelaxMaker, M3GNetRelaxMaker
 from atomate2.vasp.jobs.core import RelaxMaker
 
 if TYPE_CHECKING:
@@ -29,7 +29,6 @@ class CHGNetVaspRelaxMaker(Maker):
         Maker to generate a CHGNet relaxation job.
     vasp_maker : .BaseVaspMaker
         Maker to generate a VASP relaxation job.
-
     """
 
     name: str = "CHGNet relax followed by a VASP relax"
@@ -49,11 +48,51 @@ def make(self, structure: Structure):
         -------
         Flow
             A flow containing a CHGNet relaxation followed by a VASP relaxation
-
         """
         chgnet_relax_job = self.chgnet_maker.make(structure)
         chgnet_relax_job.name = "CHGNet pre-relax"
 
         vasp_job = self.vasp_maker.make(chgnet_relax_job.output.structure)
 
         return Flow([chgnet_relax_job, vasp_job], vasp_job.output, name=self.name)
+
+
+@dataclass
+class M3GNetVaspRelaxMaker(Maker):
+    """
+    Maker to (pre)relax a structure using M3GNet and then run VASP.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flow produced by this maker.
+    m3gnet_maker : .M3GNetRelaxMaker
+        Maker to generate a M3GNet relaxation job.
+    vasp_maker : .BaseVaspMaker
+        Maker to generate a VASP relaxation job.
+    """
+
+    name: str = "M3GNet relax followed by a VASP relax"
+    m3gnet_maker: M3GNetRelaxMaker = field(default_factory=M3GNetRelaxMaker)
+    vasp_maker: BaseVaspMaker = field(default_factory=RelaxMaker)
+
+    def make(self, structure: Structure):
+        """
+        Create a flow with a M3GNet (pre)relaxation followed by a VASP relaxation.
+
+        Parameters
+        ----------
+        structure : .Structure
+            A pymatgen structure.
+
+        Returns
+        -------
+        Flow
+            A flow containing a M3GNet relaxation followed by a VASP relaxation
+        """
+        m3gnet_relax_job = self.m3gnet_maker.make(structure)
+        m3gnet_relax_job.name = "M3GNet pre-relax"
+
+        vasp_job = self.vasp_maker.make(m3gnet_relax_job.output.structure)
+
+        return Flow([m3gnet_relax_job, vasp_job], vasp_job.output, name=self.name)

src/atomate2/forcefields/jobs.py

@@ -67,7 +67,8 @@ def make(self, structure: Structure):
             structure, relax_cell=self.relax_cell, steps=self.steps, **self.relax_kwargs
         )
 
-        return ForceFieldTaskDocument.from_chgnet_result(
+        return ForceFieldTaskDocument.from_ase_compatible_result(
+            "CHGNet",
             result,
             self.relax_cell,
             self.steps,
@@ -104,3 +105,106 @@ class CHGNetStaticMaker(CHGNetRelaxMaker):
     relax_kwargs: dict = field(default_factory=dict)
     optimizer_kwargs: dict = field(default_factory=dict)
     task_document_kwargs: dict = field(default_factory=dict)
+
+
+@dataclass
+class M3GNetRelaxMaker(Maker):
+    """
+    Maker to perform a relaxation using the M3GNet universal ML force field.
+
+    Parameters
+    ----------
+    name : str
+        The job name.
+    relax_cell : bool
+        Whether to allow the cell shape/volume to change during relaxation.
+    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()`.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+    """
+
+    name: str = "M3GNet relax"
+    relax_cell: bool = False
+    steps: int = 500
+    relax_kwargs: dict = field(default_factory=dict)
+    optimizer_kwargs: dict = field(default_factory=dict)
+    task_document_kwargs: dict = field(default_factory=dict)
+
+    @job(output_schema=ForceFieldTaskDocument)
+    def make(self, structure: Structure):
+        """
+        Perform a relaxation of a structure using M3GNet.
+
+        Parameters
+        ----------
+        structure: .Structure
+            A pymatgen structure.
+        """
+        import matgl
+        from matgl.ext.ase import Relaxer
+
+        if self.steps < 0:
+            logger.warning(
+                "WARNING: A negative number of steps is not possible. "
+                "Behavior may vary..."
+            )
+
+        # Note: the below code was taken from the matgl repo examples.
+        # Load pre-trained M3GNet model (currently uses the MP-2021.2.8 database)
+        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+
+        relaxer = Relaxer(
+            potential=pot,
+            relax_cell=self.relax_cell,
+            **self.optimizer_kwargs,
+        )
+
+        result = relaxer.relax(
+            structure,
+            steps=self.steps,
+            **self.relax_kwargs,
+        )
+
+        return ForceFieldTaskDocument.from_ase_compatible_result(
+            "M3GNet",
+            result,
+            self.relax_cell,
+            self.steps,
+            self.relax_kwargs,
+            self.optimizer_kwargs,
+            **self.task_document_kwargs,
+        )
+
+
+@dataclass
+class M3GNetStaticMaker(M3GNetRelaxMaker):
+    """
+    Maker to calculate forces and stresses using the M3GNet force field.
+
+    Parameters
+    ----------
+    name : str
+        The job name.
+    relax_cell : bool
+        Whether to allow the cell shape/volume to change during relaxation.
+    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()`.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+    """
+
+    name: str = "M3GNet static"
+    relax_cell: bool = False
+    steps: int = 1
+    relax_kwargs: dict = field(default_factory=dict)
+    optimizer_kwargs: dict = field(default_factory=dict)
+    task_document_kwargs: dict = field(default_factory=dict)

src/atomate2/forcefields/schemas.py

@@ -97,8 +97,9 @@ class ForceFieldTaskDocument(StructureMetadata):
     )
 
     @classmethod
-    def from_chgnet_result(
+    def from_ase_compatible_result(
         cls,
+        forcefield_name: str,
         result: dict,
         relax_cell: bool,
         steps: int,
@@ -107,20 +108,22 @@ def from_chgnet_result(
         ionic_step_data: tuple = ("energy", "forces", "magmoms", "stress", "structure"),
     ):
         """
-        Create a ForceFieldTaskDocument for a CHGNet Task.
+        Create a ForceFieldTaskDocument for a Task that has ASE-compatible outputs.
 
         Parameters
         ----------
+        forcefield_name : str
+            Name of the forcefield used.
         result : dict
             The outputted results from the task.
         relax_cell : bool
             Whether the cell shape/volume was allowed to change during the task.
         steps : int
             Maximum number of ionic steps allowed during relaxation.
         relax_kwargs : dict
-            Keyword arguments that will get passed to :obj:`StructOptimizer.relax`.
+            Keyword arguments that will get passed to :obj:`Relaxer.relax`.
         optimizer_kwargs : dict
-            Keyword arguments that will get passed to :obj:`StructOptimizer()`.
+            Keyword arguments that will get passed to :obj:`Relaxer()`.
         ionic_step_data : tuple
             Which data to save from each ionic step.
         """
@@ -175,11 +178,6 @@ def from_chgnet_result(
                 if "forces" in ionic_step_data
                 else None
             )
-            cur_magmoms = (
-                trajectory["magmoms"][i].tolist()
-                if "magmoms" in ionic_step_data
-                else None
-            )
             cur_stress = (
                 trajectory["stresses"][i].tolist()
                 if "stress" in ionic_step_data
@@ -196,15 +194,30 @@ def from_chgnet_result(
             else:
                 cur_structure = None
 
-            ionic_steps.append(
-                IonicStep(
+            # include "magmoms" in :obj:`cur_ionic_step` if the trajectory has "magmoms"
+            if "magmoms" in trajectory:
+                cur_ionic_step = IonicStep(
                     energy=cur_energy,
                     forces=cur_forces,
-                    magmoms=cur_magmoms,
+                    magmoms=(
+                        trajectory["magmoms"][i].tolist()
+                        if "magmoms" in ionic_step_data
+                        else None
+                    ),
                     stress=cur_stress,
                     structure=cur_structure,
                 )
-            )
+
+            # otherwise do not include "magmoms" in :obj:`cur_ionic_step`
+            elif "magmoms" not in trajectory.keys():
+                cur_ionic_step = IonicStep(
+                    energy=cur_energy,
+                    forces=cur_forces,
+                    stress=cur_stress,
+                    structure=cur_structure,
+                )
+
+            ionic_steps.append(cur_ionic_step)
 
         output_doc = OutputDoc(
             structure=output_structure,
@@ -216,15 +229,20 @@ def from_chgnet_result(
             n_steps=n_steps,
         )
 
-        import chgnet
+        if forcefield_name == "M3GNet":
+            import matgl
+
+            version = matgl.__version__
+        elif forcefield_name == "CHGNet":
+            import chgnet
 
-        version = chgnet.__version__
+            version = chgnet.__version__
 
         return cls.from_structure(
             meta_structure=output_structure,
             structure=output_structure,
             input=input_doc,
             output=output_doc,
-            forcefield_name="CHGNet",
+            forcefield_name=forcefield_name,
             forcefield_version=version,
         )

tests/forcefields/test_jobs.py

@@ -1,7 +1,7 @@
 from pytest import approx
 
 
-def test_static_maker(si_structure):
+def test_chgnet_static_maker(si_structure):
     from jobflow import run_locally
 
     from atomate2.forcefields.jobs import CHGNetStaticMaker
@@ -23,7 +23,7 @@ def test_static_maker(si_structure):
     assert output1.output.n_steps == 1
 
 
-def test_relax_maker(si_structure):
+def test_chgnet_relax_maker(si_structure):
     from jobflow import run_locally
 
     from atomate2.forcefields.jobs import CHGNetRelaxMaker
@@ -46,3 +46,46 @@ def test_relax_maker(si_structure):
         0.002112872898578644, rel=1e-4
     )
     assert output1.output.n_steps == 12
+
+
+def test_m3gnet_static_maker(si_structure):
+    from jobflow import run_locally
+
+    from atomate2.forcefields.jobs import M3GNetStaticMaker
+    from atomate2.forcefields.schemas import ForceFieldTaskDocument
+
+    task_doc_kwargs = {"ionic_step_data": ("structure", "energy")}
+
+    # generate job
+    job = M3GNetStaticMaker(task_document_kwargs=task_doc_kwargs).make(si_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(-10.711267471313477, rel=1e-4)
+    assert output1.output.n_steps == 1
+
+
+def test_m3gnet_relax_maker(si_structure):
+    from jobflow import run_locally
+
+    from atomate2.forcefields.jobs import M3GNetRelaxMaker
+    from atomate2.forcefields.schemas import ForceFieldTaskDocument
+
+    # translate one atom to ensure a small number of relaxation steps are taken
+    si_structure.translate_sites(0, [0, 0, 0.1])
+
+    # generate job
+    job = M3GNetRelaxMaker(steps=25).make(si_structure)
+
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(job, ensure_success=True)
+
+    # validating the outputs of the job
+    output1 = responses[job.uuid][1].output
+    assert isinstance(output1, ForceFieldTaskDocument)
+    assert output1.output.energy == approx(-10.710836410522461, rel=1e-4)
+    assert output1.output.n_steps == 14