Source code for aiida_quantumespresso.workflows.pw.relax

# -*- coding: utf-8 -*-
"""Workchain to relax a structure using Quantum ESPRESSO pw.x."""
from aiida import orm
from aiida.common import AttributeDict, exceptions
from aiida.common.lang import type_check
from aiida.engine import ToContext, WorkChain, append_, if_, while_
from aiida.plugins import CalculationFactory, WorkflowFactory

from aiida_quantumespresso.common.types import RelaxType
from aiida_quantumespresso.utils.mapping import prepare_process_inputs

from ..protocols.utils import ProtocolMixin

[docs]PwCalculation = CalculationFactory('quantumespresso.pw')
[docs]PwBaseWorkChain = WorkflowFactory('quantumespresso.pw.base')
[docs]def validate_inputs(inputs, _): """Validate the top level namespace.""" parameters = inputs['base']['pw']['parameters'].get_dict() if 'calculation' not in parameters.get('CONTROL', {}): return 'The parameters in `base.pw.parameters` do not specify the required key `CONTROL.calculation`.'
[docs]class PwRelaxWorkChain(ProtocolMixin, WorkChain): """Workchain to relax a structure using Quantum ESPRESSO pw.x.""" @classmethod
[docs] def define(cls, spec): """Define the process specification.""" # yapf: disable super().define(spec) spec.expose_inputs(PwBaseWorkChain, namespace='base', exclude=('clean_workdir', 'pw.structure', 'pw.parent_folder'), namespace_options={'help': 'Inputs for the `PwBaseWorkChain` for the main relax loop.'}) spec.expose_inputs(PwBaseWorkChain, namespace='base_final_scf', exclude=('clean_workdir', 'pw.structure', 'pw.parent_folder'), namespace_options={'required': False, 'populate_defaults': False, 'help': 'Inputs for the `PwBaseWorkChain` for the final scf.'}) spec.input('structure', valid_type=orm.StructureData, help='The inputs structure.') spec.input('meta_convergence', valid_type=orm.Bool, default=lambda: orm.Bool(True), help='If `True` the workchain will perform a meta-convergence on the cell volume.') spec.input('max_meta_convergence_iterations', valid_type=orm.Int, default=lambda: orm.Int(5), help='The maximum number of variable cell relax iterations in the meta convergence cycle.') spec.input('volume_convergence', valid_type=orm.Float, default=lambda: orm.Float(0.01), help='The volume difference threshold between two consecutive meta convergence iterations.') spec.input('clean_workdir', valid_type=orm.Bool, default=lambda: orm.Bool(False), help='If `True`, work directories of all called calculation will be cleaned at the end of execution.') spec.inputs.validator = validate_inputs spec.outline( cls.setup, while_(cls.should_run_relax)( cls.run_relax, cls.inspect_relax, ), if_(cls.should_run_final_scf)( cls.run_final_scf, cls.inspect_final_scf, ), cls.results, ) spec.exit_code(401, 'ERROR_SUB_PROCESS_FAILED_RELAX', message='the relax PwBaseWorkChain sub process failed') spec.exit_code(402, 'ERROR_SUB_PROCESS_FAILED_FINAL_SCF', message='the final scf PwBaseWorkChain sub process failed') spec.expose_outputs(PwBaseWorkChain, exclude=('output_structure',)) spec.output('output_structure', valid_type=orm.StructureData, required=False, help='The successfully relaxed structure.')
# yapf: enable @classmethod
[docs] def get_protocol_filepath(cls): """Return ``pathlib.Path`` to the ``.yaml`` file that defines the protocols.""" from importlib_resources import files from ..protocols import pw as pw_protocols return files(pw_protocols) / 'relax.yaml'
@classmethod
[docs] def get_builder_from_protocol( cls, code, structure, protocol=None, overrides=None, relax_type=RelaxType.POSITIONS_CELL, options=None, **kwargs ): """Return a builder prepopulated with inputs selected according to the chosen protocol. :param code: the ``Code`` instance configured for the ``quantumespresso.pw`` plugin. :param structure: the ``StructureData`` instance to use. :param protocol: protocol to use, if not specified, the default will be used. :param overrides: optional dictionary of inputs to override the defaults of the protocol. :param relax_type: the relax type to use: should be a value of the enum ``common.types.RelaxType``. :param options: A dictionary of options that will be recursively set for the ``metadata.options`` input of all the ``CalcJobs`` that are nested in this work chain. :param kwargs: additional keyword arguments that will be passed to the ``get_builder_from_protocol`` of all the sub processes that are called by this workchain. :return: a process builder instance with all inputs defined ready for launch. """ type_check(relax_type, RelaxType) inputs = cls.get_protocol_inputs(protocol, overrides) args = (code, structure, protocol) base = PwBaseWorkChain.get_builder_from_protocol( *args, overrides=inputs.get('base', None), options=options, **kwargs ) base_final_scf = PwBaseWorkChain.get_builder_from_protocol( *args, overrides=inputs.get('base_final_scf', None), options=options, **kwargs ) base['pw'].pop('structure', None) base.pop('clean_workdir', None) base_final_scf['pw'].pop('structure', None) base_final_scf.pop('clean_workdir', None) # Quantum ESPRESSO currently only supports optimization of the volume for simple cubic systems. It requires # to set `ibrav=1` or the code will except. if relax_type in (RelaxType.VOLUME, RelaxType.POSITIONS_VOLUME): raise ValueError(f'relax type `{relax_type} is not yet supported.') if relax_type in (RelaxType.VOLUME, RelaxType.SHAPE, RelaxType.CELL): base.pw.settings = orm.Dict(PwRelaxWorkChain._fix_atomic_positions(structure, base.pw.settings)) if relax_type is RelaxType.NONE: base.pw.parameters['CONTROL']['calculation'] = 'scf' base.pw.parameters.base.attributes.delete('CELL') elif relax_type is RelaxType.POSITIONS: base.pw.parameters['CONTROL']['calculation'] = 'relax' base.pw.parameters.base.attributes.delete('CELL') else: base.pw.parameters['CONTROL']['calculation'] = 'vc-relax' if relax_type in (RelaxType.VOLUME, RelaxType.POSITIONS_VOLUME): base.pw.parameters['CELL']['cell_dofree'] = 'volume' if relax_type in (RelaxType.SHAPE, RelaxType.POSITIONS_SHAPE): base.pw.parameters['CELL']['cell_dofree'] = 'shape' if relax_type in (RelaxType.CELL, RelaxType.POSITIONS_CELL): pbc_cell_dofree_map = { (True, True, True): 'all', (True, False, False): 'x', (False, True, False): 'y', (False, False, True): 'z', (True, True, False): '2Dxy', } if structure.pbc in pbc_cell_dofree_map: base.pw.parameters['CELL']['cell_dofree'] = pbc_cell_dofree_map[structure.pbc] else: raise ValueError(f'Structures with periodic boundary conditions `{structure.pbc}` are not supported.') builder = cls.get_builder() builder.base = base builder.base_final_scf = base_final_scf builder.structure = structure builder.clean_workdir = orm.Bool(inputs['clean_workdir']) builder.max_meta_convergence_iterations = orm.Int(inputs['max_meta_convergence_iterations']) builder.meta_convergence = orm.Bool(inputs['meta_convergence']) builder.volume_convergence = orm.Float(inputs['volume_convergence']) return builder
[docs] def setup(self): """Input validation and context setup.""" self.ctx.current_number_of_bands = None self.ctx.current_structure = self.inputs.structure self.ctx.current_cell_volume = None self.ctx.is_converged = False self.ctx.iteration = 0 self.ctx.relax_inputs = AttributeDict(self.exposed_inputs(PwBaseWorkChain, namespace='base')) self.ctx.relax_inputs.pw.parameters = self.ctx.relax_inputs.pw.parameters.get_dict() self.ctx.relax_inputs.pw.parameters.setdefault('CONTROL', {}) # Set the meta_convergence and add it to the context self.ctx.meta_convergence = self.inputs.meta_convergence.value volume_cannot_change = ( self.ctx.relax_inputs.pw.parameters['CONTROL'].get('calculation', 'scf') in ('scf', 'relax') or self.ctx.relax_inputs.pw.parameters.get('CELL', {}).get('cell_dofree', None) == 'shape' ) if self.ctx.meta_convergence and volume_cannot_change: self.report( 'No change in volume possible for the provided base input parameters. Meta convergence is turned off.' ) self.ctx.meta_convergence = False # Add the final scf inputs to the context if a final scf should be run if 'base_final_scf' in self.inputs: self.ctx.final_scf_inputs = AttributeDict(self.exposed_inputs(PwBaseWorkChain, namespace='base_final_scf')) if self.ctx.relax_inputs.pw.parameters['CONTROL'].get('calculation', 'scf') == 'scf': self.report( 'Work chain will not run final SCF when `calculation` is set to `scf` for the relaxation ' '`PwBaseWorkChain`.' ) self.ctx.pop('final_scf_inputs') else: self.ctx.final_scf_inputs.pw.parameters = self.ctx.final_scf_inputs.pw.parameters.get_dict() self.ctx.final_scf_inputs.pw.parameters.setdefault('CONTROL', {}) self.ctx.final_scf_inputs.metadata.call_link_label = 'final_scf'
[docs] def should_run_relax(self): """Return whether a relaxation workchain should be run. This is the case as long as the volume change between two consecutive relaxation runs is larger than the volume convergence threshold value and the maximum number of meta convergence iterations is not exceeded. """ return not self.ctx.is_converged and self.ctx.iteration < self.inputs.max_meta_convergence_iterations.value
[docs] def should_run_final_scf(self): """Return whether after successful relaxation a final scf calculation should be run. If the maximum number of meta convergence iterations has been exceeded and convergence has not been reached, the structure cannot be considered to be relaxed and the final scf should not be run. """ return self.ctx.is_converged and 'final_scf_inputs' in self.ctx
[docs] def run_relax(self): """Run the `PwBaseWorkChain` to run a relax `PwCalculation`.""" self.ctx.iteration += 1 inputs = self.ctx.relax_inputs inputs.pw.structure = self.ctx.current_structure # If one of the nested `PwBaseWorkChains` changed the number of bands, apply it here if self.ctx.current_number_of_bands is not None: inputs.pw.parameters.setdefault('SYSTEM', {})['nbnd'] = self.ctx.current_number_of_bands # Set the `CALL` link label inputs.metadata.call_link_label = f'iteration_{self.ctx.iteration:02d}' inputs = prepare_process_inputs(PwBaseWorkChain, inputs) running = self.submit(PwBaseWorkChain, **inputs) self.report(f'launching PwBaseWorkChain<{running.pk}>') return ToContext(workchains=append_(running))
[docs] def inspect_relax(self): """Inspect the results of the last `PwBaseWorkChain`. Compare the cell volume of the relaxed structure of the last completed workchain with the previous. If the difference ratio is less than the volume convergence threshold we consider the cell relaxation converged. """ workchain = self.ctx.workchains[-1] acceptable_statuses = ['ERROR_IONIC_CONVERGENCE_REACHED_EXCEPT_IN_FINAL_SCF'] if workchain.is_excepted or workchain.is_killed: self.report('relax PwBaseWorkChain was excepted or killed') return self.exit_codes.ERROR_SUB_PROCESS_FAILED_RELAX if workchain.is_failed and workchain.exit_status not in PwBaseWorkChain.get_exit_statuses(acceptable_statuses): self.report(f'relax PwBaseWorkChain failed with exit status {workchain.exit_status}') return self.exit_codes.ERROR_SUB_PROCESS_FAILED_RELAX try: structure = workchain.outputs.output_structure except exceptions.NotExistent: # If the calculation is set to 'scf', this is expected, so we are done if self.ctx.relax_inputs.pw.parameters['CONTROL']['calculation'] == 'scf': self.ctx.is_converged = True return self.report('`vc-relax` or `relax` PwBaseWorkChain finished successfully but without output structure') return self.exit_codes.ERROR_SUB_PROCESS_FAILED_RELAX prev_cell_volume = self.ctx.current_cell_volume curr_cell_volume = structure.get_cell_volume() # Set relaxed structure as input structure for next iteration self.ctx.current_structure = structure self.ctx.current_number_of_bands = workchain.outputs.output_parameters.get_dict()['number_of_bands'] self.report(f'after iteration {self.ctx.iteration} cell volume of relaxed structure is {curr_cell_volume}') # After first iteration, simply set the cell volume and restart the next base workchain if not prev_cell_volume: self.ctx.current_cell_volume = curr_cell_volume # If meta convergence is switched off we are done if not self.ctx.meta_convergence: self.ctx.is_converged = True return # Check whether the cell volume is converged volume_threshold = self.inputs.volume_convergence.value volume_difference = abs(prev_cell_volume - curr_cell_volume) / prev_cell_volume if volume_difference < volume_threshold: self.ctx.is_converged = True self.report( f'relative cell volume difference {volume_difference} smaller than threshold {volume_threshold}' ) else: self.report( f'current relative cell volume difference {volume_difference} larger than threshold {volume_threshold}' ) self.ctx.current_cell_volume = curr_cell_volume return
[docs] def run_final_scf(self): """Run the `PwBaseWorkChain` to run a final scf `PwCalculation` for the relaxed structure.""" inputs = self.ctx.final_scf_inputs inputs.pw.structure = self.ctx.current_structure inputs_nbnd = inputs.pw.parameters.get('SYSTEM', {}).get('nbnd', None) if self.ctx.current_number_of_bands is not None and inputs_nbnd is None: inputs.pw.parameters.setdefault('SYSTEM', {})['nbnd'] = self.ctx.current_number_of_bands inputs = prepare_process_inputs(PwBaseWorkChain, inputs) running = self.submit(PwBaseWorkChain, **inputs) self.report(f'launching PwBaseWorkChain<{running.pk}> for final scf') return ToContext(workchain_scf=running)
[docs] def inspect_final_scf(self): """Inspect the result of the final scf `PwBaseWorkChain`.""" workchain = self.ctx.workchain_scf if not workchain.is_finished_ok: self.report(f'final scf PwBaseWorkChain failed with exit status {workchain.exit_status}') return self.exit_codes.ERROR_SUB_PROCESS_FAILED_FINAL_SCF
[docs] def results(self): """Attach the output parameters and structure of the last workchain to the outputs.""" if self.ctx.is_converged and self.ctx.iteration <= self.inputs.max_meta_convergence_iterations.value: self.report(f'workchain completed after {self.ctx.iteration} iterations') else: self.report('maximum number of meta convergence iterations exceeded') # Get the latest relax workchain and pass the outputs final_relax_workchain = self.ctx.workchains[-1] if self.ctx.relax_inputs.pw.parameters['CONTROL']['calculation'] != 'scf': self.out('output_structure', final_relax_workchain.outputs.output_structure) try: self.out_many(self.exposed_outputs(self.ctx.workchain_scf, PwBaseWorkChain)) except AttributeError: self.out_many(self.exposed_outputs(final_relax_workchain, PwBaseWorkChain))
[docs] def on_terminated(self): """Clean the working directories of all child calculations if `clean_workdir=True` in the inputs.""" super().on_terminated() if self.inputs.clean_workdir.value is False: self.report('remote folders will not be cleaned') return cleaned_calcs = [] for called_descendant in self.node.called_descendants: if isinstance(called_descendant, orm.CalcJobNode): try: called_descendant.outputs.remote_folder._clean() # pylint: disable=protected-access cleaned_calcs.append(called_descendant.pk) except (IOError, OSError, KeyError): pass if cleaned_calcs: self.report(f"cleaned remote folders of calculations: {' '.join(map(str, cleaned_calcs))}")
@staticmethod
[docs] def _fix_atomic_positions(structure, settings): """Fix the atomic positions, by setting the `FIXED_COORDS` key in the `settings` input node.""" if settings is not None: settings = settings.get_dict() else: settings = {} settings['FIXED_COORDS'] = [[True, True, True]] * len(structure.sites) return settings