diff --git a/src/quacc/atoms/deformation.py b/src/quacc/atoms/deformation.py index 1bb53febb2..5df0160b24 100644 --- a/src/quacc/atoms/deformation.py +++ b/src/quacc/atoms/deformation.py @@ -18,7 +18,7 @@ def make_deformations_from_bulk( norm_strains: Sequence[float] = (-0.01, -0.005, 0.005, 0.01), shear_strains: Sequence[float] = (-0.06, -0.03, 0.03, 0.06), symmetry: bool = False, -) -> list[Atoms]: +) -> DeformedStructureSet: """ Function to generate deformed structures from a bulk atoms object. @@ -38,13 +38,11 @@ def make_deformations_from_bulk( list[Atoms] All generated deformed structures """ - struct = AseAtomsAdaptor.get_structure(atoms) + struct = AseAtomsAdaptor.get_structure(atoms) # type: ignore - deformed_set = DeformedStructureSet( + return DeformedStructureSet( struct, norm_strains=norm_strains, shear_strains=shear_strains, symmetry=symmetry, ) - - return [structure.to_ase_atoms() for structure in deformed_set] diff --git a/src/quacc/recipes/common/elastic.py b/src/quacc/recipes/common/elastic.py index 33d5535cd5..ebeb434106 100644 --- a/src/quacc/recipes/common/elastic.py +++ b/src/quacc/recipes/common/elastic.py @@ -4,24 +4,66 @@ from typing import TYPE_CHECKING -from quacc import subflow +from ase import units +from ase.stress import voigt_6_to_full_3x3_stress +from emmet.core.elasticity import ElasticityDoc +from emmet.core.mpid import MPID +from pymatgen.analysis.elasticity.stress import Stress +from pymatgen.io.ase import AseAtomsAdaptor + +from quacc import job, subflow from quacc.atoms.deformation import make_deformations_from_bulk if TYPE_CHECKING: from typing import Any - from ase.atoms import Atoms + from pymatgen.analysis.elasticity.strain import DeformedStructureSet from quacc import Job + from quacc.types import ElasticSchema, OptSchema, RunSchema + + +@job +def deformations_to_elastic_tensor( + undeformed_result: OptSchema | RunSchema, + deformed_structure_set: DeformedStructureSet, + results: list[dict], +) -> ElasticityDoc: + structure = AseAtomsAdaptor.get_structure(undeformed_result["atoms"]) # type: ignore + return ElasticityDoc.from_deformations_and_stresses( + structure, + material_id=MPID("quacc-00"), + deformations=deformed_structure_set.deformations, + equilibrium_stress=Stress( + ( + voigt_6_to_full_3x3_stress(undeformed_result["results"]["stress"]) + if len(undeformed_result["results"]["stress"]) == 6 + else undeformed_result["results"]["stress"] + ) + / units.GPa + ), + stresses=[ + Stress( + ( + voigt_6_to_full_3x3_stress(relax_result["results"]["stress"]) + if len(relax_result["results"]["stress"]) == 6 + else relax_result["results"]["stress"] + ) + / units.GPa + ) + for relax_result in results + ], + ) @subflow def bulk_to_deformations_subflow( - atoms: Atoms, + undeformed_result: OptSchema | RunSchema, relax_job: Job, - static_job: Job | None = None, + static_job: Job, + run_static: bool = False, deform_kwargs: dict[str, Any] | None = None, -) -> list[dict]: +) -> ElasticSchema: """ Workflow consisting of: @@ -33,10 +75,12 @@ def bulk_to_deformations_subflow( Parameters ---------- - atoms - Atoms object + undeformed_result + Result of a static or optimization calculation relax_job The relaxation function. + static_job + The static function static_job The static function. deform_kwargs @@ -50,15 +94,28 @@ def bulk_to_deformations_subflow( """ deform_kwargs = deform_kwargs or {} - deformations = make_deformations_from_bulk(atoms, **deform_kwargs) + deformed_structure_set = make_deformations_from_bulk( + undeformed_result["atoms"], **deform_kwargs + ) results = [] - for deformed in deformations: - result = relax_job(deformed) + for deformed in deformed_structure_set: + result = relax_job(deformed.to_ase_atoms()) - if static_job is not None: + if run_static: result = static_job(result["atoms"]) results.append(result) - return results + elasticity_doc = deformations_to_elastic_tensor( + undeformed_result=undeformed_result, + deformed_structure_set=deformed_structure_set, + results=results, + ) + + return { + "deformed_structure_set": deformed_structure_set, + "deformed_results": results, + "undeformed_result": undeformed_result, + "elasticity_doc": elasticity_doc, + } diff --git a/src/quacc/recipes/emt/elastic.py b/src/quacc/recipes/emt/elastic.py index 3f2e9dbbb1..65721a6610 100644 --- a/src/quacc/recipes/emt/elastic.py +++ b/src/quacc/recipes/emt/elastic.py @@ -15,17 +15,18 @@ from ase.atoms import Atoms - from quacc.types import OptSchema, RunSchema + from quacc.types import ElasticSchema @flow def bulk_to_deformations_flow( atoms: Atoms, - run_static: bool = True, + run_static: bool = False, + pre_relax: bool = True, deform_kwargs: dict[str, Any] | None = None, job_params: dict[str, dict[str, Any]] | None = None, job_decorators: dict[str, Callable | None] | None = None, -) -> list[RunSchema | OptSchema]: +) -> ElasticSchema: """ Workflow consisting of: @@ -66,11 +67,17 @@ def bulk_to_deformations_flow( [relax_job, static_job], param_swaps=job_params, decorators=job_decorators, - ) + ) # type: ignore + + if pre_relax: + undeformed_result = relax_job_(atoms, relax_cell=True) + else: + undeformed_result = static_job_(atoms) return bulk_to_deformations_subflow( - atoms, - relax_job_, - static_job=static_job_ if run_static else None, + undeformed_result=undeformed_result, + relax_job=relax_job_, + static_job=static_job_, + run_static=run_static, deform_kwargs=deform_kwargs, ) diff --git a/src/quacc/recipes/mlp/elastic.py b/src/quacc/recipes/mlp/elastic.py new file mode 100644 index 0000000000..c9290ec829 --- /dev/null +++ b/src/quacc/recipes/mlp/elastic.py @@ -0,0 +1,85 @@ +"""Elastic constants recipes for MLPs.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from quacc import flow +from quacc.recipes.common.elastic import bulk_to_deformations_subflow +from quacc.recipes.mlp.core import relax_job, static_job +from quacc.wflow_tools.customizers import customize_funcs + +if TYPE_CHECKING: + from collections.abc import Callable + from typing import Any + + from ase.atoms import Atoms + + from quacc.types import ElasticSchema + + +@flow +def bulk_to_deformations_flow( + atoms: Atoms, + run_static: bool = False, + pre_relax: bool = True, + deform_kwargs: dict[str, Any] | None = None, + job_params: dict[str, dict[str, Any]] | None = None, + job_decorators: dict[str, Callable | None] | None = None, +) -> ElasticSchema: + """ + Workflow consisting of: + + 1. Deformed structures generation + + 2. Deformed structures relaxations + - name: "relax_job" + - job: [quacc.recipes.mlp.core.relax_job][] + + 3. Deformed structures statics (optional) + - name: "static_job" + - job: [quacc.recipes.mlp.core.static_job][] + + Parameters + ---------- + atoms + Atoms object + run_static + Whether to run static calculations after the relaxations + pre_relax + Whether to pre-relax the input atoms as is common + deform_kwargs + Additional keyword arguments to pass to [quacc.atoms.deformation.make_deformations_from_bulk][] + job_params + Custom parameters to pass to each Job in the Flow. This is a dictionary where + the keys are the names of the jobs and the values are dictionaries of parameters. + job_decorators + Custom decorators to apply to each Job in the Flow. This is a dictionary where + the keys are the names of the jobs and the values are decorators. + + Returns + ------- + list[RunSchema | OptSchema] + [RunSchema][quacc.schemas.ase.Summarize.run] or + [OptSchema][quacc.schemas.ase.Summarize.opt] for each deformation. + See the return type-hint for the data structure. + """ + relax_job_, static_job_ = customize_funcs( + ["relax_job", "static_job"], + [relax_job, static_job], + param_swaps=job_params, + decorators=job_decorators, + ) # type: ignore + + if pre_relax: + undeformed_result = relax_job_(atoms, relax_cell=True) + else: + undeformed_result = static_job_(atoms) + + return bulk_to_deformations_subflow( + undeformed_result=undeformed_result, + relax_job=relax_job_, + static_job=static_job_, + run_static=run_static, + deform_kwargs=deform_kwargs, + ) diff --git a/src/quacc/types.py b/src/quacc/types.py index e75dc01c53..dfec953aab 100644 --- a/src/quacc/types.py +++ b/src/quacc/types.py @@ -24,6 +24,7 @@ class DefaultSetting(BaseSettings): from ase.atoms import Atoms from ase.md.md import MolecularDynamics from ase.optimize.optimize import Dynamics + from emmet.core.elasticity import ElasticityDoc from emmet.core.math import ListMatrix3D, Matrix3D, Vector3D from emmet.core.symmetry import CrystalSystem from emmet.core.vasp.calc_types import CalcType @@ -32,6 +33,7 @@ class DefaultSetting(BaseSettings): from emmet.core.vasp.task_valid import TaskState from numpy.random import Generator from numpy.typing import ArrayLike, NDArray + from pymatgen.analysis.elasticity.strain import DeformedStructureSet from pymatgen.core.composition import Composition from pymatgen.core.lattice import Lattice from pymatgen.core.periodic_table import Element @@ -528,6 +530,12 @@ class ThermoSchema(AtomsSchema): parameters_thermo: ParametersThermo results: ThermoResults + class ElasticSchema(TypedDict): + deformed_structure_set: DeformedStructureSet + deformed_results: list[RunSchema | OptSchema] + undeformed_result: RunSchema | OptSchema + elasticity_doc: ElasticityDoc + class VibThermoSchema(VibSchema, ThermoSchema): """Combined Vibrations and Thermo schema""" diff --git a/tests/core/atoms/test_elastic.py b/tests/core/atoms/test_elastic.py index 8a04e78604..d522ca64ce 100644 --- a/tests/core/atoms/test_elastic.py +++ b/tests/core/atoms/test_elastic.py @@ -16,8 +16,13 @@ def test_make_deformations_from_bulk(): atoms.info["test"] = "hi" deformations = make_deformations_from_bulk(atoms) assert len(deformations) == 24 - assert deformations[0].get_volume() != pytest.approx(atoms.get_volume()) + assert deformations[0].to_ase_atoms().get_volume() != pytest.approx( + atoms.get_volume() + ) for deformation in deformations: - assert_equal(deformation.get_atomic_numbers(), [30, 30, 30, 30, 52, 52, 52, 52]) - assert_equal(deformation.get_chemical_formula(), "Te4Zn4") - assert deformation.info["test"] == "hi" + assert_equal( + deformation.to_ase_atoms().get_atomic_numbers(), + [30, 30, 30, 30, 52, 52, 52, 52], + ) + assert_equal(deformation.to_ase_atoms().get_chemical_formula(), "Te4Zn4") + assert deformation.to_ase_atoms().info["test"] == "hi" diff --git a/tests/core/recipes/emt_recipes/test_emt_elastic.py b/tests/core/recipes/emt_recipes/test_emt_elastic.py index 43d7a76ae3..154769cd60 100644 --- a/tests/core/recipes/emt_recipes/test_emt_elastic.py +++ b/tests/core/recipes/emt_recipes/test_emt_elastic.py @@ -12,21 +12,29 @@ def test_elastic_jobs(tmp_path, monkeypatch): atoms = bulk("Cu") outputs = bulk_to_deformations_flow(atoms, run_static=False) - assert outputs[0]["atoms"].get_volume() != pytest.approx(atoms.get_volume()) - for output in outputs: + assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx( + atoms.get_volume() + ) + assert outputs["elasticity_doc"].bulk_modulus.voigt == pytest.approx(134.579) + for output in outputs["deformed_results"]: assert output["parameters"]["asap_cutoff"] is False assert output["name"] == "EMT Relax" assert output["nelements"] == 1 assert output["nsites"] == 1 - assert len(outputs) == 24 + assert len(outputs["deformed_results"]) == 24 outputs = bulk_to_deformations_flow( atoms, run_static=True, job_params={"static_job": {"asap_cutoff": True}} ) - assert outputs[0]["atoms"].get_volume() != pytest.approx(atoms.get_volume()) - for output in outputs: + assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx( + atoms.get_volume() + ) + assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx( + atoms.get_volume() + ) + for output in outputs["deformed_results"]: assert output["parameters"]["asap_cutoff"] is True assert output["name"] == "EMT Static" assert output["nelements"] == 1 assert output["nsites"] == 1 - assert len(outputs) == 24 + assert len(outputs["deformed_results"]) == 24 diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py index 210569e9f8..750b995bd2 100644 --- a/tests/core/recipes/mlp_recipes/test_core_recipes.py +++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py @@ -31,6 +31,12 @@ methods.append("fairchem") +def _set_dtype(size, type_="float"): + globals()[f"{type_}_th"] = getattr(torch, f"{type_}{size}") + globals()[f"{type_}_np"] = getattr(np, f"{type_}{size}") + torch.set_default_dtype(getattr(torch, f"float{size}")) + + @pytest.mark.skipif(has_chgnet is None, reason="chgnet not installed") def test_bad_method(): atoms = bulk("Cu") @@ -38,12 +44,6 @@ def test_bad_method(): static_job(atoms, method="bad_method") -def _set_dtype(size, type_="float"): - globals()[f"{type_}_th"] = getattr(torch, f"{type_}{size}") - globals()[f"{type_}_np"] = getattr(np, f"{type_}{size}") - torch.set_default_dtype(getattr(torch, f"float{size}")) - - @pytest.mark.parametrize("method", methods) def test_static_job(tmp_path, monkeypatch, method): monkeypatch.chdir(tmp_path) diff --git a/tests/core/recipes/mlp_recipes/test_elastic_recipes.py b/tests/core/recipes/mlp_recipes/test_elastic_recipes.py new file mode 100644 index 0000000000..61344c2624 --- /dev/null +++ b/tests/core/recipes/mlp_recipes/test_elastic_recipes.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +from ase.build import bulk + +from quacc.recipes.mlp.elastic import bulk_to_deformations_flow + +torch = pytest.importorskip("torch") + +from importlib.util import find_spec + +methods = [] +if has_mace := find_spec("mace"): + methods.append("mace-mp-0") + +if has_matgl := find_spec("matgl"): + methods.append("m3gnet") + +if has_chgnet := find_spec("chgnet"): + methods.append("chgnet") + +if has_sevennet := find_spec("sevenn"): + methods.append("sevennet") + +if has_orb := find_spec("orb_models"): + methods.append("orb") + +if has_fairchem := find_spec("fairchem"): + methods.append("fairchem") + + +def _set_dtype(size, type_="float"): + globals()[f"{type_}_th"] = getattr(torch, f"{type_}{size}") + globals()[f"{type_}_np"] = getattr(np, f"{type_}{size}") + torch.set_default_dtype(getattr(torch, f"float{size}")) + + +@pytest.mark.parametrize("method", methods) +def test_elastic_jobs(tmp_path, monkeypatch, method): + monkeypatch.chdir(tmp_path) + + if method == "mace-mp-0": + _set_dtype(64) + else: + _set_dtype(32) + + if method == "fairchem": + calc_kwargs = { + "checkpoint_path": Path(__file__).parent / "eqV2_31M_omat_mp_salex.pt", + "seed": 0, + "disable_amp": True, + } + else: + calc_kwargs = {} + + ref_elastic_modulus = { + "chgnet": 199, + "m3gnet": 109.369, + "mace-mp-0": 130.727, + "sevennet": 142.296, + "orb": 187.107, + "fairchem": 105, + } + + atoms = bulk("Cu") + + outputs = bulk_to_deformations_flow( + atoms, + run_static=False, + pre_relax=True, + job_params={ + "all": {"method": method, **calc_kwargs}, + "relax_job": {"opt_params": {"fmax": 0.01}}, + }, + ) + assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx( + atoms.get_volume() + ) + assert outputs["undeformed_result"]["results"]["stress"] == pytest.approx( + 0, abs=1e-2 + ) + assert outputs["elasticity_doc"].bulk_modulus.voigt == pytest.approx( + ref_elastic_modulus[method], abs=2 + ) + for output in outputs["deformed_results"]: + assert output["nelements"] == 1 + assert output["nsites"] == 1 + assert len(outputs["deformed_results"]) == 24 + + outputs = bulk_to_deformations_flow( + atoms, + run_static=True, + pre_relax=True, + job_params={ + "all": {"method": method, **calc_kwargs}, + "relax_job": {"opt_params": {"fmax": 0.01}}, + }, + ) + assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx( + atoms.get_volume() + ) + + for output in outputs["deformed_results"]: + assert output["nelements"] == 1 + assert output["nsites"] == 1 + assert len(outputs["deformed_results"]) == 24