"""
Master workchain.
-----------------
This is the master workchain, which most users should call when issuing day-day
calculations.
"""
# pylint: disable=attribute-defined-outside-init
from aiida.common.extendeddicts import AttributeDict
from aiida.engine import WorkChain, append_, if_
from aiida.plugins import WorkflowFactory
from aiida_vasp.utils.aiida_utils import get_data_class, get_data_node
from aiida_vasp.utils.workchains import compose_exit_code, prepare_process_inputs
[docs]
class MasterWorkChain(WorkChain):
"""The master workchain that selects sub workchains to perform necessary calculations."""
_verbose = False
_base_workchains_string = 'vasp.converge'
_bands_workchain_string = 'vasp.bands'
_dos_workchain_string = 'vasp.vasp'
_base_workchain = WorkflowFactory(_base_workchains_string)
_bands_workchain = WorkflowFactory(_bands_workchain_string)
_dos_workchain = WorkflowFactory(_dos_workchain_string)
[docs]
@classmethod
def define(cls, spec):
super(MasterWorkChain, cls).define(spec)
spec.expose_inputs(cls._base_workchain, exclude=['settings', 'clean_workdir'])
spec.input(
'settings',
valid_type=get_data_class('core.dict'),
required=False,
)
spec.input(
'kpoints',
valid_type=get_data_class('core.array.kpoints'),
required=False,
)
spec.input_namespace(
'relax',
required=False,
dynamic=True,
)
spec.input(
'extract_bands',
valid_type=get_data_class('core.bool'),
required=False,
default=lambda: get_data_node('core.bool', False),
help="""Do you want to extract the band structure?""",
)
spec.input(
'extract_dos',
valid_type=get_data_class('core.bool'),
required=False,
default=lambda: get_data_node('core.bool', False),
help="""Do you want to extract the density of states?""",
)
spec.input(
'dos.kpoints_distance',
valid_type=get_data_class('core.float'),
required=False,
default=lambda: get_data_node('core.float', 0.1),
help="""
The target k-point distance for density of states extraction.
""",
)
spec.input(
'dos.kpoints',
valid_type=get_data_class('core.array.kpoints'),
required=False,
help="""
The target k-point distance for density of states extraction.
""",
)
spec.input(
'kpoints_distance',
valid_type=get_data_class('core.float'),
required=False,
help="""The maximum distance between k-points in inverse AA.""",
)
spec.outline(
cls.initialize,
cls.init_prerun,
cls.init_workchain,
cls.run_next_workchain,
cls.verify_next_workchain,
if_(cls.extract_bands)(
cls.init_bands,
cls.init_workchain,
cls.run_next_workchain,
cls.verify_next_workchain
),
if_(cls.extract_dos)(
cls.init_dos,
cls.init_workchain,
cls.run_next_workchain,
cls.verify_next_workchain
),
cls.finalize
) # yapf: disable
spec.expose_outputs(
cls._bands_workchain,
namespace='bands',
namespace_options={
'required': False,
'populate_defaults': False
},
)
spec.expose_outputs(
cls._dos_workchain,
namespace='dos',
namespace_options={
'required': False,
'populate_defaults': False
},
)
spec.exit_code(
0,
'NO_ERROR',
message='the sun is shining',
)
spec.exit_code(
420,
'ERROR_NO_CALLED_WORKCHAIN',
message='no called workchain detected',
)
spec.exit_code(
500,
'ERROR_UNKNOWN',
message='unknown error detected in the master workchain',
)
[docs]
def initialize(self):
"""Initialize."""
self._init_context()
self._init_inputs()
self._init_settings()
self._set_base_workchain()
def _init_context(self):
"""Initialize context variables."""
self.ctx.exit_code = self.exit_codes.ERROR_UNKNOWN # pylint: disable=no-member
self.ctx.inputs = AttributeDict()
def _init_inputs(self):
"""Initialize inputs."""
try:
self._verbose = self.inputs.verbose.value
self.ctx.inputs.verbose = self.inputs.verbose
except AttributeError:
pass
# If we want to keep previous outputs for relaunch, do not clean remote folders
if self.extract_bands() or self.extract_dos():
self.ctx.inputs.clean_workdir = get_data_node('core.bool', False)
self._init_structure()
self._init_kpoints()
def _init_structure(self):
"""Initialize the structure."""
self.ctx.inputs.structure = self.inputs.structure
def _init_kpoints(self):
"""
Initialize the kpoints.
Since the reciprocal lattice is set in KpointsData, we need, if not performing
convergence tests to access the single StructureData, update the
supplied KpointsData with the right reciprocal lattice.
"""
kpoints = self._get_kpoints(self.inputs)
if kpoints is not None:
self.inputs.kpoints = kpoints
def _get_kpoints(self, inputs):
"""Decides if we are going to generate the k-point grid based on spacing or mesh."""
# First check if a k-point distance is supplied
try:
distance = inputs.kpoints_distance.value
except AttributeError:
# Then check if a KpointsData has been supplied
try:
kpoints = inputs.kpoints
except AttributeError:
# If neither, return, we need to run convergence tests
return None
return kpoints
kpoints = get_data_class('core.array.kpoints')()
kpoints.set_cell_from_structure(self.ctx.inputs.structure)
kpoints.set_kpoints_mesh_from_density(distance)
return kpoints
def _init_settings(self):
"""Initialize the settings."""
if 'settings' in self.inputs:
self.ctx.inputs.settings = self.inputs.settings
def _set_base_workchain(self):
"""Set the base workchain to be called."""
self._next_workchain = self._base_workchain
[docs]
def init_prerun(self):
"""Initialize the prerun."""
# Add relax inputs if they exists on the input
try:
self.ctx.inputs.relax = self.inputs.relax
except KeyError:
pass
[docs]
def init_bands(self):
"""Initialize the run to extract the band structure."""
self._next_workchain = self._bands_workchain
self._enable_charge_density_restart()
self._clean_inputs(exclude=['converge', 'relax', 'verify', 'kpoints'])
[docs]
def init_dos(self):
"""Initialize the run to extract the density of states at a denser k-point grid."""
self._next_workchain = self._dos_workchain
self._enable_charge_density_restart()
self._clean_inputs(exclude=['converge', 'relax', 'verify', 'dos'])
# Fetch density of states k-points
self.ctx.inputs.kpoints = self._get_kpoints(self.inputs.dos)
# Make sure we parse the density of states
if 'settings' in self.inputs:
settings = AttributeDict(self.inputs.settings.get_dict())
else:
settings = AttributeDict({'parser_settings': {}})
dict_entry = {'add_dos': True}
try:
settings.parser_settings.update(dict_entry)
except AttributeError:
settings.parser_settings = dict_entry
self.ctx.inputs.settings = settings
def _enable_charge_density_restart(self):
"""Enables a restart from a previously generated charge density object."""
# Make sure we set the restart folder (the charge density object is not
# copied locally, but is present in the folder of the previous remote directory)
self.ctx.inputs.restart_folder = self.ctx.workchains[-1].outputs.remote_folder
# Also enable the clean_workdir again
self.ctx.inputs.clean_workdir = get_data_node('core.bool', True)
def _clean_inputs(self, exclude=None):
"""Clean the inputs for the next workchain in order not to pass redundant inputs."""
# Now make sure we clean the inputs for redundant inputs not needed for the bands workchain
if exclude is None:
exclude = ['converge', 'relax', 'verify', 'kpoints']
self.ctx.inputs = AttributeDict({k: v for k, v in self.ctx.inputs.items() if k not in exclude})
[docs]
def init_workchain(self):
"""Initialize the base workchain."""
try:
self.ctx.inputs
except AttributeError as no_inputs:
raise ValueError('No input dictionary was defined in self.ctx.inputs') from no_inputs
# Add exposed inputs
self.ctx.inputs.update(self.exposed_inputs(self._next_workchain))
# Make sure we do not have any floating dict (convert to Dict)
self.ctx.inputs = prepare_process_inputs(
self.ctx.inputs, namespaces=['dynamics', 'relax', 'converge', 'verify']
)
[docs]
def run_next_workchain(self):
"""Run the next workchain."""
inputs = self.ctx.inputs
running = self.submit(self._next_workchain, **inputs)
self.report(f'launching {self._next_workchain.__name__}<{running.pk}> ')
return self.to_context(workchains=append_(running))
[docs]
def verify_next_workchain(self):
"""Inherit exit status from child workchains."""
try:
workchain = self.ctx.workchains[-1]
except IndexError:
self.report(f'There is no {self._next_workchain.__name__} in the called workchain list.')
return self.exit_codes.ERROR_NO_CALLED_WORKCHAIN # pylint: disable=no-member
# Inherit exit status from last workchain (supposed to be
# successfull)
next_workchain_exit_status = workchain.exit_status
next_workchain_exit_message = workchain.exit_message
if not next_workchain_exit_status:
self.ctx.exit_code = self.exit_codes.NO_ERROR # pylint: disable=no-member
else:
self.ctx.exit_code = compose_exit_code(next_workchain_exit_status, next_workchain_exit_message)
self.report(
'The called {}<{}> returned a non-zero exit status. '
'The exit status {} is inherited'.format(
workchain.__class__.__name__, workchain.pk, self.ctx.exit_code
)
)
return self.ctx.exit_code
[docs]
def extract_bands(self):
"""Determines if we should extract the band structure."""
return self.inputs.extract_bands.value
[docs]
def loop_structures(self):
"""Determine if we should continue to calculate structures."""
return len(self.ctx.structures.get_list()) > 0
[docs]
def finalize(self):
"""Finalize the workchain."""
workchain = self.ctx.workchains[-1]
if self.extract_bands():
self.out_many(self.exposed_outputs(workchain, self._bands_workchain, namespace='bands'))
if self.extract_dos():
self.out_many(self.exposed_outputs(workchain, self._dos_workchain, namespace='dos'))