Skip to content

Commit 8fd52c2

Browse files
ramirezfranciscoframirezfranciscof
committed
Add transfer calculation for density restart
1 parent da94161 commit 8fd52c2

2 files changed

Lines changed: 346 additions & 0 deletions

File tree

‎aiida_quantumespresso/workflows/density/transfer.py‎

Lines changed: 345 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,345 @@
1+
# -*- coding: utf-8 -*-
2+
##############################################################################################################
3+
"""Workchain to compute a band structure for a given structure using Quantum ESPRESSO pw.x."""
4+
import warnings
5+
6+
from aiida import orm
7+
from aiida.engine import WorkChain, ToContext, if_, calcfunction
8+
from aiida.plugins import WorkflowFactory, CalculationFactory
9+
10+
from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin
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 TransferDensityWorkChain(ProtocolMixin, WorkChain):
20+
"""Workchain to transfer the charge density (and other restart data) from a Quantum ESPRESSO calculation.
21+
22+
If a `RemoteData` node is provided, this has to point to the folder where a PW calculation was last ran.
23+
This workchain will then take the density (and other restart data files) and copy it to a local node.
24+
25+
If a `FolderData` node is provided, you will also need to specify either a `remote_computer` or all the
26+
inputs for a `PwBaseWorkChain` run. With the first, this workchain will just transfer the density (and
27+
other restart data files) to the remote machine and end. If the inputs for a `PwBaseWorkChain` are
28+
provided instead, the workchain will also perform an NSCF calculation to re-generate the wavefunctions.
29+
30+
"""
31+
32+
@classmethod
33+
def define(cls, spec):
34+
"""Define the process specification."""
35+
super().define(spec)
36+
37+
spec.input(
38+
'data_source',
39+
valid_type=(orm.RemoteData, orm.FolderData),
40+
help='Node containing the density file and restart data (or path to it in a remote).'
41+
)
42+
43+
spec.input(
44+
'remote_computer',
45+
valid_type=orm.Computer,
46+
non_db=True,
47+
required=False,
48+
help='Computer to which to transfer the information.'
49+
)
50+
51+
spec.expose_inputs(
52+
PwBaseWorkChain,
53+
namespace='nscf',
54+
exclude=('clean_workdir', 'pw.parent_folder'),
55+
namespace_options={
56+
'required': False,
57+
'populate_defaults': False,
58+
'help': 'Inputs for the `PwBaseWorkChain` for the NSCF calculation.',
59+
}
60+
)
61+
62+
spec.inputs.validator = validate_inputs
63+
64+
spec.outline(
65+
cls.run_transfer,
66+
cls.inspect_transfer,
67+
if_(cls.should_run_nscf)(
68+
cls.run_nscf,
69+
cls.inspect_nscf,
70+
),
71+
cls.results,
72+
)
73+
74+
spec.output(
75+
'output_data',
76+
valid_type=(orm.RemoteData, orm.FolderData),
77+
help='The output node with the final data.',
78+
)
79+
80+
spec.exit_code(401, 'ERROR_SUB_PROCESS_FAILED_TRANSFER', message='The TransferCalcjob sub process failed.')
81+
spec.exit_code(402, 'ERROR_SUB_PROCESS_FAILED_NSCF', message='The ncf PwBasexWorkChain sub process failed.')
82+
83+
@classmethod
84+
def get_protocol_filepath(cls):
85+
"""Return ``pathlib.Path`` to the ``.yaml`` file that defines the protocols."""
86+
raise NotImplementedError(f'`get_protocol_filepath` method not yet implemented in `TransferDensityWorkChain`')
87+
88+
@classmethod
89+
def get_builder_from_protocol(
90+
cls, data_source, remote_computer=None, code=None, structure=None, protocol=None, overrides=None, **kwargs
91+
):
92+
"""Return a builder prepopulated with inputs selected according to the chosen protocol.
93+
94+
:param data_source: the source of the density (and the rest of the restart data).
95+
:param remote_computer: the ``Computer`` to which the data will be transfered. This is necessary
96+
if the ``data_source`` is a local ``FolderData`` and no code is going to be provided for
97+
running an NSCF calculation.
98+
:param code: the ``Code`` instance configured for the ``quantumespresso.pw`` plugin, required to
99+
run the NSF calculation.
100+
:param structure: the ``StructureData`` instance required to run the NSF calculation.
101+
:param protocol: protocol to use, if not specified, the default will be used.
102+
:param overrides: optional dictionary of inputs to override the defaults of the protocol.
103+
:param kwargs: additional keyword arguments that will be passed to the ``get_builder_from_protocol``
104+
of all the sub processes that are called by this workchain.
105+
:return: a process builder instance with all inputs defined ready for launch.
106+
"""
107+
# inputs = cls.get_protocol_inputs(protocol, overrides)
108+
109+
builder = cls.get_builder()
110+
builder.data_source = data_source
111+
112+
if isinstance(data_source, orm.RemoteData):
113+
if (remote_computer is not None) or (code is not None) or (structure is not None):
114+
warnings.warn(
115+
f'\nWhen providing a data source of type `RemoteData`, the other 3 inputs '
116+
f'(remote_computer, code and structure) will be ignored:'
117+
f'\n - remote_computer: {remote_computer}'
118+
f'\n - structure: {structure}'
119+
f'\n - code: {code}'
120+
)
121+
builder.pop('nscf')
122+
123+
if isinstance(data_source, orm.FolderData):
124+
125+
# If the code and structure are given: prepare NSCF and check computer compatibility
126+
# If the code and structure are abscent: check that a computer is provided
127+
# If only one was given: problematic
128+
129+
if (code is not None) and (structure is not None):
130+
131+
nscf_args = (code, structure, protocol)
132+
nscf_kwargs = kwargs
133+
134+
nscf_kwargs['overrides'] = {}
135+
if overrides is not None:
136+
nscf_kwargs['overrides'] = overrides.get('nscf', None)
137+
138+
# This is for easily setting defaults at each level of:
139+
# [overrides].nscf.pw.parameters.CONTROL.calculation
140+
last_layer = nscf_kwargs['overrides']
141+
last_layer = last_layer.setdefault('pw', {})
142+
last_layer = last_layer.setdefault('parameters', {})
143+
last_layer = last_layer.setdefault('CONTROL', {})
144+
145+
if last_layer.setdefault('calculation', 'nscf') != 'nscf':
146+
bad_value = last_layer['calculation']
147+
raise ValueError(
148+
f'The internal PwBaseWorkChain is for running an NSCF calculation, this should not\n'
149+
f'be overriden.\n'
150+
f'(Found overrides.nscf.pw.parametersCONTROL.calculation=`{bad_value}`)'
151+
)
152+
153+
nscf = PwBaseWorkChain.get_builder_from_protocol(*nscf_args, **nscf_kwargs)
154+
nscf['pw'].pop('parent_folder', None)
155+
nscf.pop('clean_workdir', None)
156+
builder.nscf = nscf
157+
158+
builder.remote_computer = code.computer
159+
if remote_computer is not None:
160+
warnings.warn(
161+
f'\nWhen providing a code for running the NSCF, any remote_computer given '
162+
f'will be ignored:'
163+
f'\n - remote_computer: `{remote_computer}`'
164+
f'\n - code provided: `{code}`'
165+
)
166+
167+
elif (code is None) and (structure is None):
168+
169+
if remote_computer is None:
170+
raise ValueError(
171+
f'If the `data_source` is a `FolderData` node, a `remote_computer` must also be\n'
172+
f'specified, or at least inferred from a `code` provided for running the NSCF.\n'
173+
f'(Currently remote_computer=`{remote_computer}` and code=`{code}`)'
174+
)
175+
builder.remote_computer = remote_computer
176+
builder.pop('nscf')
177+
178+
else:
179+
180+
raise ValueError(
181+
f'To run the NSCF both the code and structure must be specified.\n'
182+
f'(Currently code=`{code}` and structure=`{structure}`)'
183+
)
184+
185+
return builder
186+
187+
def should_run_nscf(self):
188+
"""If the 'nscf' input namespace was specified, we reconstruct the wave functions."""
189+
return 'nscf' in self.inputs
190+
191+
def run_transfer(self):
192+
"""Run the TransferCalcjob."""
193+
source_folder = self.inputs.data_source
194+
inputs = {
195+
'instructions': generate_instructions(source_folder)['instructions'],
196+
'source_nodes': {
197+
'source_node': source_folder
198+
},
199+
'metadata': {}
200+
}
201+
202+
if isinstance(source_folder, orm.FolderData):
203+
inputs['metadata']['call_link_label'] = 'transfer_put'
204+
if 'nscf' in self.inputs:
205+
inputs['metadata']['computer'] = self.inputs.nscf.pw.code.computer
206+
else:
207+
inputs['metadata']['computer'] = self.inputs.remote_computer
208+
209+
elif isinstance(source_folder, orm.RemoteData):
210+
inputs['metadata']['call_link_label'] = 'transfer_get'
211+
inputs['metadata']['computer'] = source_folder.computer
212+
213+
running = self.submit(TransferCalcjob, **inputs)
214+
self.report(f'launching TransferCalcjob<{running.pk}>')
215+
return ToContext(transfer_calcjob=running)
216+
217+
def inspect_transfer(self):
218+
"""Verify that the TransferCalcjob finished successfully."""
219+
source0_node = self.inputs.data_source
220+
calcjob_node = self.ctx.transfer_calcjob
221+
222+
if not calcjob_node.is_finished_ok:
223+
self.report(f'TransferCalcjob failed with exit status {calcjob_node.exit_status}')
224+
return self.exit_codes.ERROR_SUB_PROCESS_FAILED_TRANSFER
225+
226+
if isinstance(source0_node, orm.FolderData):
227+
self.ctx.last_output = calcjob_node.outputs.remote_folder
228+
229+
elif isinstance(source0_node, orm.RemoteData):
230+
self.ctx.last_output = calcjob_node.outputs.retrieved
231+
232+
def run_nscf(self):
233+
"""Run the PwBaseWorkChain in nscf mode on the restart folder."""
234+
inputs = self.exposed_inputs(PwBaseWorkChain, namespace='nscf')
235+
inputs['metadata']['call_link_label'] = 'nscf'
236+
inputs['pw']['parent_folder'] = self.ctx.last_output
237+
238+
running = self.submit(PwBaseWorkChain, **inputs)
239+
self.report(f'launching PwBaseWorkChain<{running.pk}> in nscf mode')
240+
return ToContext(workchain_nscf=running)
241+
242+
def inspect_nscf(self):
243+
"""Verify that the PwBaseWorkChain for the scf run finished successfully."""
244+
workchain = self.ctx.workchain_nscf
245+
246+
if not workchain.is_finished_ok:
247+
self.report(f'scf PwBaseWorkChain failed with exit status {workchain.exit_status}')
248+
return self.exit_codes.ERROR_SUB_PROCESS_FAILED_NSCF
249+
250+
self.ctx.last_output = workchain.outputs.remote_folder
251+
252+
def results(self):
253+
"""Attach the desired output nodes directly as outputs of the workchain."""
254+
self.report('workchain succesfully completed')
255+
self.out('output_data', self.ctx.last_output)
256+
257+
258+
##############################################################################################################
259+
def validate_inputs(inputs, _):
260+
"""Validate the inputs of the entire input namespace."""
261+
source_folder = inputs['data_source']
262+
263+
# FolderData: files must be there and code/computer compatibility
264+
if isinstance(source_folder, orm.FolderData):
265+
266+
error_message = ''
267+
if 'data-file-schema.xml' not in source_folder.list_object_names():
268+
error_message += f'Missing `data-file-schema.xml` on node {source_folder.pk}\n'
269+
if 'charge-density.dat' not in source_folder.list_object_names():
270+
error_message += f'Missing `charge-density.dat` on node {source_folder.pk}\n'
271+
if len(error_message) > 0:
272+
return error_message
273+
274+
if ('nscf' in inputs) and ('remote_computer' in inputs):
275+
computer_remote = inputs['remote_computer']
276+
computer_pwcode = inputs['nscf']['pw']['code'].computer
277+
if computer_remote.pk != computer_pwcode.pk:
278+
return (
279+
f'\nSome of the inputs provided are associated to different computers:'
280+
f'\n - remote_computer: {computer_remote}'
281+
f'\n - nscf.pw.code: {computer_pwcode}'
282+
)
283+
284+
elif ('nscf' not in inputs) and ('remote_computer' not in inputs):
285+
return 'The source is a FolderData and no code or remote_computer was provided.'
286+
287+
# RemoteData: RemoteData and remote_computer compatibility and warn if NSCF was provided
288+
if isinstance(source_folder, orm.RemoteData):
289+
290+
if 'remote_computer' in inputs:
291+
computer_remote = inputs['remote_computer']
292+
computer_source = source_folder.computer
293+
if computer_remote.pk != computer_source.pk:
294+
return (
295+
f'\nSome of the inputs provided are associated to different computers:'
296+
f'\n - remote_computer: {computer_remote}'
297+
f'\n - source_folder: {computer_source}'
298+
)
299+
300+
if 'nscf' in inputs:
301+
warnings.warn(
302+
f'\nThe `source_folder` (PK={source_folder.pk}) is a RemoteData node, so the data will be'
303+
f'\nretrieved and the NSCF input will be ignored'
304+
)
305+
306+
307+
def validate_nscf(value, _):
308+
"""Validate the inputs of the nscf input namespace."""
309+
parameters = value['pw']['parameters'].get_dict()
310+
if parameters.get('CONTROL', {}).get('calculation', 'scf') != 'nscf':
311+
return '`CONTOL.calculation` in `nscf.pw.parameters` is not set to `nscf`.'
312+
313+
314+
##############################################################################################################
315+
@calcfunction
316+
def generate_instructions(source_folder):
317+
"""Generate the instruction node to be used for copying the files."""
318+
319+
# Paths in the QE run folder
320+
schema_qepath = 'out/aiida.save/data-file-schema.xml'
321+
charge_qepath = 'out/aiida.save/charge-density.dat'
322+
pawtxt_qepath = 'out/aiida.save/paw.txt'
323+
324+
# Paths in the local node
325+
schema_dbpath = 'data-file-schema.xml'
326+
charge_dbpath = 'charge-density.dat'
327+
pawtxt_dbpath = 'paw.txt'
328+
329+
# Transfer from local to remote
330+
if isinstance(source_folder, orm.FolderData):
331+
instructions = {'retrieve_files': False, 'local_files': []}
332+
instructions['local_files'].append(('source_node', schema_dbpath, schema_qepath))
333+
instructions['local_files'].append(('source_node', charge_dbpath, charge_qepath))
334+
335+
if 'paw.txt' in source_folder.list_object_names():
336+
instructions['local_files'].append(('source_node', pawtxt_dbpath, pawtxt_qepath))
337+
338+
# Transfer from remote to local
339+
elif isinstance(source_folder, orm.RemoteData):
340+
instructions = {'retrieve_files': True, 'symlink_files': []}
341+
instructions['symlink_files'].append(('source_node', schema_qepath, schema_dbpath))
342+
instructions['symlink_files'].append(('source_node', charge_qepath, charge_dbpath))
343+
instructions['symlink_files'].append(('source_node', pawtxt_qepath, pawtxt_dbpath))
344+
345+
return {'instructions': orm.Dict(dict=instructions)}

‎setup.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
"quantumespresso.pw.band_structure = aiida_quantumespresso.workflows.pw.band_structure:PwBandStructureWorkChain",
6666
"quantumespresso.q2r.base = aiida_quantumespresso.workflows.q2r.base:Q2rBaseWorkChain",
6767
"quantumespresso.matdyn.base = aiida_quantumespresso.workflows.matdyn.base:MatdynBaseWorkChain",
68+
"quantumespresso.density.transfer = aiida_quantumespresso.workflows.density.transfer:TransferDensityWorkChain",
6869
"quantumespresso.pdos = aiida_quantumespresso.workflows.pdos:PdosWorkChain"
6970
],
7071
"console_scripts": [

0 commit comments

Comments
 (0)