Add support for electron-phonon calculations

src/aiida_quantumespresso/calculations/matdyn.py

@@ -1,8 +1,12 @@
 # -*- coding: utf-8 -*-
 """`CalcJob` implementation for the matdyn.x code of Quantum ESPRESSO."""
+from pathlib import Path
+
 from aiida import orm
 
+from aiida_quantumespresso.calculations import _uppercase_dict
 from aiida_quantumespresso.calculations.namelists import NamelistsCalculation
+from aiida_quantumespresso.calculations.ph import PhCalculation
 from aiida_quantumespresso.data.force_constants import ForceConstantsData
 
 
@@ -18,7 +22,6 @@ class MatdynCalculation(NamelistsCalculation):
         ('INPUT', 'flfrq', _PHONON_FREQUENCIES_NAME),  # output frequencies
         ('INPUT', 'flvec', _PHONON_MODES_NAME),  # output displacements
         ('INPUT', 'fldos', _PHONON_DOS_NAME),  # output density of states
-        ('INPUT', 'q_in_cryst_coord', True),  # kpoints always in crystal coordinates
     ]
 
     _internal_retrieve_list = [_PHONON_FREQUENCIES_NAME, _PHONON_DOS_NAME]
@@ -31,10 +34,12 @@ def define(cls, spec):
         super().define(spec)
         spec.input('force_constants', valid_type=ForceConstantsData, required=True)
         spec.input('kpoints', valid_type=orm.KpointsData, help='Kpoints on which to calculate the phonon frequencies.')
+        spec.input('parent_folder', valid_type=orm.RemoteData, required=False)
         spec.inputs.validator = cls._validate_inputs
 
         spec.output('output_parameters', valid_type=orm.Dict)
-        spec.output('output_phonon_bands', valid_type=orm.BandsData)
+        spec.output('output_phonon_bands', valid_type=orm.BandsData, required=False)
+        spec.output('output_phonon_dos', valid_type=orm.XyData, required=False)
         spec.default_output_node = 'output_parameters'
 
         spec.exit_code(310, 'ERROR_OUTPUT_STDOUT_READ',
@@ -43,6 +48,8 @@ def define(cls, spec):
             message='The stdout output file was incomplete probably because the calculation got interrupted.')
         spec.exit_code(330, 'ERROR_OUTPUT_FREQUENCIES',
             message='The output frequencies file could not be read from the retrieved folder.')
+        spec.exit_code(330, 'ERROR_OUTPUT_DOS',
+            message='The output DOS file could not be read from the retrieved folder.')
         spec.exit_code(410, 'ERROR_OUTPUT_KPOINTS_MISSING',
             message='Number of kpoints not found in the output data')
         spec.exit_code(411, 'ERROR_OUTPUT_KPOINTS_INCOMMENSURATE',
@@ -55,11 +62,20 @@ def _validate_inputs(value, _):
         if 'parameters' in value:
             parameters = value['parameters'].get_dict()
         else:
-            parameters = {}
+            parameters = {'INPUT': {}}
 
-        if parameters.get('INPUT', {}).get('flfrc', None) is not None:
+        if parameters['INPUT'].get('flfrc', None) is not None:
             return '`INPUT.flfrc` is set automatically from the `force_constants` input.'
 
+        if parameters['INPUT'].get('q_in_cryst_coord', None) is not None:
+            return '`INPUT.q_in_cryst_coords` is always set to `.true.` if `kpoints` input corresponds to list.'
+
+        if 'parent_folder' in value and not parameters.get('INPUT').get('la2F', False):
+            return (
+                'The `parent_folder` input is only used to calculate the el-ph coefficients but `la2F` is not set '
+                'to `.true.`'
+            )
+
     def generate_input_file(self, parameters):  # pylint: disable=arguments-differ
         """Generate namelist input_file content given a dict of parameters.
 
@@ -68,21 +84,32 @@ def generate_input_file(self, parameters):  # pylint: disable=arguments-differ
         :return: 'str' containing the input_file content a plain text.
         """
         kpoints = self.inputs.kpoints
+        append_string = ''
+
         parameters.setdefault('INPUT', {})['flfrc'] = self.inputs.force_constants.filename
-        file_content = super().generate_input_file(parameters)
 
-        try:
-            kpoints_list = kpoints.get_kpoints()
-        except AttributeError:
-            kpoints_list = kpoints.get_kpoints_mesh(print_list=True)
+        # Calculating DOS requires (nk1,nk2,nk3), see
+        # https://gitlab.com/QEF/q-e/-/blob/b231a0d0174ad1853f191160389029aa14fba6e9/PHonon/PH/matdyn.f90#L82
+        if parameters['INPUT'].get('dos', False):
+            kpoints_mesh = kpoints.get_kpoints_mesh()[0]
+            parameters['INPUT']['nk1'] = kpoints_mesh[0]
+            parameters['INPUT']['nk2'] = kpoints_mesh[1]
+            parameters['INPUT']['nk3'] = kpoints_mesh[2]
+        else:
+            parameters['INPUT']['q_in_cryst_coord'] = True
+            try:
+                kpoints_list = kpoints.get_kpoints()
+            except AttributeError:
+                kpoints_list = kpoints.get_kpoints_mesh(print_list=True)
 
-        kpoints_string = [f'{len(kpoints_list)}']
-        for kpoint in kpoints_list:
-            kpoints_string.append('{:18.10f} {:18.10f} {:18.10f}'.format(*kpoint))  # pylint: disable=consider-using-f-string
+            kpoints_string = [f'{len(kpoints_list)}']
+            for kpoint in kpoints_list:
+                kpoints_string.append('{:18.10f} {:18.10f} {:18.10f}'.format(*kpoint))  # pylint: disable=consider-using-f-string
+            append_string = '\n'.join(kpoints_string) + '\n'
 
-        file_content += '\n'.join(kpoints_string) + '\n'
+        file_content = super().generate_input_file(parameters)
 
-        return file_content
+        return file_content + append_string
 
     def prepare_for_submission(self, folder):
         """Prepare the calculation job for submission by transforming input nodes into input files.
@@ -91,6 +118,10 @@ def prepare_for_submission(self, folder):
         contains lists of files that need to be copied to the remote machine before job submission, as well as file
         lists that are to be retrieved after job completion.
 
+        After calling the method of the parent `NamelistsCalculation` class, the input parameters are checked to see
+        if the `la2F` tag is set to true. In this case the remote symlink or copy list is set to the electron-phonon
+        directory, depending on the settings.
+
         :param folder: a sandbox folder to temporarily write files on disk.
         :return: :class:`~aiida.common.datastructures.CalcInfo` instance.
         """
@@ -99,4 +130,31 @@ def prepare_for_submission(self, folder):
         force_constants = self.inputs.force_constants
         calcinfo.local_copy_list.append((force_constants.uuid, force_constants.filename, force_constants.filename))
 
+        if 'settings' in self.inputs:
+            settings = _uppercase_dict(self.inputs.settings.get_dict(), dict_name='settings')
+        else:
+            settings = {}
+
+        if 'parameters' in self.inputs:
+            parameters = _uppercase_dict(self.inputs.parameters.get_dict(), dict_name='parameters')
+        else:
+            parameters = {}
+
+        source = self.inputs.get('parent_folder', None)
+
+        if source is not None and parameters.get('INPUT').get('la2F', False):
+
+            # pylint: disable=protected-access
+            dirpath = Path(source.get_remote_path()) / PhCalculation._FOLDER_ELECTRON_PHONON
+            remote_list = [(source.computer.uuid, str(dirpath), PhCalculation._FOLDER_ELECTRON_PHONON)]
+
+            # For el-ph calculations, _only_ the `elph_dir` should be copied from the parent folder
+            if settings.pop('PARENT_FOLDER_SYMLINK', False):
+                calcinfo.remote_symlink_list = remote_list
+            else:
+                calcinfo.remote_copy_list = remote_list
+
+            calcinfo.retrieve_list += [f'a2F.dos{i}' for i in range(1, 11)]
+            calcinfo.retrieve_list.append('lambda')
+
         return calcinfo

src/aiida_quantumespresso/calculations/ph.py

@@ -36,6 +36,7 @@ class PhCalculation(CalcJob):
     _DVSCF_PREFIX = 'dvscf'
     _DRHO_STAR_EXT = 'drho_rot'
     _FOLDER_DYNAMICAL_MATRIX = 'DYN_MAT'
+    _FOLDER_ELECTRON_PHONON = 'elph_dir'
     _VERBOSITY = 'high'
     _OUTPUT_DYNAMICAL_MATRIX_PREFIX = os.path.join(_FOLDER_DYNAMICAL_MATRIX, 'dynamical-matrix-')
 
@@ -228,25 +229,6 @@ def prepare_for_submission(self, folder):
         if not restart_flag:  # if it is a restart, it will be copied over
             folder.get_subfolder(self._FOLDER_DYNAMICAL_MATRIX, create=True)
 
-        with folder.open(self.metadata.options.input_filename, 'w') as infile:
-            for namelist_name in namelists_toprint:
-                infile.write(f'&{namelist_name}\n')
-                # namelist content; set to {} if not present, so that we leave an empty namelist
-                namelist = parameters.pop(namelist_name, {})
-                for key, value in sorted(namelist.items()):
-                    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:
@@ -284,14 +266,44 @@ def prepare_for_submission(self, folder):
                     os.path.join(parent_folder.get_remote_path(),
                                  self._FOLDER_DYNAMICAL_MATRIX), self._FOLDER_DYNAMICAL_MATRIX
                 ))
-
+                if parameters['INPUTPH'].get('electron_phonon', None) is not None:
+                    remote_symlink_list.append((
+                        parent_folder.computer.uuid,
+                        os.path.join(parent_folder.get_remote_path(),
+                                     self._FOLDER_ELECTRON_PHONON), self._FOLDER_ELECTRON_PHONON
+                    ))
             else:
                 # copy the dynamical matrices
                 # no need to copy the _ph0, since I copied already the whole ./out folder
                 remote_copy_list.append((
                     parent_folder.computer.uuid,
                     os.path.join(parent_folder.get_remote_path(), self._FOLDER_DYNAMICAL_MATRIX), '.'
                 ))
+                if parameters['INPUTPH'].get('electron_phonon', None) is not None:
+                    remote_copy_list.append((
+                        parent_folder.computer.uuid,
+                        os.path.join(parent_folder.get_remote_path(),
+                                     self._FOLDER_ELECTRON_PHONON), self._FOLDER_ELECTRON_PHONON
+                    ))
+
+        with folder.open(self.metadata.options.input_filename, 'w') as infile:
+            for namelist_name in namelists_toprint:
+                infile.write(f'&{namelist_name}\n')
+                # namelist content; set to {} if not present, so that we leave an empty namelist
+                namelist = parameters.pop(namelist_name, {})
+                for key, value in sorted(namelist.items()):
+                    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()))}'
+            )
 
         # Create an `.EXIT` file if `only_initialization` flag in `settings` is set to `True`
         if settings.pop('ONLY_INITIALIZATION', False):

src/aiida_quantumespresso/calculations/q2r.py

@@ -1,8 +1,10 @@
 # -*- coding: utf-8 -*-
 """`CalcJob` implementation for the q2r.x code of Quantum ESPRESSO."""
+from pathlib import Path
 
 from aiida import orm
 
+from aiida_quantumespresso.calculations import _uppercase_dict
 from aiida_quantumespresso.calculations.namelists import NamelistsCalculation
 from aiida_quantumespresso.calculations.ph import PhCalculation
 from aiida_quantumespresso.data.force_constants import ForceConstantsData
@@ -39,3 +41,44 @@ def define(cls, spec):
         spec.exit_code(330, 'ERROR_READING_FORCE_CONSTANTS_FILE',
             message='The force constants file could not be read.')
         # yapf: enable
+
+    def prepare_for_submission(self, folder):
+        """Prepare the calculation job for submission by transforming input nodes into input files.
+
+        In addition to the input files being written to the sandbox folder, a `CalcInfo` instance will be returned that
+        contains lists of files that need to be copied to the remote machine before job submission, as well as file
+        lists that are to be retrieved after job completion.
+
+        After calling the method of the parent `NamelistsCalculation` class, the input parameters are checked to see
+        if the `la2F` tag is set to true. In this case the electron-phonon directory is added to the remote symlink or
+        copy list, depending on the settings.
+
+        :param folder: a sandbox folder to temporarily write files on disk.
+        :return: :class:`~aiida.common.datastructures.CalcInfo` instance.
+        """
+        calcinfo = super().prepare_for_submission(folder)
+
+        if 'settings' in self.inputs:
+            settings = _uppercase_dict(self.inputs.settings.get_dict(), dict_name='settings')
+        else:
+            settings = {}
+
+        if 'parameters' in self.inputs:
+            parameters = _uppercase_dict(self.inputs.parameters.get_dict(), dict_name='parameters')
+        else:
+            parameters = {'INPUT': {}}
+
+        source = self.inputs.get('parent_folder', None)
+
+        if source is not None and 'parameters' in self.inputs:
+
+            if parameters.get('INPUT').get('la2F', False):
+
+                symlink = settings.pop('PARENT_FOLDER_SYMLINK', False)
+                remote_list = calcinfo.remote_symlink_list if symlink else calcinfo.remote_copy_list
+
+                # pylint: disable=protected-access
+                dirpath = Path(source.get_remote_path()) / PhCalculation._FOLDER_ELECTRON_PHONON
+                remote_list.append((source.computer.uuid, str(dirpath), PhCalculation._FOLDER_ELECTRON_PHONON))
+
+        return calcinfo

src/aiida_quantumespresso/parsers/matdyn.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 from aiida import orm
+import numpy
 from qe_tools import CONSTANTS
 
+from aiida_quantumespresso.calculations import _uppercase_dict
 from aiida_quantumespresso.calculations.matdyn import MatdynCalculation
 
 from .base import Parser
@@ -15,6 +17,7 @@ def parse(self, **kwargs):
         retrieved = self.retrieved
         filename_stdout = self.node.get_option('output_filename')
         filename_frequencies = MatdynCalculation._PHONON_FREQUENCIES_NAME
+        filename_dos = MatdynCalculation._PHONON_DOS_NAME
 
         if filename_stdout not in retrieved.base.repository.list_object_names():
             return self.exit(self.exit_codes.ERROR_OUTPUT_STDOUT_READ)
@@ -23,7 +26,7 @@ def parse(self, **kwargs):
             return self.exit(self.exit_codes.ERROR_OUTPUT_STDOUT_INCOMPLETE)
 
         if filename_frequencies not in retrieved.base.repository.list_object_names():
-            return self.exit(self.exit_codes.ERROR_OUTPUT_STDOUT_READ)
+            return self.exit(self.exit_codes.ERROR_OUTPUT_FREQUENCIES)
 
         # Extract the kpoints from the input data and create the `KpointsData` for the `BandsData`
         try:
@@ -36,23 +39,46 @@ def parse(self, **kwargs):
 
         parsed_data = parse_raw_matdyn_phonon_file(retrieved.base.repository.get_object_content(filename_frequencies))
 
-        try:
-            num_kpoints = parsed_data.pop('num_kpoints')
-        except KeyError:
-            return self.exit(self.exit_codes.ERROR_OUTPUT_KPOINTS_MISSING)
+        if 'parameters' in self.node.inputs:
+            parameters = _uppercase_dict(self.node.inputs.parameters.get_dict(), dict_name='parameters')
+        else:
+            parameters = {'INPUT': {}}
+
+        if parameters['INPUT'].get('dos', False):
+
+            if filename_dos not in retrieved.base.repository.list_object_names():
+                return self.exit(self.exit_codes.ERROR_OUTPUT_DOS)
+
+            parsed_data.pop('phonon_bands', None)
+
+            with retrieved.open(filename_dos) as handle:
+                dos_array = numpy.genfromtxt(handle)
+
+            output_dos = orm.XyData()
+            output_dos.set_x(dos_array[:, 0], 'frequency', 'cm^(-1)')
+            output_dos.set_y(dos_array[:, 1], 'dos', 'states * cm')
+
+            self.out('output_phonon_dos', output_dos)
+
+        else:
+            try:
+                num_kpoints = parsed_data.pop('num_kpoints')
+            except KeyError:
+                return self.exit(self.exit_codes.ERROR_OUTPUT_KPOINTS_MISSING)
+
+            if num_kpoints != kpoints.shape[0]:
+                return self.exit(self.exit_codes.ERROR_OUTPUT_KPOINTS_INCOMMENSURATE)
 
-        if num_kpoints != kpoints.shape[0]:
-            return self.exit(self.exit_codes.ERROR_OUTPUT_KPOINTS_INCOMMENSURATE)
+            output_bands = orm.BandsData()
+            output_bands.set_kpointsdata(kpoints_for_bands)
+            output_bands.set_bands(parsed_data.pop('phonon_bands'), units='THz')
 
-        output_bands = orm.BandsData()
-        output_bands.set_kpointsdata(kpoints_for_bands)
-        output_bands.set_bands(parsed_data.pop('phonon_bands'), units='THz')
+            self.out('output_phonon_bands', output_bands)
 
         for message in parsed_data['warnings']:
             self.logger.error(message)
 
         self.out('output_parameters', orm.Dict(parsed_data))
-        self.out('output_phonon_bands', output_bands)
 
         return