# coding: utf-8
"""
Workflows for calculations within the VZISA approximation to the quasi-harmonic approximation.
See [Phys. Rev. B 110, 014103](https://doi.org/10.1103/PhysRevB.110.014103)
"""
from __future__ import annotations
import dataclasses
import numpy as np
from abipy.tools.serialization import Serializable
from abipy.core.structure import Structure
from abipy.tools.typing import PathLike, VectorLike
from abipy.abio.inputs import AbinitInput
from abipy.dfpt.vzsisa import Vzsisa
from abipy.flowtk.works import Work, PhononWork
from abipy.flowtk.flows import Flow
[docs]
class VzsisaFlow(Flow):
"""
Flow for QHA calculations with the V-ZSISA approach.
This is the main entry point for client code.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VzsisaFlow
"""
[docs]
def finalize(self):
"""
This method is called by the scheduler when the flow finishes.
It generates a `vzsisa.json` file in the output directory, which can be used
to build the Vzsisa object for post-processing.
"""
work = self[0]
data = {"bo_vol_scales": work.bo_vol_scales, "ph_vol_scales": work.ph_vol_scales}
data["initial_structure"] = work.initial_scf_input.structure
# Build list of strings with paths to the relevant output files ordered by volume.
data["gsr_relax_paths"] = [task.gsr_path for task in work.relax_tasks_vol]
gsr_relax_entries, gsr_relax_volumes = [], []
for task in work.relax_tasks_vol:
with task.open_gsr() as gsr:
gsr_relax_entries.append(dict(
volume=gsr.structure.volume,
energy_eV=float(gsr.energy),
pressure_GPa=float(gsr.pressure),
#structure=gsr.structure,
))
gsr_relax_volumes.append(gsr.structure.volume)
data["gsr_relax_entries"] = gsr_relax_entries
data["gsr_relax_volumes_ang3"] = gsr_relax_volumes
data["ddb_relax_paths"] = [ph_work.outdir.has_abiext("DDB") for ph_work in work.ph_works]
data["ddb_relax_volumes_ang3"] = [ph_work[0].input.structure.volume for ph_work in work.ph_works]
data["gsr_relax_edos_paths"] = [] if not work.edos_work else [task.gsr_path for task in work.edos_work]
data["gsr_relax_ebands_paths"] = []
if work.ndivsm != 0:
data["gsr_relax_ebands_paths"] = [ph_work.ebands_task.gsr_path
for ph_work in work.ph_works if ph_work.ebands_task is not None]
# Write json file.
VzsisaResults(**data).json_write(self.outdir.path_in("vzsisa.json"), indent=4)
return super().finalize()
[docs]
class VzsisaWork(Work):
"""
This work performs the structural relaxation of the initial structure,
then a set of distorted structures is generated and the relaxed structures
are used to compute phonons, BECS and the dielectric tensor with DFPT.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VzsisaWork
"""
[docs]
def on_ok(self, sender):
"""
This method is called when one task reaches status `S_OK`.
It executes on_all_ok when all tasks in self have reached `S_OK`.
"""
if sender == self.initial_relax_task:
# Get relaxed structure
relaxed_structure = sender.get_final_structure()
v0 = relaxed_structure.volume
# build new tasks for structural relaxation at fixed volume.
relax_template = self.relax_template
self.relax_tasks_vol = []
for bo_scale in self.bo_vol_scales:
new_structure = relax_template.structure.scale_lattice(v0 * bo_scale)
task = self.register_relax_task(relax_template.new_with_structure(new_structure))
self.relax_tasks_vol.append(task)
self.flow.allocate(build=True)
return super().on_ok(sender)
[docs]
def on_all_ok(self):
"""
This callback is called when all tasks in the Work reach status `S_OK`.
Here we add a new PhononWork for each volume using the relaxed structure.
"""
self.edos_work = None if self.edos_ngkpt is None else Work()
self.ph_works = []
# Build phonon works for the different relaxed structures associated to ph_vol_scales.
for task, bo_scale in zip(self.relax_tasks_vol, self.bo_vol_scales):
if all(abs(bo_scale - self.ph_vol_scales)) > 1e-3: continue
relaxed_structure = task.get_final_structure()
scf_input = self.initial_scf_input.new_with_structure(relaxed_structure)
ph_work = PhononWork.from_scf_input(scf_input, self.ngqpt, is_ngqpt=True, tolerance=None,
with_becs=self.with_becs, with_quad=self.with_quad,
ndivsm=0 if bo_scale != 1.0 else self.ndivsm)
ph_work.set_name(f"PH for {bo_scale=}")
# Reduce the number of files produced in the DFPT tasks to avoid possible disk quota issues.
for task in ph_work[1:]:
task.input.set_vars(prtpot=0)
self.flow.register_work(ph_work)
self.ph_works.append(ph_work)
# Add task for electron DOS calculation to edos_work.
if self.edos_work is not None:
edos_input = scf_input.make_edos_input(self.edos_ngkpt, prtwf=-1)
self.edos_work.register_nscf_task(edos_input, deps={ph_work[0]: "DEN"})
if self.edos_ngkpt is not None:
self.flow.register_work(self.edos_work)
self.flow.allocate(build=True)
return super().on_all_ok()
#@dataclasses.dataclass(kw_only=True)
#class GsrRelaxEntry:
# volume: float
# energy_eV: float
# pressure_GPa: float
# # structure: Optional[Any] = None # Uncomment if you want to store structure
[docs]
@dataclasses.dataclass(kw_only=True)
class VzsisaResults(Serializable):
"""
Main entry point for post-processing and visualizing the results of a Zsisa calculation.
The JSON file contains:
- BO and phonon volumetric scaling factors.
- Paths to GSR files of relaxed structures and their corresponding volumes.
- Paths to DDB files for phonon calculations.
- Paths to electronic DOS data, if available.
.. code-block:: python
data = VzsisaResults.json_load("json_filepath")
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: VzsisaResults
"""
bo_vol_scales: np.ndarray
ph_vol_scales: np.ndarray
initial_structure: Structure
gsr_relax_paths: list[str]
#gsr_relax_entries: list[GsrRelaxEntry]
gsr_relax_entries: list[dict]
gsr_relax_volumes_ang3: list[float]
ddb_relax_paths: list[str]
ddb_relax_volumes_ang3: list[str]
gsr_relax_edos_paths: list[str]
gsr_relax_ebands_paths: list[str]
[docs]
def get_vzsisa(self,
nqsmall_or_qppa: int,
anaget_kwargs: dict | None = None,
smearing_ev: float | None = None,
verbose: int = 0) -> Vzsisa:
"""
Build an instance of Vzsisa to plot the results.
Args:
nqsmall_or_qppa: Define the q-mesh for the computation of the PHDOS.
if > 0, it is interpreted as nqsmall
if < 0, it is interpreted as qppa.
anaget_kwargs: dict with extra arguments passed to anaget_phdoses_with_gauss.
smearing_ev: Gaussian smearing in eV.
verbose: Verbosity level.
"""
return Vzsisa.from_gsr_ddb_paths(nqsmall_or_qppa, self.gsr_relax_paths, self.ddb_relax_paths,
anaget_kwargs, smearing_ev, verbose)