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",
@@ -66,6 +66,7 @@ strict = [
     "emmet-core==0.55.2",
     "jobflow==0.1.11",
     "lobsterpy==0.2.9",
+    "matgl==0.5.6",
     "monty==2023.5.8",
     "mp-api==0.33.3",
     "numpy",

src/atomate2/common/schemas/elastic.py

@@ -148,6 +148,7 @@ def from_stresses(
         order: Optional[int] = None,
         equilibrium_stress: Optional[Matrix3D] = None,
         symprec: float = SETTINGS.SYMPREC,
+        allow_elastically_unstable_structs: bool = True,
     ):
         """
         Create an elastic document from strains and stresses.
@@ -177,6 +178,9 @@ def from_stresses(
         symprec : float
             Symmetry precision for deriving symmetry equivalent deformations. If
             ``symprec=None``, then no symmetry operations will be applied.
+        allow_elastically_unstable_structs : bool
+            Whether to allow the ElasticDocument to still complete in the event that
+            the structure is elastically unstable.
         """
         if symprec is not None:
             deformations, stresses, uuids, job_dirs = _expand_deformations(
@@ -212,7 +216,12 @@ def from_stresses(
 
         ieee = result.convert_to_ieee(structure)
         property_tensor = ieee if order == 2 else ElasticTensor(ieee[0])
-        property_dict = property_tensor.get_structure_property_dict(structure)
+
+        ignore_errors = bool(allow_elastically_unstable_structs)
+        property_dict = property_tensor.get_structure_property_dict(
+            structure, ignore_errors=ignore_errors
+        )
+
         derived_properties = DerivedProperties(**property_dict)
 
         eq_stress = eq_stress.tolist() if eq_stress is not None else eq_stress

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,
         )

src/atomate2/vasp/flows/elastic.py

@@ -66,6 +66,8 @@ class ElasticMaker(Maker):
         Keyword arguments passed to :obj:`generate_elastic_deformations`.
     fit_elastic_tensor_kwargs : dict
         Keyword arguments passed to :obj:`fit_elastic_tensor`.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ElasticDocument.from_stresses()`.
     """
 
     name: str = "elastic"
@@ -78,6 +80,7 @@ class ElasticMaker(Maker):
     elastic_relax_maker: BaseVaspMaker = field(default_factory=ElasticRelaxMaker)
     generate_elastic_deformations_kwargs: dict = field(default_factory=dict)
     fit_elastic_tensor_kwargs: dict = field(default_factory=dict)
+    task_document_kwargs: dict = field(default_factory=dict)
 
     def make(
         self,
@@ -128,6 +131,7 @@ def make(
             order=self.order,
             symprec=self.symprec if self.sym_reduce else None,
             **self.fit_elastic_tensor_kwargs,
+            **self.task_document_kwargs,
         )
 
         # allow some of the deformations to fail

src/atomate2/vasp/jobs/elastic.py

@@ -240,6 +240,7 @@ def fit_elastic_tensor(
     order: int = 2,
     fitting_method: str = SETTINGS.ELASTIC_FITTING_METHOD,
     symprec: float = SETTINGS.SYMPREC,
+    allow_elastically_unstable_structs: bool = True,
 ):
     """
     Analyze stress/strain data to fit the elastic tensor and related properties.
@@ -265,6 +266,9 @@ def fit_elastic_tensor(
     symprec : float
         Symmetry precision for deriving symmetry equivalent deformations. If
         ``symprec=None``, then no symmetry operations will be applied.
+    allow_elastically_unstable_structs : bool
+        Whether to allow the ElasticDocument to still complete in the event that
+        the structure is elastically unstable.
     """
     stresses = []
     deformations = []
@@ -292,4 +296,5 @@ def fit_elastic_tensor(
         order=order,
         equilibrium_stress=equilibrium_stress,
         symprec=symprec,
+        allow_elastically_unstable_structs=allow_elastically_unstable_structs,
     )