diff --git a/pyproject.toml b/pyproject.toml index b3b4c230e7..c4981a76fa 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/atomate2/common/schemas/elastic.py b/src/atomate2/common/schemas/elastic.py index 49b8e37d76..bca7deb6cb 100644 --- a/src/atomate2/common/schemas/elastic.py +++ b/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 diff --git a/src/atomate2/forcefields/flows.py b/src/atomate2/forcefields/flows.py index de8ef902ab..b35da50622 100644 --- a/src/atomate2/forcefields/flows.py +++ b/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,7 +48,6 @@ 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" @@ -57,3 +55,44 @@ def make(self, structure: Structure): 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) diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index 1e9c85ed43..99ab078363 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/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) diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index 7ada2b2b0f..56f2e022c9 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/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,10 +108,12 @@ 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 @@ -118,9 +121,9 @@ def from_chgnet_result( 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, ) diff --git a/src/atomate2/vasp/flows/elastic.py b/src/atomate2/vasp/flows/elastic.py index b4c2c6a636..3235605ad9 100644 --- a/src/atomate2/vasp/flows/elastic.py +++ b/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 diff --git a/src/atomate2/vasp/jobs/elastic.py b/src/atomate2/vasp/jobs/elastic.py index 0e9e47100b..9a34f7001b 100644 --- a/src/atomate2/vasp/jobs/elastic.py +++ b/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, ) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index fe209d7978..1502c9138b 100644 --- a/tests/forcefields/test_jobs.py +++ b/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