Skip to content

Commit a0b4bda

Browse files
frariccifraricci
authored
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
1 parent 68f7e07 commit a0b4bda

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

45 files changed

+776
-0
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,3 +72,4 @@ docs/reference/atomate2.*
7272
*.doctrees*
7373

7474
.ipynb_checkpoints
75+
.aider*

‎src/atomate2/vasp/flows/ferroelectric.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
"""Flows for calculating the polarization of a polar material."""
2+
3+
from __future__ import annotations
4+
5+
import logging
6+
from dataclasses import dataclass, field
7+
from typing import TYPE_CHECKING
8+
9+
from jobflow import Flow, Maker
10+
11+
if TYPE_CHECKING:
12+
from pathlib import Path
13+
14+
from pymatgen.core.structure import Structure
15+
16+
from atomate2.vasp.jobs.base import BaseVaspMaker
17+
18+
from atomate2.vasp.flows.core import DoubleRelaxMaker
19+
from atomate2.vasp.jobs.core import PolarizationMaker, RelaxMaker
20+
from atomate2.vasp.jobs.ferroelectric import (
21+
add_interpolation_flow,
22+
get_polarization_output,
23+
interpolate_structures,
24+
polarization_analysis,
25+
)
26+
27+
__all__ = ["FerroelectricMaker"]
28+
29+
logger = logging.getLogger(__name__)
30+
31+
32+
@dataclass
33+
class FerroelectricMaker(Maker):
34+
"""
35+
Maker to calculate polarization of a polar material.
36+
37+
Parameters
38+
----------
39+
name : str
40+
Name of the flows produced by this maker.
41+
nimages: int
42+
Number of interpolated structures calculated from polar to nonpolar structures
43+
relax_maker: BaseVaspMaker or None or tuple
44+
None to avoid relaxation of both polar and nonpolar structures
45+
BaseVaspMaker to relax both structures (default)
46+
tuple of BaseVaspMaker and None to control relaxation for each structure
47+
lcalcpol_maker: BaseVaspMaker
48+
Vasp maker to compute the polarization of each structure
49+
"""
50+
51+
name: str = "ferroelectric"
52+
nimages: int = 8
53+
relax_maker: BaseVaspMaker | None | tuple = field(
54+
default_factory=lambda: DoubleRelaxMaker.from_relax_maker(RelaxMaker())
55+
)
56+
lcalcpol_maker: BaseVaspMaker = field(default_factory=PolarizationMaker)
57+
58+
def make(
59+
self,
60+
polar_structure: Structure,
61+
nonpolar_structure: Structure,
62+
prev_vasp_dir: str | Path | None = None,
63+
) -> Flow:
64+
"""
65+
Make flow to calculate the polarization.
66+
67+
Parameters
68+
----------
69+
polar_structure : .Structure
70+
A pymatgen structure of the polar phase.
71+
nonpolar_structure : .Structure
72+
A pymatgen structure of the nonpolar phase.
73+
prev_vasp_dir : str or Path or None
74+
A previous vasp calculation directory to use for copying outputs.
75+
"""
76+
jobs = []
77+
prev_vasp_dir_p, prev_vasp_dir_np = None, None
78+
79+
if not isinstance(self.relax_maker, tuple):
80+
self.relax_maker = (self.relax_maker, self.relax_maker)
81+
82+
if self.relax_maker[0]:
83+
# optionally relax the polar structure
84+
relax_p = self.relax_maker[0].make(polar_structure)
85+
relax_p.append_name(" polar")
86+
jobs.append(relax_p)
87+
polar_structure = relax_p.output.structure
88+
prev_vasp_dir_p = relax_p.output.dir_name
89+
90+
logger.info(f"{type(polar_structure)}")
91+
92+
polar_lcalcpol = self.lcalcpol_maker.make(
93+
polar_structure, prev_dir=prev_vasp_dir_p
94+
)
95+
polar_lcalcpol.append_name(" polar")
96+
jobs.append(polar_lcalcpol)
97+
polar_structure = polar_lcalcpol.output.structure
98+
99+
if self.relax_maker[1]:
100+
# optionally relax the nonpolar structure
101+
relax_np = self.relax_maker[1].make(nonpolar_structure)
102+
relax_np.append_name(" nonpolar")
103+
jobs.append(relax_np)
104+
nonpolar_structure = relax_np.output.structure
105+
prev_vasp_dir_np = relax_np.output.dir_name
106+
107+
nonpolar_lcalcpol = self.lcalcpol_maker.make(
108+
nonpolar_structure, prev_dir=prev_vasp_dir_np
109+
)
110+
nonpolar_lcalcpol.append_name(" nonpolar")
111+
jobs.append(nonpolar_lcalcpol)
112+
nonpolar_structure = nonpolar_lcalcpol.output.structure
113+
114+
interp_structs_job = interpolate_structures(
115+
polar_structure, nonpolar_structure, self.nimages
116+
)
117+
jobs.append(interp_structs_job)
118+
119+
interp_structs = interp_structs_job.output
120+
add_interp_flow = add_interpolation_flow(interp_structs, self.lcalcpol_maker)
121+
122+
pol_analysis = polarization_analysis(
123+
get_polarization_output(nonpolar_lcalcpol),
124+
get_polarization_output(polar_lcalcpol),
125+
add_interp_flow.output,
126+
)
127+
128+
jobs.append(add_interp_flow)
129+
jobs.append(pol_analysis)
130+
131+
return Flow(
132+
jobs=jobs,
133+
output=pol_analysis.output,
134+
name=self.name,
135+
)

‎src/atomate2/vasp/jobs/core.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,6 +533,45 @@ class DielectricMaker(BaseVaspMaker):
533533
)
534534

535535

536+
@dataclass
537+
class PolarizationMaker(BaseVaspMaker):
538+
"""
539+
Maker to create polarization calculation VASP jobs.
540+
541+
.. Note::
542+
If starting from a previous calculation, magnetism will be disabled if all
543+
MAGMOMs are less than 0.02.
544+
545+
Parameters
546+
----------
547+
name : str
548+
The job name.
549+
input_set_generator : .StaticSetGenerator
550+
A generator used to make the input set.
551+
write_input_set_kwargs : dict
552+
Keyword arguments that will get passed to :obj:`.write_vasp_input_set`.
553+
copy_vasp_kwargs : dict
554+
Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`.
555+
run_vasp_kwargs : dict
556+
Keyword arguments that will get passed to :obj:`.run_vasp`.
557+
task_document_kwargs : dict
558+
Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`.
559+
stop_children_kwargs : dict
560+
Keyword arguments that will get passed to :obj:`.should_stop_children`.
561+
write_additional_data : dict
562+
Additional data to write to the current directory. Given as a dict of
563+
{filename: data}. Note that if using FireWorks, dictionary keys cannot contain
564+
the "." character which is typically used to denote file extensions. To avoid
565+
this, use the ":" character, which will automatically be converted to ".". E.g.
566+
``{"my_file:txt": "contents of the file"}``.
567+
"""
568+
569+
name: str = "polarization"
570+
input_set_generator: StaticSetGenerator = field(
571+
default_factory=lambda: StaticSetGenerator(lcalcpol=True, auto_ispin=True)
572+
)
573+
574+
536575
@dataclass
537576
class TransmuterMaker(BaseVaspMaker):
538577
"""

‎src/atomate2/vasp/jobs/ferroelectric.py‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
"""Job used in the Ferroelectric wflow."""
2+
3+
from __future__ import annotations
4+
5+
import logging
6+
from typing import TYPE_CHECKING, Any
7+
8+
from jobflow import Flow, Job, Response, job
9+
from pymatgen.analysis.ferroelectricity.polarization import get_total_ionic_dipole
10+
11+
from atomate2.vasp.schemas.ferroelectric import PolarizationDocument
12+
13+
if TYPE_CHECKING:
14+
from pymatgen.core.structure import Structure
15+
16+
from atomate2.vasp.jobs.base import BaseVaspMaker
17+
18+
logger = logging.getLogger(__name__)
19+
20+
__all__ = ["polarization_analysis"]
21+
22+
23+
@job(output_schema=PolarizationDocument)
24+
def polarization_analysis(
25+
np_lcalcpol_output: dict[str, Any],
26+
p_lcalcpol_output: dict[str, Any],
27+
interp_lcalcpol_outputs: dict[str, Any],
28+
) -> PolarizationDocument:
29+
"""
30+
Recover the same branch polarization and the spontaneous polarization.
31+
32+
Parameters
33+
----------
34+
np_lcalcpol_output : dict
35+
Output from previous nonpolar lcalcpol job.
36+
p_lcalcpol_output : dict
37+
Output from previous polar lcalcpol job.
38+
interp_lcalcpol_outputs : dict
39+
Output from previous interpolation lcalcpol jobs.
40+
41+
Returns
42+
-------
43+
PolarizationDocument
44+
Document containing the polarization analysis results.
45+
46+
"""
47+
# order previous calculations from nonpolar to polar
48+
ordered_keys = [
49+
f"interpolation_{i}" for i in reversed(range(len(interp_lcalcpol_outputs)))
50+
]
51+
52+
polarization_tasks = [np_lcalcpol_output]
53+
polarization_tasks += [interp_lcalcpol_outputs[k] for k in ordered_keys]
54+
polarization_tasks += [p_lcalcpol_output]
55+
56+
task_lbls = []
57+
structures = []
58+
energies_per_atom = []
59+
energies = []
60+
job_dirs = []
61+
uuids = []
62+
63+
for i, p in enumerate(polarization_tasks):
64+
energies_per_atom.append(p["energy_per_atom"])
65+
energies.append(p["energy"])
66+
task_lbls.append(p["task_label"] or str(i))
67+
structures.append(p["structure"])
68+
job_dirs.append(p["job_dir"])
69+
uuids.append(p["uuid"])
70+
71+
# If LCALCPOL = True then Outcar will parse and store the pseudopotential zvals.
72+
zval_dict = p["zval_dict"]
73+
74+
# Assumes that we want to calculate the ionic contribution to the dipole moment.
75+
# VASP's ionic contribution is sometimes strange.
76+
# See pymatgen.analysis.ferroelectricity.polarization.Polarization for details.
77+
p_elecs = [p["p_elecs"] for p in polarization_tasks]
78+
p_ions = [get_total_ionic_dipole(st, zval_dict) for st in structures]
79+
80+
return PolarizationDocument.from_pol_output(
81+
p_elecs,
82+
p_ions,
83+
structures,
84+
energies,
85+
energies_per_atom,
86+
zval_dict,
87+
task_lbls,
88+
job_dirs,
89+
uuids,
90+
)
91+
92+
93+
@job
94+
def interpolate_structures(p_st: Structure, np_st: Structure, nimages: int) -> list:
95+
"""
96+
Interpolate linearly the polar and the nonpolar structures with nimages structures.
97+
98+
Parameters
99+
----------
100+
p_st : Structure
101+
A pymatgen structure of polar phase.
102+
np_st : Structure
103+
A pymatgen structure of nonpolar phase.
104+
nimages : int
105+
Number of interpolatated structures calculated
106+
from polar to nonpolar structures.
107+
108+
Returns
109+
-------
110+
List of interpolated structures
111+
"""
112+
# adding +1 to nimages to match convention used in the interpolate
113+
# func where nonpolar is (weirdly) included in the nimages count
114+
return p_st.interpolate(
115+
np_st, nimages + 1, interpolate_lattices=True, autosort_tol=0.0
116+
)
117+
118+
119+
@job
120+
def add_interpolation_flow(
121+
interp_structures: list[Structure], lcalcpol_maker: BaseVaspMaker
122+
) -> Response:
123+
"""
124+
Generate the interpolations jobs and add them to the main ferroelectric flow.
125+
126+
Parameters
127+
----------
128+
interp_structures: List[Structure]
129+
List of interpolated structures
130+
lcalcpol_maker : BaseVaspMaker
131+
Vasp maker to compute the polarization of each structure.
132+
133+
Returns
134+
-------
135+
Response
136+
Job response containing the interpolation flow.
137+
"""
138+
jobs = []
139+
outputs = {}
140+
141+
for i, interp_structure in enumerate(interp_structures[1:-1]):
142+
lcalcpol_maker.write_additional_data["structures:json"] = {
143+
"st_polar": interp_structures[0],
144+
"st_nonpolar": interp_structures[-1],
145+
"st_interp_idx": i + 1,
146+
}
147+
interpolation = lcalcpol_maker.make(interp_structure)
148+
interpolation.append_name(f" interpolation_{i}")
149+
jobs.append(interpolation)
150+
output = get_polarization_output(interpolation)
151+
outputs.update({f"interpolation_{i}": output})
152+
153+
interp_flow = Flow(jobs, outputs)
154+
return Response(replace=interp_flow)
155+
156+
157+
def get_polarization_output(job: Job) -> dict:
158+
"""
159+
Extract from lcalcpol job all the relevant output to compute the polarization.
160+
161+
Parameters
162+
----------
163+
job : Job
164+
Job from which to extract relevant quantities.
165+
166+
Returns
167+
-------
168+
dict
169+
Dictionary containing the extracted polarization data.
170+
"""
171+
p = job.output
172+
outcar = p.calcs_reversed[0].output.outcar
173+
174+
return {
175+
"energy_per_atom": p.calcs_reversed[0].output.energy_per_atom,
176+
"energy": p.calcs_reversed[0].output.energy,
177+
"task_label": p.task_label,
178+
"structure": p.structure,
179+
"zval_dict": outcar["zval_dict"],
180+
"p_elecs": outcar["p_elec"],
181+
"job_dir": p.dir_name,
182+
"uuid": p.uuid,
183+
}

0 commit comments

Comments
 (0)