Ferroelectric Workflow 2 (#1012)

* start writing fe wflow

* added makers and analysis job

* added schema for polarization data

* order interpolation outputs

* fix typing

* fix nimages

* fix polarization_analysis variable

* fix polarization_analysis variable

* structure interpolation as flow

* fixing the connections between jobs, testing replace arg in Responce

* TaskDocument to dict

* Fixing polarization analysis and document

* tests added; fix PolarizationDocument; some doc

* some doc lines

* pre-commit fixes

* TaskDocument to TaskDoc

* more more little improvements

* pre-commit fix

* suggestions from Alex implemented

* uuid, job_dirs added to Pol Doc

* update tests

* create output dict with uuid outside pol_analysis job

* fix typo

* remove kspacing from incar, add kpoints test inputs

* start writing fe wflow

* syncing to recent upstream

* added schema for polarization data

* order interpolation outputs

* fix typing

* fix nimages

* fix polarization_analysis variable

* fix polarization_analysis variable

* structure interpolation as flow

* fixing the connections between jobs, testing replace arg in Responce

* TaskDocument to dict

* Fixing polarization analysis and document

* tests added; fix PolarizationDocument; some doc

* more more little improvements

* some doc lines

* pre-commit fixes

* TaskDocument to TaskDoc

* pre-commit fix

* suggestions from Alex implemented

* uuid, job_dirs added to Pol Doc

* update tests

* create output dict with uuid outside pol_analysis job

* fix typo

* remove kspacing from incar, add kpoints test inputs

* fix mistake in merging

* some manual typing fix

* fix prev_dir in BaseVaspMaker.make()

* fix test

* passing interp_structures instead of a path

* fixing variable name in flow

* added write_additional_data; nimages correspond to # of interp structs

* fixed write_additional_data file

* update pol norm value

* regenerate test files

* update energy value

* additional tests

.gitignore

@@ -72,3 +72,4 @@ docs/reference/atomate2.*
 *.doctrees*
 
 .ipynb_checkpoints
+.aider*

src/atomate2/vasp/flows/ferroelectric.py

@@ -0,0 +1,135 @@
+"""Flows for calculating the polarization of a polar material."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import Flow, Maker
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from pymatgen.core.structure import Structure
+
+    from atomate2.vasp.jobs.base import BaseVaspMaker
+
+from atomate2.vasp.flows.core import DoubleRelaxMaker
+from atomate2.vasp.jobs.core import PolarizationMaker, RelaxMaker
+from atomate2.vasp.jobs.ferroelectric import (
+    add_interpolation_flow,
+    get_polarization_output,
+    interpolate_structures,
+    polarization_analysis,
+)
+
+__all__ = ["FerroelectricMaker"]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FerroelectricMaker(Maker):
+    """
+    Maker to calculate polarization of a polar material.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    nimages: int
+        Number of interpolated structures calculated from polar to nonpolar structures
+    relax_maker: BaseVaspMaker or None or tuple
+        None to avoid relaxation of both polar and nonpolar structures
+        BaseVaspMaker to relax both structures (default)
+        tuple of BaseVaspMaker and None to control relaxation for each structure
+    lcalcpol_maker: BaseVaspMaker
+       Vasp maker to compute the polarization of each structure
+    """
+
+    name: str = "ferroelectric"
+    nimages: int = 8
+    relax_maker: BaseVaspMaker | None | tuple = field(
+        default_factory=lambda: DoubleRelaxMaker.from_relax_maker(RelaxMaker())
+    )
+    lcalcpol_maker: BaseVaspMaker = field(default_factory=PolarizationMaker)
+
+    def make(
+        self,
+        polar_structure: Structure,
+        nonpolar_structure: Structure,
+        prev_vasp_dir: str | Path | None = None,
+    ) -> Flow:
+        """
+        Make flow to calculate the polarization.
+
+        Parameters
+        ----------
+        polar_structure : .Structure
+            A pymatgen structure of the polar phase.
+        nonpolar_structure : .Structure
+            A pymatgen structure of the nonpolar phase.
+        prev_vasp_dir : str or Path or None
+            A previous vasp calculation directory to use for copying outputs.
+        """
+        jobs = []
+        prev_vasp_dir_p, prev_vasp_dir_np = None, None
+
+        if not isinstance(self.relax_maker, tuple):
+            self.relax_maker = (self.relax_maker, self.relax_maker)
+
+        if self.relax_maker[0]:
+            # optionally relax the polar structure
+            relax_p = self.relax_maker[0].make(polar_structure)
+            relax_p.append_name(" polar")
+            jobs.append(relax_p)
+            polar_structure = relax_p.output.structure
+            prev_vasp_dir_p = relax_p.output.dir_name
+
+        logger.info(f"{type(polar_structure)}")
+
+        polar_lcalcpol = self.lcalcpol_maker.make(
+            polar_structure, prev_dir=prev_vasp_dir_p
+        )
+        polar_lcalcpol.append_name(" polar")
+        jobs.append(polar_lcalcpol)
+        polar_structure = polar_lcalcpol.output.structure
+
+        if self.relax_maker[1]:
+            # optionally relax the nonpolar structure
+            relax_np = self.relax_maker[1].make(nonpolar_structure)
+            relax_np.append_name(" nonpolar")
+            jobs.append(relax_np)
+            nonpolar_structure = relax_np.output.structure
+            prev_vasp_dir_np = relax_np.output.dir_name
+
+        nonpolar_lcalcpol = self.lcalcpol_maker.make(
+            nonpolar_structure, prev_dir=prev_vasp_dir_np
+        )
+        nonpolar_lcalcpol.append_name(" nonpolar")
+        jobs.append(nonpolar_lcalcpol)
+        nonpolar_structure = nonpolar_lcalcpol.output.structure
+
+        interp_structs_job = interpolate_structures(
+            polar_structure, nonpolar_structure, self.nimages
+        )
+        jobs.append(interp_structs_job)
+
+        interp_structs = interp_structs_job.output
+        add_interp_flow = add_interpolation_flow(interp_structs, self.lcalcpol_maker)
+
+        pol_analysis = polarization_analysis(
+            get_polarization_output(nonpolar_lcalcpol),
+            get_polarization_output(polar_lcalcpol),
+            add_interp_flow.output,
+        )
+
+        jobs.append(add_interp_flow)
+        jobs.append(pol_analysis)
+
+        return Flow(
+            jobs=jobs,
+            output=pol_analysis.output,
+            name=self.name,
+        )

src/atomate2/vasp/jobs/core.py

@@ -533,6 +533,45 @@ class DielectricMaker(BaseVaspMaker):
     )
 
 
+@dataclass
+class PolarizationMaker(BaseVaspMaker):
+    """
+    Maker to create polarization calculation VASP jobs.
+
+    .. Note::
+        If starting from a previous calculation, magnetism will be disabled if all
+        MAGMOMs are less than 0.02.
+
+    Parameters
+    ----------
+    name : str
+        The job name.
+    input_set_generator : .StaticSetGenerator
+        A generator used to make the input set.
+    write_input_set_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.write_vasp_input_set`.
+    copy_vasp_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`.
+    run_vasp_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.run_vasp`.
+    task_document_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`.
+    stop_children_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.should_stop_children`.
+    write_additional_data : dict
+        Additional data to write to the current directory. Given as a dict of
+        {filename: data}. Note that if using FireWorks, dictionary keys cannot contain
+        the "." character which is typically used to denote file extensions. To avoid
+        this, use the ":" character, which will automatically be converted to ".". E.g.
+        ``{"my_file:txt": "contents of the file"}``.
+    """
+
+    name: str = "polarization"
+    input_set_generator: StaticSetGenerator = field(
+        default_factory=lambda: StaticSetGenerator(lcalcpol=True, auto_ispin=True)
+    )
+
+
 @dataclass
 class TransmuterMaker(BaseVaspMaker):
     """

src/atomate2/vasp/jobs/ferroelectric.py

@@ -0,0 +1,183 @@
+"""Job used in the Ferroelectric wflow."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from jobflow import Flow, Job, Response, job
+from pymatgen.analysis.ferroelectricity.polarization import get_total_ionic_dipole
+
+from atomate2.vasp.schemas.ferroelectric import PolarizationDocument
+
+if TYPE_CHECKING:
+    from pymatgen.core.structure import Structure
+
+    from atomate2.vasp.jobs.base import BaseVaspMaker
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["polarization_analysis"]
+
+
+@job(output_schema=PolarizationDocument)
+def polarization_analysis(
+    np_lcalcpol_output: dict[str, Any],
+    p_lcalcpol_output: dict[str, Any],
+    interp_lcalcpol_outputs: dict[str, Any],
+) -> PolarizationDocument:
+    """
+    Recover the same branch polarization and the spontaneous polarization.
+
+    Parameters
+    ----------
+    np_lcalcpol_output : dict
+        Output from previous nonpolar lcalcpol job.
+    p_lcalcpol_output : dict
+        Output from previous polar lcalcpol job.
+    interp_lcalcpol_outputs : dict
+        Output from previous interpolation lcalcpol jobs.
+
+    Returns
+    -------
+    PolarizationDocument
+        Document containing the polarization analysis results.
+
+    """
+    # order previous calculations from nonpolar to polar
+    ordered_keys = [
+        f"interpolation_{i}" for i in reversed(range(len(interp_lcalcpol_outputs)))
+    ]
+
+    polarization_tasks = [np_lcalcpol_output]
+    polarization_tasks += [interp_lcalcpol_outputs[k] for k in ordered_keys]
+    polarization_tasks += [p_lcalcpol_output]
+
+    task_lbls = []
+    structures = []
+    energies_per_atom = []
+    energies = []
+    job_dirs = []
+    uuids = []
+
+    for i, p in enumerate(polarization_tasks):
+        energies_per_atom.append(p["energy_per_atom"])
+        energies.append(p["energy"])
+        task_lbls.append(p["task_label"] or str(i))
+        structures.append(p["structure"])
+        job_dirs.append(p["job_dir"])
+        uuids.append(p["uuid"])
+
+    # If LCALCPOL = True then Outcar will parse and store the pseudopotential zvals.
+    zval_dict = p["zval_dict"]
+
+    # Assumes that we want to calculate the ionic contribution to the dipole moment.
+    # VASP's ionic contribution is sometimes strange.
+    # See pymatgen.analysis.ferroelectricity.polarization.Polarization for details.
+    p_elecs = [p["p_elecs"] for p in polarization_tasks]
+    p_ions = [get_total_ionic_dipole(st, zval_dict) for st in structures]
+
+    return PolarizationDocument.from_pol_output(
+        p_elecs,
+        p_ions,
+        structures,
+        energies,
+        energies_per_atom,
+        zval_dict,
+        task_lbls,
+        job_dirs,
+        uuids,
+    )
+
+
+@job
+def interpolate_structures(p_st: Structure, np_st: Structure, nimages: int) -> list:
+    """
+    Interpolate linearly the polar and the nonpolar structures with nimages structures.
+
+    Parameters
+    ----------
+    p_st : Structure
+        A pymatgen structure of polar phase.
+    np_st : Structure
+        A pymatgen structure of nonpolar phase.
+    nimages : int
+        Number of interpolatated structures calculated
+        from polar to nonpolar structures.
+
+    Returns
+    -------
+    List of interpolated structures
+    """
+    # adding +1 to nimages to match convention used in the interpolate
+    # func where nonpolar is (weirdly) included in the nimages count
+    return p_st.interpolate(
+        np_st, nimages + 1, interpolate_lattices=True, autosort_tol=0.0
+    )
+
+
+@job
+def add_interpolation_flow(
+    interp_structures: list[Structure], lcalcpol_maker: BaseVaspMaker
+) -> Response:
+    """
+    Generate the interpolations jobs and add them to the main ferroelectric flow.
+
+    Parameters
+    ----------
+    interp_structures: List[Structure]
+        List of interpolated structures
+    lcalcpol_maker : BaseVaspMaker
+        Vasp maker to compute the polarization of each structure.
+
+    Returns
+    -------
+    Response
+        Job response containing the interpolation flow.
+    """
+    jobs = []
+    outputs = {}
+
+    for i, interp_structure in enumerate(interp_structures[1:-1]):
+        lcalcpol_maker.write_additional_data["structures:json"] = {
+            "st_polar": interp_structures[0],
+            "st_nonpolar": interp_structures[-1],
+            "st_interp_idx": i + 1,
+        }
+        interpolation = lcalcpol_maker.make(interp_structure)
+        interpolation.append_name(f" interpolation_{i}")
+        jobs.append(interpolation)
+        output = get_polarization_output(interpolation)
+        outputs.update({f"interpolation_{i}": output})
+
+    interp_flow = Flow(jobs, outputs)
+    return Response(replace=interp_flow)
+
+
+def get_polarization_output(job: Job) -> dict:
+    """
+    Extract from lcalcpol job all the relevant output to compute the polarization.
+
+    Parameters
+    ----------
+    job : Job
+        Job from which to extract relevant quantities.
+
+    Returns
+    -------
+    dict
+        Dictionary containing the extracted polarization data.
+    """
+    p = job.output
+    outcar = p.calcs_reversed[0].output.outcar
+
+    return {
+        "energy_per_atom": p.calcs_reversed[0].output.energy_per_atom,
+        "energy": p.calcs_reversed[0].output.energy,
+        "task_label": p.task_label,
+        "structure": p.structure,
+        "zval_dict": outcar["zval_dict"],
+        "p_elecs": outcar["p_elec"],
+        "job_dir": p.dir_name,
+        "uuid": p.uuid,
+    }