Add transfer calculation for density restart

aiida_quantumespresso/utils/transfer.py

@@ -0,0 +1,91 @@
+# -*- coding: utf-8 -*-
+##############################################################################################################
+"""Return process builders ready for transferring Quantum ESPRESSO density restart files."""
+
+import warnings
+
+from aiida.orm import FolderData, RemoteData, Dict
+from aiida.engine import calcfunction
+from aiida.plugins import CalculationFactory
+
+
+def get_transfer_builder(data_source, computer=None, track=False):
+ """Create a `ProcessBuilder` for `TransferCalcjob`from a data_source.
+
+ The data_source can be of either `RemoteData` or `FolderData`:
+
+ - `RemoteData`: generate a set of instructions so that the density restart data will be taken from the
+ remote computer specified by the node and into the local aiida DB.
+
+ - `FolderData`: generate a set of instructions so that the density restart data will be taken from the
+ local aiida DB and into the provided computer (which has to be given as an extra parameter).
+
+ :param data_source: the node instance from which to take the density data
+ :param computer: if `data_source` is a `FolderData` node, the remote computer to which to transfer the data must
+ be specified here
+ :param track: boolean, if True, the generation of the instructions will be done through a calcfunction from the
+ data_source as input (and thus be tracked as such in the provenance)
+ :return: a `ProcessBuilder` instance configured for launching a `TransferCalcjob`
+ """
+ builder = CalculationFactory('core.transfer').get_builder()
+ builder.source_nodes = {'source_node': data_source}
+
+ if isinstance(data_source, FolderData):
+ if computer is None:
+ raise ValueError('No computer was provided for setting up a transfer to a remote.')
+ builder.metadata['computer'] = computer
+
+ elif isinstance(data_source, RemoteData):
+ if computer is not None:
+ warnings.warn(
+ f'Computer `{computer}` provided will be ignored '
+ f'(using `{data_source.computer}` from the RemoteData input `{data_source}`)'
+ )
+ builder.metadata['computer'] = data_source.computer
+
+ if track:
+ builder.instructions = generate_instructions(data_source)['instructions']
+ else:
+ builder.instructions = generate_instructions_untracked(data_source)
+
+ return builder
+
+
+##############################################################################################################
+def generate_instructions_untracked(source_folder):
+ """Generate the instruction node to be used for copying the files."""
+
+ # Paths in the QE run folder
+ schema_qepath = 'out/aiida.save/data-file-schema.xml'
+ charge_qepath = 'out/aiida.save/charge-density.dat'
+ pawtxt_qepath = 'out/aiida.save/paw.txt'
+
+ # Paths in the local node
+ schema_dbpath = 'data-file-schema.xml'
+ charge_dbpath = 'charge-density.dat'
+ pawtxt_dbpath = 'paw.txt'
+
+ # Transfer from local to remote
+ if isinstance(source_folder, FolderData):
+ instructions = {'retrieve_files': False, 'local_files': []}
+ instructions['local_files'].append(('source_node', schema_dbpath, schema_qepath))
+ instructions['local_files'].append(('source_node', charge_dbpath, charge_qepath))
+
+ if 'paw.txt' in source_folder.list_object_names():
+ instructions['local_files'].append(('source_node', pawtxt_dbpath, pawtxt_qepath))
+
+ # Transfer from remote to local
+ elif isinstance(source_folder, RemoteData):
+ instructions = {'retrieve_files': True, 'symlink_files': []}
+ instructions['symlink_files'].append(('source_node', schema_qepath, schema_dbpath))
+ instructions['symlink_files'].append(('source_node', charge_qepath, charge_dbpath))
+ instructions['symlink_files'].append(('source_node', pawtxt_qepath, pawtxt_dbpath))
+
+ return Dict(dict=instructions)
+
+
+@calcfunction
+def generate_instructions(source_folder):
+ """Auxiliary function to keep provenance track of the generation of the instructions."""
+ output_node = generate_instructions_untracked(source_folder)
+ return {'instructions': output_node}

aiida_quantumespresso/workflows/restart_setup.py

@@ -0,0 +1,227 @@
+# -*- coding: utf-8 -*-
+##############################################################################################################
+"""Workchain to generate a restart remote folder for a Quantum ESPRESSO calculation."""
+
+from aiida.orm import RemoteData, FolderData
+from aiida.engine import WorkChain, ToContext
+from aiida.plugins import WorkflowFactory, CalculationFactory
+
+from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin
+from aiida_quantumespresso.utils.transfer import get_transfer_builder
+
+TransferCalcjob = CalculationFactory('core.transfer')
+PwBaseWorkChain = WorkflowFactory('quantumespresso.pw.base')
+
+# pylint: disable=f-string-without-interpolation
+
+
+##############################################################################################################
+class RestartSetupWorkChain(ProtocolMixin, WorkChain):
+ """Workchain to generate a restart remote folder for a Quantum ESPRESSO calculation.
+
+ It consists of two steps:
+
+ 1. TransferCalcjob: takes the content of a ``FolderData`` node and copies it into the remote computer
+ into a ``RemoteData`` folder with the correct folder structure. The original ``FolderData`` needs
+ to have the necessary files in the right internal path (check the ``get_transfer_builder`` utility
+ function and/or use it to retrieve the densities to have this already taken care of)
+
+ 2. PwBaseWorkChain: it runs an NSCF calculation using the previously created ``RemoteData`` as its
+ ``parent_folder``. This will re-generate all the wavefunctions in the running directory, which
+ are necessary for launching any other kind of QE calculation in restart mode.
+
+ """
+
+ @classmethod
+ def define(cls, spec):
+ """Define the process specification."""
+ super().define(spec)
+
+ spec.expose_inputs(
+ TransferCalcjob,
+ namespace='transfer',
+ namespace_options={
+ 'validator': validate_transfer,
+ 'populate_defaults': False,
+ 'help': 'Inputs for the `TransferCalcjob` to put the data on the cluster.',
+ }
+ )
+
+ spec.expose_inputs(
+ PwBaseWorkChain,
+ namespace='nscf',
+ exclude=('clean_workdir', 'pw.parent_folder'),
+ namespace_options={
+ 'validator': validate_nscf,
+ 'populate_defaults': False,
+ 'help': 'Inputs for the `PwBaseWorkChain` for the NSCF calculation.',
+ }
+ )
+
+ spec.inputs.validator = validate_inputs
+
+ spec.outline(
+ cls.run_transfer,
+ cls.inspect_transfer,
+ cls.run_nscf,
+ cls.inspect_nscf,
+ cls.results,
+ )
+
+ spec.output(
+ 'remote_data',
+ valid_type=RemoteData,
+ help='The output node with the folder to be used as parent folder for other calculations.',
+ )
+
+ spec.exit_code(401, 'ERROR_SUB_PROCESS_FAILED_TRANSFER', message='The TransferCalcjob sub process failed.')
+ spec.exit_code(402, 'ERROR_SUB_PROCESS_FAILED_NSCF', message='The ncf PwBasexWorkChain sub process failed.')
+
+ @classmethod
+ def get_protocol_filepath(cls):
+ """Return ``pathlib.Path`` to the ``.yaml`` file that defines the protocols."""
+ raise NotImplementedError(f'`get_protocol_filepath` method not yet implemented in `RestartSetupWorkChain`')
+
+ @classmethod
+ def get_builder_from_protocol(cls, folder_data, structure, code, protocol=None, overrides=None, **kwargs):
+ """Return a builder prepopulated with inputs selected according to the chosen protocol.
+
+ :param data_source: the ``FolderData`` node containing the density (and the rest of the restart data).
+ :param structure: the ``StructureData`` instance required to run the NSF calculation.
+ :param code: the ``Code`` instance configured for the ``quantumespresso.pw`` plugin, required to
+ run the NSF calculation.
+ :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 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.
+ """
+ # inputs = cls.get_protocol_inputs(protocol, overrides)
+
+ builder = cls.get_builder()
+
+ track = kwargs.get('track', False)
+ transfer = get_transfer_builder(folder_data, computer=code.computer, track=track)
+ # The cls.builder seems to be checking that the metadata.options of each calcjob has
+ # a resources entry? There should be an exception to this for the CalcJobs that skip
+ # the submission, but until that is implemented in a supported version of aiida-core
+ # this patch needs to remain here:
+ transfer['metadata']['options']['resources'] = {}
+ builder.transfer = transfer
+
+ nscf_args = (code, structure, protocol)
+ nscf_kwargs = kwargs
+ nscf_kwargs['overrides'] = {}
+ if overrides is not None:
+ nscf_kwargs['overrides'] = overrides.get('nscf', None)
+
+ # This is for easily setting defaults at each level of:
+ # [overrides.nscf].pw.parameters.CONTROL.calculation
+ last_layer = nscf_kwargs['overrides']
+ last_layer = last_layer.setdefault('pw', {})
+ last_layer = last_layer.setdefault('parameters', {})
+ last_layer = last_layer.setdefault('CONTROL', {})
+
+ if last_layer.setdefault('calculation', 'nscf') != 'nscf':
+ bad_value = last_layer['calculation']
+ raise ValueError(
+ f'The internal PwBaseWorkChain is for running an NSCF calculation, '
+ f'this should not be overriden. '
+ f'(Found overrides.nscf.pw.parameters.CONTROL.calculation=`{bad_value}`)'
+ )
+
+ nscf = PwBaseWorkChain.get_builder_from_protocol(*nscf_args, **nscf_kwargs)
+ nscf['pw'].pop('parent_folder', None)
+ nscf.pop('clean_workdir', None)
+ builder.nscf = nscf
+
+ return builder
+
+ def run_transfer(self):
+ """Run the TransferCalcjob to put the data in the remote computer."""
+ inputs = self.exposed_inputs(TransferCalcjob, namespace='transfer')
+ running = self.submit(TransferCalcjob, **inputs)
+ self.report(f'launching TransferCalcjob<{running.pk}> for put the data into the remote computer')
+ return ToContext(transfer_calcjob=running)
+
+ def inspect_transfer(self):
+ """Verify that the TransferCalcjob to get data finished successfully."""
+ calcjob_node = self.ctx.transfer_calcjob
+
+ if not calcjob_node.is_finished_ok:
+ self.report(f'TransferCalcjob failed with exit status {calcjob_node.exit_status}')
+ return self.exit_codes.ERROR_SUB_PROCESS_FAILED_TRANSFER
+
+ self.ctx.remote_parent = calcjob_node.outputs.remote_folder
+
+ def run_nscf(self):
+ """Run the PwBaseWorkChain in nscf mode on the restart folder."""
+ inputs = self.exposed_inputs(PwBaseWorkChain, namespace='nscf')
+ inputs['metadata']['call_link_label'] = 'nscf'
+ inputs['pw']['parent_folder'] = self.ctx.remote_parent
+
+ running = self.submit(PwBaseWorkChain, **inputs)
+ self.report(f'launching PwBaseWorkChain<{running.pk}> in nscf mode')
+ return ToContext(workchain_nscf=running)
+
+ def inspect_nscf(self):
+ """Verify that the PwBaseWorkChain for the scf run finished successfully."""
+ workchain = self.ctx.workchain_nscf
+
+ if not workchain.is_finished_ok:
+ self.report(f'scf PwBaseWorkChain failed with exit status {workchain.exit_status}')
+ return self.exit_codes.ERROR_SUB_PROCESS_FAILED_NSCF
+
+ self.ctx.remote_data = workchain.outputs.remote_folder
+
+ def results(self):
+ """Attach the desired output nodes directly as outputs of the workchain."""
+ self.report('workchain succesfully completed')
+ self.out('remote_data', self.ctx.remote_data)
+
+
+##############################################################################################################
+def validate_transfer(value, _):
+ """Validate the inputs of the transfer input namespace."""
+
+ # Check that the source node is there and is of right type
+ if 'source_nodes' not in value:
+ return f'The inputs of the transfer namespace were not set correctly: {value}'
+
+ source_nodes = value['source_nodes']
+ if 'source_node' not in source_nodes:
+ return f'The `source_nodes` in the transfer namespace was not set correctly: {source_nodes}'
+
+ source_node = source_nodes['source_node']
+ if not isinstance(source_node, FolderData):
+ return f'The `source_node` in the transfer namespace is not `FolderData`: {source_node}'
+
+ # Check that the files are in the source node
+ error_message = ''
+ if 'data-file-schema.xml' not in source_node.list_object_names():
+ error_message += f'Missing `data-file-schema.xml` on node PK={source_node.pk}\n'
+
+ if 'charge-density.dat' not in source_node.list_object_names():
+ error_message += f'Missing `charge-density.dat` on node PK={source_node.pk}\n'
+
+ if len(error_message) > 0:
+ return error_message
+
+
+def validate_nscf(value, _):
+ """Validate the inputs of the nscf input namespace."""
+ parameters = value['pw']['parameters'].get_dict()
+ if parameters.get('CONTROL', {}).get('calculation', 'scf') != 'nscf':
+ return '`CONTOL.calculation` in `nscf.pw.parameters` is not set to `nscf`.'
+
+
+def validate_inputs(inputs, _):
+ """Validate the inputs of the entire input namespace."""
+ computer_transfer = inputs['transfer']['metadata']['computer']
+ computer_nscf = inputs['nscf']['pw']['code'].computer
+
+ if computer_transfer.pk != computer_nscf.pk:
+ return (
+ f'The computer where the files are being copied ({computer_transfer}) '
+ f'is not where the code resides ({computer_nscf})'
+ )

setup.json

@@ -65,6 +65,7 @@
  "quantumespresso.pw.band_structure = aiida_quantumespresso.workflows.pw.band_structure:PwBandStructureWorkChain",
  "quantumespresso.q2r.base = aiida_quantumespresso.workflows.q2r.base:Q2rBaseWorkChain",
  "quantumespresso.matdyn.base = aiida_quantumespresso.workflows.matdyn.base:MatdynBaseWorkChain",
+ "quantumespresso.restart_setup = aiida_quantumespresso.workflows.restart_setup:RestartSetupWorkChain",
  "quantumespresso.pdos = aiida_quantumespresso.workflows.pdos:PdosWorkChain"
  ],
  "console_scripts": [