Skip to content

Commit

Permalink
EpwCalculation: Allow restart from previous epw.x run
Browse files Browse the repository at this point in the history
  • Loading branch information
@mbercx
mbercx committed Apr 16, 2023
1 parent 74d25d1 commit 38d2242
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 125 deletions.
207 changes: 88 additions & 119 deletions src/aiida_quantumespresso/calculations/epw.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
# -*- coding: utf-8 -*-
"""Plugin to create a Quantum Espresso epw.x input file."""
import os
from pathlib import Path

from aiida import orm
from aiida.common import datastructures, exceptions
import numpy as np

from aiida_quantumespresso.calculations import _lowercase_dict, _uppercase_dict
from aiida_quantumespresso.calculations.ph import PhCalculation
from aiida_quantumespresso.calculations.pw import PwCalculation
from aiida_quantumespresso.utils.convert import convert_input_to_namelist_entry

from .base import CalcJob
Expand All @@ -30,9 +31,7 @@ class EpwCalculation(CalcJob):
_DEFAULT_OUTPUT_FILE = 'aiida.out'
_OUTPUT_XML_TENSOR_FILE_NAME = 'tensors.xml'
_OUTPUT_SUBFOLDER = './out/'
_SAVE_PREFIX = '/save/'
_FOLDER_SAVE = 'save'
_VERBOSITY = 'high'
_FOLDER_DYNAMICAL_MATRIX = 'DYN_MAT'

# Not using symlink in pw to allow multiple nscf to run on top of the same scf
Expand All @@ -52,9 +51,12 @@ def define(cls, spec):
spec.input('qfpoints', valid_type=orm.KpointsData, help='fine qpoint mesh')
spec.input('parameters', valid_type=orm.Dict, help='')
spec.input('settings', valid_type=orm.Dict, required=False, help='')
spec.input('parent_folder_nscf', valid_type=orm.RemoteData,
help='the folder of a completed nscf `PwCalculation`')
spec.input('parent_folder_ph', valid_type=orm.RemoteData, help='the folder of a completed `PhCalculation`')
spec.input('parent_folder_nscf', required=False, valid_type=orm.RemoteData,
help='the folder of a completed nscf `PwCalculation`')
spec.input('parent_folder_ph', required=False, valid_type=orm.RemoteData,
help='the folder of a completed `PhCalculation`')
spec.input('parent_folder_epw', required=False, valid_type=(orm.RemoteData, orm.RemoteStashFolderData),
help='folder that contains all files required to restart an `EpwCalculation`')
# yapf: enable

def prepare_for_submission(self, folder):
Expand All @@ -68,7 +70,7 @@ def prepare_for_submission(self, folder):
:return: :class:`~aiida.common.datastructures.CalcInfo` instance.
"""

# pylint: disable=too-many-statements,too-many-branches
# pylint: disable=too-many-statements,too-many-branches, protected-access

def test_offset(offset):
"""Check if the grid has an offset."""
Expand All @@ -82,48 +84,99 @@ def test_offset(offset):
remote_copy_list = []
remote_symlink_list = []

parameters = _uppercase_dict(self.inputs.parameters.get_dict(), dict_name='parameters')
parameters = {k: _lowercase_dict(v, dict_name=k) for k, v in parameters.items()}

if 'INPUTEPW' not in parameters:
raise exceptions.InputValidationError('required namelist INPUTEPW not specified')

if 'settings' in self.inputs:
settings = _uppercase_dict(self.inputs.settings.get_dict(), dict_name='settings')
else:
settings = {}

# Copy nscf folder
parent_folder_nscf = self.inputs.parent_folder_nscf
parent_calc_nscf = parent_folder_nscf.creator
remote_list = remote_symlink_list if settings.pop(
'PARENT_FOLDER_SYMLINK', self._default_symlink_usage
) else remote_copy_list

if parent_calc_nscf is None:
raise exceptions.NotExistent(f'parent_folder<{parent_folder_nscf.pk}> has no parent calculation')
if 'parent_folder_nscf' in self.inputs:
parent_folder_nscf = self.inputs.parent_folder_nscf

# Also, the parent calculation must be on the same computer
if not self.node.computer.uuid == parent_calc_nscf.computer.uuid:
computer_label = parent_calc_nscf.computer.get_name()
raise exceptions.InputValidationError(
f'Calculation has to be launched on the same computer as that of the parent: {computer_label}'
)
remote_list.append((
parent_folder_nscf.computer.uuid,
Path(parent_folder_nscf.get_remote_path(), PwCalculation._OUTPUT_SUBFOLDER).as_posix(),
self._OUTPUT_SUBFOLDER,
))

# put by default, default_parent_output_folder = ./out
parent_calc_out_subfolder_nscf = parent_calc_nscf.process_class._OUTPUT_SUBFOLDER # pylint: disable=protected-access
if 'parent_folder_ph' in self.inputs:
parent_folder_ph = self.inputs.parent_folder_ph

# Now phonon folder
parent_folder_ph = self.inputs.parent_folder_ph
parent_calc_ph = parent_folder_ph.creator
# Create the save folder with dvscf and dyn files
folder.get_subfolder(self._FOLDER_SAVE, create=True)

# Also, the parent calculation must be on the same computer
if not self.node.computer.uuid == parent_calc_ph.computer.uuid:
computer_label = parent_calc_nscf.computer.get_name()
raise exceptions.InputValidationError(
f'Calculation has to be launched on the same computer as that of the parent: {computer_label}'
# List of IBZ q-point to be added below EPW. To be removed when removed from EPW.
qibz_ar = []
for key, value in sorted(parent_folder_ph.creator.outputs.output_parameters.get_dict().items()):
if key.startswith('dynamical_matrix_'):
qibz_ar.append(value['q_point'])

nqpt = len(qibz_ar)

# Append the required contents of the `save` folder to the remove copy list, copied from the `ph.x`
# calculation

prefix = self._PREFIX
outdir = self._OUTPUT_SUBFOLDER
fildvscf = PhCalculation._DVSCF_PREFIX
fildyn = PhCalculation._OUTPUT_DYNAMICAL_MATRIX_PREFIX

ph_path = Path(parent_folder_ph.get_remote_path())

remote_list.append(
(parent_folder_ph.computer.uuid, Path(ph_path, outdir, '_ph0', f'{prefix}.phsave').as_posix(), 'save')
)

# I put the first-level keys as uppercase (i.e., namelist and card names) and the second-level keys as lowercase
parameters = _uppercase_dict(self.inputs.parameters.get_dict(), dict_name='parameters')
parameters = {k: _lowercase_dict(v, dict_name=k) for k, v in parameters.items()}
for iqpt in range(1, nqpt + 1):
remote_list.append((
parent_folder_ph.computer.uuid,
Path(ph_path, outdir, '_ph0', '' if iqpt == 1 else f'{prefix}.q_{iqpt}',
f'{prefix}.{fildvscf}1').as_posix(), Path('save', f'{prefix}.dvscf_q{iqpt}').as_posix()
))
# remote_copy_list.append((
# parent_folder_ph.computer.uuid,
# Path(
# ph_path, outdir, '_ph0', '' if iqpt == 1 else f'{prefix}.q_{iqpt}', f'{prefix}.{fildvscf}_paw1'
# ).as_posix(),
# Path('save', f"{prefix}.dvscf_paw_q{iqpt}").as_posix()
# ))
remote_list.append((
parent_folder_ph.computer.uuid, Path(ph_path, f'{fildyn}{iqpt}').as_posix(),
Path('save', f'{prefix}.dyn_q{iqpt}').as_posix()
))

if 'INPUTEPW' not in parameters:
raise exceptions.InputValidationError('required namelist INPUTEPW not specified')
if 'parent_folder_epw' in self.inputs:

parent_folder_epw = self.inputs.parent_folder_epw
if isinstance(parent_folder_epw, orm.RemoteStashFolderData):
epw_path = Path(parent_folder_epw.target_basepath)
else:
epw_path = Path(parent_folder_epw.get_remote_path())

vme_fmt_dict = {
'dipole': 'dmedata.fmt',
'wannier': 'vmedata.fmt',
}

for filename in (
'crystal.fmt', 'epwdata.fmt', vme_fmt_dict[parameters['INPUTEPW']['vme']], f'{self._PREFIX}.kgmap',
f'{self._PREFIX}.kmap', f'{self._PREFIX}.ukk', self._OUTPUT_SUBFOLDER, self._FOLDER_SAVE
):
remote_list.append(
(parent_folder_epw.computer.uuid, Path(epw_path, filename).as_posix(), Path(filename).as_posix())
)

parameters['INPUTEPW']['outdir'] = self._OUTPUT_SUBFOLDER
parameters['INPUTEPW']['verbosity'] = self._VERBOSITY
parameters['INPUTEPW']['dvscf_dir'] = self._FOLDER_SAVE
parameters['INPUTEPW']['prefix'] = self._PREFIX

try:
Expand All @@ -132,7 +185,6 @@ def test_offset(offset):
parameters['INPUTEPW']['nq1'] = mesh[0]
parameters['INPUTEPW']['nq2'] = mesh[1]
parameters['INPUTEPW']['nq3'] = mesh[2]
postpend_text = None
except NotImplementedError as exception:
raise exceptions.InputValidationError('Cannot get the coarse q-point grid') from exception

Expand All @@ -142,7 +194,6 @@ def test_offset(offset):
parameters['INPUTEPW']['nk1'] = mesh[0]
parameters['INPUTEPW']['nk2'] = mesh[1]
parameters['INPUTEPW']['nk3'] = mesh[2]
postpend_text = None
except NotImplementedError as exception:
raise exceptions.InputValidationError('Cannot get the coarse k-point grid') from exception

Expand All @@ -152,7 +203,6 @@ def test_offset(offset):
parameters['INPUTEPW']['nqf1'] = mesh[0]
parameters['INPUTEPW']['nqf2'] = mesh[1]
parameters['INPUTEPW']['nqf3'] = mesh[2]
postpend_text = None
except NotImplementedError as exception:
raise exceptions.InputValidationError('Cannot get the fine q-point grid') from exception

Expand All @@ -162,7 +212,6 @@ def test_offset(offset):
parameters['INPUTEPW']['nkf1'] = mesh[0]
parameters['INPUTEPW']['nkf2'] = mesh[1]
parameters['INPUTEPW']['nkf3'] = mesh[2]
postpend_text = None
except NotImplementedError as exception:
raise exceptions.InputValidationError('Cannot get the fine k-point grid') from exception

Expand All @@ -177,28 +226,6 @@ def test_offset(offset):
except KeyError: # list of namelists not specified in the settings; do automatic detection
namelists_toprint = self._compulsory_namelists

# create the save folder with dvscf and dyn files.
folder.get_subfolder(self._FOLDER_SAVE, create=True)

# List of IBZ q-point to be added below EPW. To be removed when removed from EPW.
qibz_ar = []
for key, value in sorted(parent_folder_ph.creator.outputs.output_parameters.get_dict().items()):
if key.startswith('dynamical_matrix_'):
qibz_ar.append(value['q_point'])

qibz_node = orm.ArrayData()
qibz_node.set_array('qibz', np.array(qibz_ar))

list_of_points = qibz_node.get_array('qibz')
# Number of q-point in the irreducible Brillouin Zone.
nqpt = len(list_of_points[0, :])

# add here the list of point coordinates
if len(list_of_points) > 1:
postpend_text = f'{len(list_of_points)} cartesian\n'
for points in list_of_points:
postpend_text += '{0:18.10f} {1:18.10f} {2:18.10f} \n'.format(*points) # pylint: disable=consider-using-f-string

with folder.open(self.metadata.options.input_filename, 'w') as infile:
for namelist_name in namelists_toprint:
infile.write(f'&{namelist_name}\n')
Expand All @@ -208,70 +235,12 @@ def test_offset(offset):
infile.write(convert_input_to_namelist_entry(key, value))
infile.write('/\n')

# add list of qpoints if required
if postpend_text is not None:
infile.write(postpend_text)

if parameters:
raise exceptions.InputValidationError(
'The following namelists are specified in parameters, but are not valid namelists for the current type '
f'of calculation: {",".join(list(parameters.keys()))}'
)

# copy the parent scratch
symlink = settings.pop('PARENT_FOLDER_SYMLINK', self._default_symlink_usage) # a boolean
if symlink:
# I create a symlink to each file/folder in the parent ./out
folder.get_subfolder(self._OUTPUT_SUBFOLDER, create=True)

remote_symlink_list.append((
parent_folder_nscf.computer.uuid,
os.path.join(parent_folder_nscf.get_remote_path(), parent_calc_out_subfolder_nscf,
'*'), self._OUTPUT_SUBFOLDER
))

else:
# here I copy the whole folder ./out
remote_copy_list.append((
parent_folder_nscf.computer.uuid,
os.path.join(parent_folder_nscf.get_remote_path(),
parent_calc_out_subfolder_nscf), self._OUTPUT_SUBFOLDER
))

prefix = self._PREFIX

for iqpt in range(1, nqpt + 1):
label = str(iqpt)
tmp_path = os.path.join(self._FOLDER_DYNAMICAL_MATRIX, 'dynamical-matrix-0')
remote_copy_list.append((
parent_folder_ph.computer.uuid, os.path.join(parent_folder_ph.get_remote_path(),
tmp_path), 'save/' + prefix + '.dyn_q0'
))
tmp_path = os.path.join(self._FOLDER_DYNAMICAL_MATRIX, 'dynamical-matrix-' + label)
remote_copy_list.append((
parent_folder_ph.computer.uuid, os.path.join(parent_folder_ph.get_remote_path(),
tmp_path), 'save/' + prefix + '.dyn_q' + label
))

if iqpt == 1:
tmp_path = os.path.join(self._OUTPUT_SUBFOLDER, '_ph0/' + prefix + '.dvscf*')
remote_copy_list.append((
parent_folder_ph.computer.uuid, os.path.join(parent_folder_ph.get_remote_path(),
tmp_path), 'save/' + prefix + '.dvscf_q' + label
))
tmp_path = os.path.join(self._OUTPUT_SUBFOLDER, '_ph0/' + prefix + '.phsave')
remote_copy_list.append((
parent_folder_ph.computer.uuid, os.path.join(parent_folder_ph.get_remote_path(), tmp_path), 'save/'
))
else:
tmp_path = os.path.join(
self._OUTPUT_SUBFOLDER, '_ph0/' + prefix + '.q_' + label + '/' + prefix + '.dvscf*'
)
remote_copy_list.append((
parent_folder_ph.computer.uuid, os.path.join(parent_folder_ph.get_remote_path(),
tmp_path), 'save/' + prefix + '.dvscf_q' + label
))

codeinfo = datastructures.CodeInfo()
codeinfo.cmdline_params = (list(settings.pop('CMDLINE', [])) + ['-in', self.metadata.options.input_filename])
codeinfo.stdout_name = self.metadata.options.output_filename
Expand Down
7 changes: 1 addition & 6 deletions tests/calculations/test_epw/test_epw_default.in
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
&INPUTEPW
dis_froz_max = 8.5000000000d+00
dis_win_max = 18
dvscf_dir = './save/'
dvscf_dir = 'save'
elecselfen = .true.
elph = .true.
epbwrite = .true.
Expand All @@ -22,10 +22,5 @@
outdir = './out/'
prefix = 'aiida'
proj(1) = 'Si : sp3'
verbosity = 'high'
wannierize = .true.
/
3 cartesian
0.0000000000 0.0000000000 0.0000000000
0.0000000000 0.5000000000 0.5000000000
0.5000000000 0.5000000000 0.5000000000

0 comments on commit 38d2242

Please sign in to comment.