# coding: utf-8
"""
Works and Flows for GW calculations with the quartic-scaling implementation.
"""
from __future__ import annotations
import os
#import numpy as np
from . import wrappers
from .nodes import Node
from .tasks import NscfTask
from .works import Work
[docs]
class ScreeningWork(Work):
"""
This work parallelizes the calculation of the q-points of the screening.
It also provides the callback `on_all_ok` that calls the mrgscr tool
to merge all the partial SCR files produced.
The final SCR file is available in the `outdata` directory of the work.
To use this SCR file as dependency of other tasks, use the syntax:
deps={scr_work: "SCR"}
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: ScreeningWork
"""
[docs]
@classmethod
def from_nscf_task(cls, nscf_task: NscfTask, scr_input, manager=None) -> ScreeningWork:
"""
Construct a `ScreeningWork` from a |NscfTask| object.
Args:
nscf_task: |NscfTask| object which produced the WFK file with empty states
scr_input: |AbinitInput| object representing a SCREENING calculation.
manager: |TaskManager| object.
"""
if not isinstance(nscf_task, NscfTask):
raise TypeError("task `%s` does not inherit from NscfTask" % nscf_task)
new = cls(manager=manager)
# Call abinit to get the list of q-points for the dielectric matrix
scr_ibz = scr_input.abiget_scr_ibz()
# Now we can register the task for the different q-points
for qpoint in scr_ibz.points:
new.register_scr_task(scr_input.new_with_vars(nqptdm=1, qptdm=qpoint),
deps={nscf_task: "WFK"})
return new
[docs]
@classmethod
def from_wfk_filepath(cls, wfk_filepath, scr_input, manager=None) -> ScreeningWork:
"""
Construct a `ScreeningWork` from a WFK filepath and a screening input.
WARNING: the parameters reported in `scr_input` e.g. k-mesh, nband, ecut etc
must be consisten with those used to generate the WFK file.
No consistencty check is done at this level. You have been warned!
Args:
wfk_filepath: Path to the WFK file.
scr_input: |AbinitInput| object representing a SCREENING calculation.
manager: |TaskManager| object.
"""
new = cls(manager=manager)
wfk_node = Node.as_node(wfk_filepath)
# Get list of q-points for the dielectric matrix
scr_ibz = scr_input.abiget_scr_ibz()
# Now we can register the task for the different q-points
for qpoint in scr_ibz.points:
new.register_scr_task(scr_input.new_with_vars(nqptdm=1, qptdm=qpoint),
deps={wfk_node: "WFK"})
return new
[docs]
def merge_scr_files(self, remove_scrfiles=True, verbose=0) -> str:
"""
This method is called when all the q-points have been computed.
It runs `mrgscr` in sequential on the local machine to produce
the final SCR file in the outdir of the `Work`.
If remove_scrfiles is True, the partial SCR files are removed after the merge.
"""
scr_files = list(filter(None, [task.outdir.has_abiext("SCR") for task in self]))
self.history.info("Will call mrgscr to merge %s SCR files:\n" % len(scr_files))
assert len(scr_files) == len(self)
mrgscr = wrappers.Mrgscr(manager=self[0].manager, verbose=verbose)
final_scr = mrgscr.merge_qpoints(self.outdir.path, scr_files, out_prefix="out")
if remove_scrfiles:
for scr_file in scr_files:
try:
os.remove(scr_file)
except IOError:
pass
return final_scr
[docs]
def on_all_ok(self): # pragma: no cover
"""
This method is called when all the q-points have been computed.
It runs `mrgscr` in sequential on the local machine to produce
the final SCR file in the outdir of the `Work`.
"""
final_scr = self.merge_scr_files()
return self.Results(node=self, returncode=0, message="mrgscr done", final_scr=final_scr)