Merge pull request materialsproject#581 from JaGeo/extend_elastic_wor…

…kflow

Move elastic workflow to common and build force-field elastic workflow

src/atomate2/common/flows/elastic.py

@@ -0,0 +1,166 @@
+"""Flows for calculating elastic constants."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import Flow, Maker, OnMissing
+from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
+
+from atomate2 import SETTINGS
+from atomate2.common.jobs.elastic import (
+    fit_elastic_tensor,
+    generate_elastic_deformations,
+    run_elastic_deformations,
+)
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from emmet.core.math import Matrix3D
+    from pymatgen.core.structure import Structure
+
+    from atomate2.forcefields.jobs import ForceFieldRelaxMaker
+    from atomate2.vasp.jobs.base import BaseVaspMaker
+
+
+@dataclass
+class BaseElasticMaker(Maker, ABC):
+    """
+    Maker to calculate elastic constants.
+
+    Calculate the elastic constant of a material. Initially, a tight structural
+    relaxation is performed to obtain the structure in a state of approximately zero
+    stress. Subsequently, perturbations are applied to the lattice vectors and the
+    resulting stress tensor is calculated from DFT, while allowing for relaxation of the
+    ionic degrees of freedom. Finally, constitutive relations from linear elasticity,
+    relating stress and strain, are employed to fit the full 6x6 elastic tensor. From
+    this, aggregate properties such as Voigt and Reuss bounds on the bulk and shear
+    moduli are derived.
+
+    .. Note::
+        It is heavily recommended to symmetrize the structure before passing it to
+        this flow. Otherwise, the symmetry reduction routines will not be as
+        effective at reducing the total number of deformations needed.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    order : int
+        Order of the tensor expansion to be determined. Can be either 2 or 3.
+    sym_reduce : bool
+        Whether to reduce the number of deformations using symmetry.
+    symprec : float
+        Symmetry precision to use in the reduction of symmetry.
+    bulk_relax_maker : .BaseVaspMaker or .ForceFieldRelaxMaker or None
+        A maker to perform a tight relaxation on the bulk. Set to ``None`` to skip the
+        bulk relaxation.
+    elastic_relax_maker : .BaseVaspMaker or .ForceFieldRelaxMaker
+        Maker used to generate elastic relaxations.
+    generate_elastic_deformations_kwargs : dict
+        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"
+    order: int = 2
+    sym_reduce: bool = True
+    symprec: float = SETTINGS.SYMPREC
+    bulk_relax_maker: BaseVaspMaker | ForceFieldRelaxMaker | None = None
+    elastic_relax_maker: BaseVaspMaker | ForceFieldRelaxMaker = (
+        None  # constant volume optimization
+    )
+    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,
+        structure: Structure,
+        prev_dir: str | Path | None = None,
+        equilibrium_stress: Matrix3D = None,
+        conventional: bool = False,
+    ) -> Flow:
+        """
+        Make flow to calculate the elastic constant.
+
+        Parameters
+        ----------
+        structure : .Structure
+            A pymatgen structure.
+        prev_vasp_dir : str or Path or None
+            A previous vasp calculation directory to use for copying outputs.
+        equilibrium_stress : tuple of tuple of float
+            The equilibrium stress of the (relaxed) structure, if known.
+        conventional : bool
+            Whether to transform the structure into the conventional cell.
+        """
+        jobs = []
+
+        if self.bulk_relax_maker is not None:
+            # optionally relax the structure
+            bulk_kwargs = {}
+            if self.prev_calc_dir_argname is not None:
+                bulk_kwargs[self.prev_calc_dir_argname] = prev_dir
+            bulk = self.bulk_relax_maker.make(structure, **bulk_kwargs)
+            jobs.append(bulk)
+            structure = bulk.output.structure
+            prev_dir = bulk.output.dir_name
+            if equilibrium_stress is None:
+                equilibrium_stress = bulk.output.output.stress
+
+        if conventional:
+            sga = SpacegroupAnalyzer(structure, symprec=self.symprec)
+            structure = sga.get_conventional_standard_structure()
+
+        deformations = generate_elastic_deformations(
+            structure,
+            order=self.order,
+            sym_reduce=self.sym_reduce,
+            symprec=self.symprec,
+            **self.generate_elastic_deformations_kwargs,
+        )
+
+        vasp_deformation_calcs = run_elastic_deformations(
+            structure,
+            deformations.output,
+            elastic_relax_maker=self.elastic_relax_maker,
+            prev_dir=prev_dir,
+        )
+        fit_tensor = fit_elastic_tensor(
+            structure,
+            vasp_deformation_calcs.output,
+            equilibrium_stress=equilibrium_stress,
+            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
+        fit_tensor.config.on_missing_references = OnMissing.NONE
+
+        jobs += [deformations, vasp_deformation_calcs, fit_tensor]
+
+        return Flow(
+            jobs=jobs,
+            output=fit_tensor.output,
+            name=self.name,
+        )
+
+    @property
+    @abstractmethod
+    def prev_calc_dir_argname(self):
+        """Name of argument informing static maker of previous calculation directory.
+
+        As this differs between different DFT codes (e.g., VASP, CP2K), it
+        has been left as a property to be implemented by the inheriting class.
+        Note: this is only applicable if a relax_maker is specified; i.e., two
+        calculations are performed for each ordering (relax -> static)
+        """

src/atomate2/common/jobs/elastic.py

@@ -0,0 +1,231 @@
+"""Jobs used in the calculation of elastic tensors."""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+from typing import TYPE_CHECKING
+
+import numpy as np
+from jobflow import Flow, Response, job
+from pymatgen.alchemy.materials import TransformedStructure
+from pymatgen.analysis.elasticity import Deformation, Strain, Stress
+from pymatgen.core.tensors import symmetry_reduce
+from pymatgen.transformations.standard_transformations import (
+    DeformStructureTransformation,
+)
+
+from atomate2 import SETTINGS
+from atomate2.common.analysis.elastic import get_default_strain_states
+from atomate2.common.schemas.elastic import ElasticDocument
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from emmet.core.math import Matrix3D
+    from pymatgen.core.structure import Structure
+
+    from atomate2.forcefields.jobs import ForceFieldRelaxMaker
+    from atomate2.vasp.jobs.base import BaseVaspMaker
+
+
+logger = logging.getLogger(__name__)
+
+
+@job
+def generate_elastic_deformations(
+    structure: Structure,
+    order: int = 2,
+    strain_states: list[tuple[int, int, int, int, int, int]] | None = None,
+    strain_magnitudes: list[float] | list[list[float]] | None = None,
+    symprec: float = SETTINGS.SYMPREC,
+    sym_reduce: bool = True,
+) -> list[Deformation]:
+    """
+    Generate elastic deformations.
+
+    Parameters
+    ----------
+    structure : Structure
+        A pymatgen structure object.
+    order : int
+        Order of the tensor expansion to be determined. Can be either 2 or 3.
+    strain_states : None or list of tuple of int
+        List of Voigt-notation strains, e.g. ``[(1, 0, 0, 0, 0, 0), (0, 1, 0, 0, 0, 0),
+        etc]``.
+    strain_magnitudes : None or list of float or list of list of float
+        A list of strain magnitudes to multiply by for each strain state, e.g. ``[-0.01,
+        -0.005, 0.005, 0.01]``. Alternatively, a list of lists can be specified, where
+        each inner list corresponds to a specific strain state.
+
+    symprec : float
+        Symmetry precision.
+    sym_reduce : bool
+        Whether to reduce the number of deformations using symmetry.
+
+    Returns
+    -------
+    List[Deformation]
+        A list of deformations.
+    """
+    if strain_states is None:
+        strain_states = get_default_strain_states(order)
+
+    if strain_magnitudes is None:
+        strain_magnitudes = np.linspace(-0.01, 0.01, 5 + (order - 2) * 2)
+
+    if np.array(strain_magnitudes).ndim == 1:
+        strain_magnitudes = [strain_magnitudes] * len(strain_states)  # type: ignore[assignment]
+
+    strains = []
+    for state, magnitudes in zip(strain_states, strain_magnitudes):
+        strains.extend([Strain.from_voigt(m * np.array(state)) for m in magnitudes])
+
+    # remove zero strains
+    strains = [strain for strain in strains if (abs(strain) > 1e-10).any()]
+
+    if np.linalg.matrix_rank([strain.voigt for strain in strains]) < 6:
+        # TODO: check for sufficiency of input for nth order
+        raise ValueError("strain list is insufficient to fit an elastic tensor")
+
+    if sym_reduce:
+        strain_mapping = symmetry_reduce(strains, structure, symprec=symprec)
+        logger.info(
+            f"Using symmetry to reduce number of strains from {len(strains)} to "
+            f"{len(list(strain_mapping.keys()))}"
+        )
+        strains = list(strain_mapping.keys())
+
+    return [s.get_deformation_matrix() for s in strains]
+
+
+@job
+def run_elastic_deformations(
+    structure: Structure,
+    deformations: list[Deformation],
+    prev_dir: str | Path | None = None,
+    prev_dir_argname: str = None,
+    elastic_relax_maker: BaseVaspMaker | ForceFieldRelaxMaker = None,
+) -> Response:
+    """
+    Run elastic deformations.
+
+    Note, this job will replace itself with N relaxation calculations, where N is
+    the number of deformations.
+
+    Parameters
+    ----------
+    structure : Structure
+        A pymatgen structure.
+    deformations : list of Deformation
+        The deformations to apply.
+    prev_dir : str or Path or None
+        A previous directory to use for copying outputs.
+    prev_dir_argname: str
+        argument name for the prev_dir variable
+    elastic_relax_maker : .BaseVaspMaker or .ForceFieldRelaxMaker
+        A VaspMaker or a ForceFieldMaker to use to generate the elastic relaxation jobs.
+    """
+    relaxations = []
+    outputs = []
+    for i, deformation in enumerate(deformations):
+        # deform the structure
+        dst = DeformStructureTransformation(deformation=deformation)
+        ts = TransformedStructure(structure, transformations=[dst])
+        deformed_structure = ts.final_structure
+
+        with contextlib.suppress(Exception):
+            # write details of the transformation to the transformations.json file
+            # this file will automatically get added to the task document and allow
+            # the elastic builder to reconstruct the elastic document; note the ":" is
+            # automatically converted to a "." in the filename.
+            elastic_relax_maker.write_additional_data["transformations:json"] = ts
+
+        elastic_job_kwargs = {}
+        if prev_dir is not None and prev_dir_argname is not None:
+            elastic_job_kwargs[prev_dir_argname] = prev_dir
+        # create the job
+        relax_job = elastic_relax_maker.make(deformed_structure, **elastic_job_kwargs)
+        relax_job.append_name(f" {i + 1}/{len(deformations)}")
+        relaxations.append(relax_job)
+
+        # extract the outputs we want
+        output = {
+            "stress": relax_job.output.output.stress,
+            "deformation": deformation,
+            "uuid": relax_job.output.uuid,
+            "job_dir": relax_job.output.dir_name,
+        }
+
+        outputs.append(output)
+
+    relax_flow = Flow(relaxations, outputs)
+    return Response(replace=relax_flow)
+
+
+@job(output_schema=ElasticDocument)
+def fit_elastic_tensor(
+    structure: Structure,
+    deformation_data: list[dict],
+    equilibrium_stress: Matrix3D | None = None,
+    order: int = 2,
+    fitting_method: str = SETTINGS.ELASTIC_FITTING_METHOD,
+    symprec: float = SETTINGS.SYMPREC,
+    allow_elastically_unstable_structs: bool = True,
+) -> ElasticDocument:
+    """
+    Analyze stress/strain data to fit the elastic tensor and related properties.
+
+    Parameters
+    ----------
+    structure : ~pymatgen.core.structure.Structure
+        A pymatgen structure.
+    deformation_data : list of dict
+        The deformation data, as a list of dictionaries, each containing the keys
+        "stress", "deformation".
+    equilibrium_stress : None or tuple of tuple of float
+        The equilibrium stress of the (relaxed) structure, if known.
+    order : int
+        Order of the tensor expansion to be fitted. Can be either 2 or 3.
+    fitting_method : str
+        The method used to fit the elastic tensor. See pymatgen for more details on the
+        methods themselves. The options are:
+
+        - "finite_difference" (note this is required if fitting a 3rd order tensor)
+        - "independent"
+        - "pseudoinverse"
+    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 = []
+    uuids = []
+    job_dirs = []
+    for data in deformation_data:
+        # stress could be none if the deformation calculation failed
+        if data["stress"] is None:
+            continue
+
+        stresses.append(Stress(data["stress"]))
+        deformations.append(Deformation(data["deformation"]))
+        uuids.append(data["uuid"])
+        job_dirs.append(data["job_dir"])
+
+    logger.info("Analyzing stress/strain data")
+
+    return ElasticDocument.from_stresses(
+        structure,
+        stresses,
+        deformations,
+        uuids,
+        job_dirs,
+        fitting_method=fitting_method,
+        order=order,
+        equilibrium_stress=equilibrium_stress,
+        symprec=symprec,
+        allow_elastically_unstable_structs=allow_elastically_unstable_structs,
+    )