|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +############################################################################################################## |
| 3 | +"""Workchain to generate a restart remote folder for a Quantum ESPRESSO calculation.""" |
| 4 | + |
| 5 | +from aiida.orm import RemoteData, FolderData |
| 6 | +from aiida.engine import WorkChain, ToContext |
| 7 | +from aiida.plugins import WorkflowFactory, CalculationFactory |
| 8 | + |
| 9 | +from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin |
| 10 | +from aiida_quantumespresso.utils.transfer import get_transfer_builder |
| 11 | + |
| 12 | +TransferCalcjob = CalculationFactory('core.transfer') |
| 13 | +PwBaseWorkChain = WorkflowFactory('quantumespresso.pw.base') |
| 14 | + |
| 15 | +# pylint: disable=f-string-without-interpolation |
| 16 | + |
| 17 | + |
| 18 | +############################################################################################################## |
| 19 | +class RestartSetupWorkChain(ProtocolMixin, WorkChain): |
| 20 | + """Workchain to generate a restart remote folder for a Quantum ESPRESSO calculation. |
| 21 | +
|
| 22 | + It consists of two steps: |
| 23 | +
|
| 24 | + 1. TransferCalcjob: takes the content of a ``FolderData`` node and copies it into the remote computer |
| 25 | + into a ``RemoteData`` folder with the correct folder structure. The original ``FolderData`` needs |
| 26 | + to have the necessary files in the right internal path (check the ``get_transfer_builder`` utility |
| 27 | + function and/or use it to retrieve the densities to have this already taken care of) |
| 28 | +
|
| 29 | + 2. PwBaseWorkChain: it runs an NSCF calculation using the previously created ``RemoteData`` as its |
| 30 | + ``parent_folder``. This will re-generate all the wavefunctions in the running directory, which |
| 31 | + are necessary for launching any other kind of QE calculation in restart mode. |
| 32 | +
|
| 33 | + """ |
| 34 | + |
| 35 | + @classmethod |
| 36 | + def define(cls, spec): |
| 37 | + """Define the process specification.""" |
| 38 | + super().define(spec) |
| 39 | + |
| 40 | + spec.expose_inputs( |
| 41 | + TransferCalcjob, |
| 42 | + namespace='transfer', |
| 43 | + namespace_options={ |
| 44 | + 'validator': validate_transfer, |
| 45 | + 'populate_defaults': False, |
| 46 | + 'help': 'Inputs for the `TransferCalcjob` to put the data on the cluster.', |
| 47 | + } |
| 48 | + ) |
| 49 | + |
| 50 | + spec.expose_inputs( |
| 51 | + PwBaseWorkChain, |
| 52 | + namespace='nscf', |
| 53 | + exclude=('clean_workdir', 'pw.parent_folder'), |
| 54 | + namespace_options={ |
| 55 | + 'validator': validate_nscf, |
| 56 | + 'populate_defaults': False, |
| 57 | + 'help': 'Inputs for the `PwBaseWorkChain` for the NSCF calculation.', |
| 58 | + } |
| 59 | + ) |
| 60 | + |
| 61 | + spec.inputs.validator = validate_inputs |
| 62 | + |
| 63 | + spec.outline( |
| 64 | + cls.run_transfer, |
| 65 | + cls.inspect_transfer, |
| 66 | + cls.run_nscf, |
| 67 | + cls.inspect_nscf, |
| 68 | + cls.results, |
| 69 | + ) |
| 70 | + |
| 71 | + spec.output( |
| 72 | + 'remote_data', |
| 73 | + valid_type=RemoteData, |
| 74 | + help='The output node with the folder to be used as parent folder for other calculations.', |
| 75 | + ) |
| 76 | + |
| 77 | + spec.exit_code(401, 'ERROR_SUB_PROCESS_FAILED_TRANSFER', message='The TransferCalcjob sub process failed.') |
| 78 | + spec.exit_code(402, 'ERROR_SUB_PROCESS_FAILED_NSCF', message='The ncf PwBasexWorkChain sub process failed.') |
| 79 | + |
| 80 | + @classmethod |
| 81 | + def get_protocol_filepath(cls): |
| 82 | + """Return ``pathlib.Path`` to the ``.yaml`` file that defines the protocols.""" |
| 83 | + raise NotImplementedError(f'`get_protocol_filepath` method not yet implemented in `RestartSetupWorkChain`') |
| 84 | + |
| 85 | + @classmethod |
| 86 | + def get_builder_from_protocol(cls, folder_data, structure, code, protocol=None, overrides=None, **kwargs): |
| 87 | + """Return a builder prepopulated with inputs selected according to the chosen protocol. |
| 88 | +
|
| 89 | + :param data_source: the ``FolderData`` node containing the density (and the rest of the restart data). |
| 90 | + :param structure: the ``StructureData`` instance required to run the NSF calculation. |
| 91 | + :param code: the ``Code`` instance configured for the ``quantumespresso.pw`` plugin, required to |
| 92 | + run the NSF calculation. |
| 93 | + :param protocol: protocol to use, if not specified, the default will be used. |
| 94 | + :param overrides: optional dictionary of inputs to override the defaults of the protocol. |
| 95 | + :param kwargs: additional keyword arguments that will be passed to the ``get_builder_from_protocol`` |
| 96 | + of all the sub processes that are called by this workchain. |
| 97 | + :return: a process builder instance with all inputs defined ready for launch. |
| 98 | + """ |
| 99 | + # inputs = cls.get_protocol_inputs(protocol, overrides) |
| 100 | + |
| 101 | + builder = cls.get_builder() |
| 102 | + |
| 103 | + track = kwargs.get('track', False) |
| 104 | + transfer = get_transfer_builder(folder_data, computer=code.computer, track=track) |
| 105 | + transfer['metadata']['options']['resources'] = {} |
| 106 | + builder.transfer = transfer |
| 107 | + |
| 108 | + nscf_args = (code, structure, protocol) |
| 109 | + nscf_kwargs = kwargs |
| 110 | + nscf_kwargs['overrides'] = {} |
| 111 | + if overrides is not None: |
| 112 | + nscf_kwargs['overrides'] = overrides.get('nscf', None) |
| 113 | + |
| 114 | + # This is for easily setting defaults at each level of: |
| 115 | + # [overrides.nscf].pw.parameters.CONTROL.calculation |
| 116 | + last_layer = nscf_kwargs['overrides'] |
| 117 | + last_layer = last_layer.setdefault('pw', {}) |
| 118 | + last_layer = last_layer.setdefault('parameters', {}) |
| 119 | + last_layer = last_layer.setdefault('CONTROL', {}) |
| 120 | + |
| 121 | + if last_layer.setdefault('calculation', 'nscf') != 'nscf': |
| 122 | + bad_value = last_layer['calculation'] |
| 123 | + raise ValueError( |
| 124 | + f'The internal PwBaseWorkChain is for running an NSCF calculation, ' |
| 125 | + f'this should not be overriden. ' |
| 126 | + f'(Found overrides.nscf.pw.parameters.CONTROL.calculation=`{bad_value}`)' |
| 127 | + ) |
| 128 | + |
| 129 | + nscf = PwBaseWorkChain.get_builder_from_protocol(*nscf_args, **nscf_kwargs) |
| 130 | + nscf['pw'].pop('parent_folder', None) |
| 131 | + nscf.pop('clean_workdir', None) |
| 132 | + builder.nscf = nscf |
| 133 | + |
| 134 | + return builder |
| 135 | + |
| 136 | + def run_transfer(self): |
| 137 | + """Run the TransferCalcjob to put the data in the remote computer.""" |
| 138 | + inputs = self.exposed_inputs(TransferCalcjob, namespace='transfer') |
| 139 | + running = self.submit(TransferCalcjob, **inputs) |
| 140 | + self.report(f'launching TransferCalcjob<{running.pk}> for put the data into the remote computer') |
| 141 | + return ToContext(transfer_calcjob=running) |
| 142 | + |
| 143 | + def inspect_transfer(self): |
| 144 | + """Verify that the TransferCalcjob to get data finished successfully.""" |
| 145 | + calcjob_node = self.ctx.transfer_calcjob |
| 146 | + |
| 147 | + if not calcjob_node.is_finished_ok: |
| 148 | + self.report(f'TransferCalcjob failed with exit status {calcjob_node.exit_status}') |
| 149 | + return self.exit_codes.ERROR_SUB_PROCESS_FAILED_TRANSFER |
| 150 | + |
| 151 | + self.ctx.remote_parent = calcjob_node.outputs.remote_folder |
| 152 | + |
| 153 | + def run_nscf(self): |
| 154 | + """Run the PwBaseWorkChain in nscf mode on the restart folder.""" |
| 155 | + inputs = self.exposed_inputs(PwBaseWorkChain, namespace='nscf') |
| 156 | + inputs['metadata']['call_link_label'] = 'nscf' |
| 157 | + inputs['pw']['parent_folder'] = self.ctx.remote_parent |
| 158 | + |
| 159 | + running = self.submit(PwBaseWorkChain, **inputs) |
| 160 | + self.report(f'launching PwBaseWorkChain<{running.pk}> in nscf mode') |
| 161 | + return ToContext(workchain_nscf=running) |
| 162 | + |
| 163 | + def inspect_nscf(self): |
| 164 | + """Verify that the PwBaseWorkChain for the scf run finished successfully.""" |
| 165 | + workchain = self.ctx.workchain_nscf |
| 166 | + |
| 167 | + if not workchain.is_finished_ok: |
| 168 | + self.report(f'scf PwBaseWorkChain failed with exit status {workchain.exit_status}') |
| 169 | + return self.exit_codes.ERROR_SUB_PROCESS_FAILED_NSCF |
| 170 | + |
| 171 | + self.ctx.remote_data = workchain.outputs.remote_folder |
| 172 | + |
| 173 | + def results(self): |
| 174 | + """Attach the desired output nodes directly as outputs of the workchain.""" |
| 175 | + self.report('workchain succesfully completed') |
| 176 | + self.out('remote_data', self.ctx.remote_data) |
| 177 | + |
| 178 | + |
| 179 | +############################################################################################################## |
| 180 | +def validate_transfer(value, _): |
| 181 | + """Validate the inputs of the transfer input namespace.""" |
| 182 | + |
| 183 | + # Check that the source node is there and is of right type |
| 184 | + if 'source_nodes' not in value: |
| 185 | + return f'The inputs of the transfer namespace were not set correctly: {value}' |
| 186 | + |
| 187 | + source_nodes = value['source_nodes'] |
| 188 | + if 'source_node' not in source_nodes: |
| 189 | + return f'The `source_nodes` in the transfer namespace was not set correctly: {source_nodes}' |
| 190 | + |
| 191 | + source_node = source_nodes['source_node'] |
| 192 | + if not isinstance(source_node, FolderData): |
| 193 | + return f'The `source_node` in the transfer namespace is not `FolderData`: {source_node}' |
| 194 | + |
| 195 | + # Check that the files are in the source node |
| 196 | + error_message = '' |
| 197 | + if 'data-file-schema.xml' not in source_node.list_object_names(): |
| 198 | + error_message += f'Missing `data-file-schema.xml` on node PK={source_node.pk}\n' |
| 199 | + |
| 200 | + if 'charge-density.dat' not in source_node.list_object_names(): |
| 201 | + error_message += f'Missing `charge-density.dat` on node PK={source_node.pk}\n' |
| 202 | + |
| 203 | + if len(error_message) > 0: |
| 204 | + return error_message |
| 205 | + |
| 206 | + |
| 207 | +def validate_nscf(value, _): |
| 208 | + """Validate the inputs of the nscf input namespace.""" |
| 209 | + parameters = value['pw']['parameters'].get_dict() |
| 210 | + if parameters.get('CONTROL', {}).get('calculation', 'scf') != 'nscf': |
| 211 | + return '`CONTOL.calculation` in `nscf.pw.parameters` is not set to `nscf`.' |
| 212 | + |
| 213 | + |
| 214 | +def validate_inputs(inputs, _): |
| 215 | + """Validate the inputs of the entire input namespace.""" |
| 216 | + computer_transfer = inputs['transfer']['metadata']['computer'] |
| 217 | + computer_nscf = inputs['nscf']['pw']['code'].computer |
| 218 | + |
| 219 | + if computer_transfer.pk != computer_nscf.pk: |
| 220 | + return ( |
| 221 | + f'The computer where the files are being copied ({computer_transfer}) ' |
| 222 | + f'is not where the code resides ({computer_nscf})' |
| 223 | + ) |
0 commit comments