Merge pull request #322 from matthewkuner/add_forcefields

Add forcefield schemas/makers to atomate2

pyproject.toml

@@ -43,6 +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"]
 docs = [
     "FireWorks==2.0.3",
     "autodoc_pydantic==1.8.0",
@@ -72,6 +73,7 @@ strict = [
     "pymatgen-analysis-defects==2023.5.8",
     "pymatgen==2023.5.10",
     "seekpath==2.1.0",
+    "chgnet==0.1.3"
 ]
 
 [project.scripts]

src/atomate2/forcefields/flows.py

@@ -0,0 +1,55 @@
+"""Flows to combine a force field relaxation with another job (e.g. DFT relaxation)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from jobflow import Flow, Maker
+from pymatgen.core.structure import Structure
+
+from atomate2.forcefields.jobs import CHGNetRelaxMaker
+from atomate2.vasp.jobs.base import BaseVaspMaker
+from atomate2.vasp.jobs.core import RelaxMaker
+
+
+@dataclass
+class CHGNetVaspRelaxMaker(Maker):
+    """
+    Maker to (pre)relax a structure using CHGNet and then run VASP.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flow produced by this maker.
+    chgnet_maker : .CHGNetRelaxMaker
+        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"
+    chgnet_maker: CHGNetRelaxMaker = field(default_factory=CHGNetRelaxMaker)
+    vasp_maker: BaseVaspMaker = field(default_factory=RelaxMaker)
+
+    def make(self, structure: Structure):
+        """
+        Create a flow with a CHGNet (pre)relaxation followed by a VASP relaxation.
+
+        Parameters
+        ----------
+        structure : .Structure
+            A pymatgen structure.
+
+        Returns
+        -------
+        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)

src/atomate2/forcefields/jobs.py

@@ -0,0 +1,105 @@
+"""Job to relax a structure using a force field (aka an interatomic potential)."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+from jobflow import Maker, job
+from pymatgen.core.structure import Structure
+
+from atomate2.forcefields.schemas import ForceFieldTaskDocument
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CHGNetRelaxMaker(Maker):
+    """
+    Maker to perform a relaxation using the CHGNet 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:`StructOptimizer.relax`.
+    optimizer_kwargs : dict
+        Keyword arguments that will get passed to :obj:`StructOptimizer()`.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+    """
+
+    name: str = "CHGNet 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 CHGNet.
+
+        Parameters
+        ----------
+        structure: .Structure
+            A pymatgen structure.
+        """
+        from chgnet.model import StructOptimizer
+
+        if self.steps < 0:
+            logger.warning(
+                "WARNING: A negative number of steps is not possible. "
+                "Behavior may vary..."
+            )
+
+        relaxer = StructOptimizer(**self.optimizer_kwargs)
+        result = relaxer.relax(
+            structure, relax_cell=self.relax_cell, steps=self.steps, **self.relax_kwargs
+        )
+
+        ff_task_doc = ForceFieldTaskDocument.from_chgnet_result(
+            result,
+            self.relax_cell,
+            self.steps,
+            self.relax_kwargs,
+            self.optimizer_kwargs,
+            **self.task_document_kwargs,
+        )
+
+        return ff_task_doc
+
+
+@dataclass
+class CHGNetStaticMaker(CHGNetRelaxMaker):
+    """
+    Maker to calculate forces and stresses using the CHGNet 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:`StructOptimizer.relax`.
+    optimizer_kwargs : dict
+        Keyword arguments that will get passed to :obj:`StructOptimizer()`.
+    task_document_kwargs : dict
+        Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`.
+    """
+
+    name: str = "CHGNet 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)