From 7a8254d0d3bfe62584720a3b763c2b89f084df26 Mon Sep 17 00:00:00 2001 From: Miki Bonacci Date: Thu, 5 Dec 2024 18:39:36 +0100 Subject: [PATCH] providing support for aiida-atomistic in Pw CalcJob and BaseWorkChain --- .../calculations/__init__.py | 51 ++++++++++++-- src/aiida_quantumespresso/calculations/pw.py | 7 +- .../workflows/protocols/utils.py | 67 +++++++++++++++++++ .../workflows/pw/base.py | 54 +++++++++++++-- 4 files changed, 166 insertions(+), 13 deletions(-) diff --git a/src/aiida_quantumespresso/calculations/__init__.py b/src/aiida_quantumespresso/calculations/__init__.py index 05bd80e7b..a3733179c 100644 --- a/src/aiida_quantumespresso/calculations/__init__.py +++ b/src/aiida_quantumespresso/calculations/__init__.py @@ -18,6 +18,7 @@ from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData from aiida_quantumespresso.utils.convert import convert_input_to_namelist_entry from aiida_quantumespresso.utils.hubbard import HubbardUtils +from aiida_quantumespresso.utils.magnetic import MagneticUtils from .base import CalcJob from .helpers import QEInputValidationError @@ -25,6 +26,15 @@ LegacyUpfData = DataFactory('core.upf') UpfData = DataFactory('pseudo.upf') +LegacyStructureData = DataFactory('core.structure') # pylint: disable=invalid-name + +try: + StructureData = DataFactory('atomistic.structure') +except exceptions.MissingEntryPointError: + structures_classes = (LegacyStructureData,) +else: + structures_classes = (LegacyStructureData, StructureData) + class BasePwCpInputGenerator(CalcJob): """Base `CalcJob` for implementations for pw.x and cp.x of Quantum ESPRESSO.""" @@ -94,6 +104,8 @@ class BasePwCpInputGenerator(CalcJob): _use_kpoints = False + supported_properties = ['magmoms', 'hubbard'] + @classproperty def xml_filenames(cls): """Return a list of XML output filenames that can be written by a calculation. @@ -116,7 +128,7 @@ def define(cls, spec): spec.input('metadata.options.input_filename', valid_type=str, default=cls._DEFAULT_INPUT_FILE) spec.input('metadata.options.output_filename', valid_type=str, default=cls._DEFAULT_OUTPUT_FILE) spec.input('metadata.options.withmpi', valid_type=bool, default=True) # Override default withmpi=False - spec.input('structure', valid_type=orm.StructureData, + spec.input('structure', valid_type=(structures_classes), help='The input structure.') spec.input('parameters', valid_type=orm.Dict, help='The input parameters that are to be used to construct the input file.') @@ -168,6 +180,21 @@ def validate_inputs(cls, value, port_namespace): if any(key not in port_namespace for key in ('pseudos', 'structure')): return + if not isinstance(value['structure'], LegacyStructureData): + # we have the atomistic StructureData, so we need to check if all the defined properties are supported + plugin_check = value['structure'].check_plugin_support(cls.supported_properties) + if len(plugin_check) > 0: + raise NotImplementedError( + f'The input structure contains one or more unsupported properties \ + for this process: {plugin_check}' + ) + + if value['structure'].is_alloy or value['structure'].has_vacancies: + raise exceptions.InputValidationError( + 'The structure is an alloy or has vacancies. This is not allowed for \ + aiida-quantumespresso input structures.' + ) + # At this point, both ports are part of the namespace, and both are required so return an error message if any # of the two is missing. for key in ('pseudos', 'structure'): @@ -702,9 +729,17 @@ def _generate_PWCPinputdata(cls, parameters, settings, pseudos, structure, kpoin kpoints_card = ''.join(kpoints_card_list) del kpoints_card_list - # HUBBARD CARD - hubbard_card = HubbardUtils(structure).get_hubbard_card() if isinstance(structure, HubbardStructureData) \ - else None + # HUBBARD CARD and MAGNETIC NAMELIST + hubbard_card = None + magnetic_namelist = None + if isinstance(structure, HubbardStructureData): + hubbard_card = HubbardUtils(structure).get_hubbard_card() + elif len(structures_classes) == 2 and not isinstance(structure, LegacyStructureData): + # this means that we have the atomistic StructureData. + hubbard_card = HubbardUtils(structure).get_hubbard_card() if 'hubbard' \ + in structure.get_defined_properties() else None + magnetic_namelist = MagneticUtils(structure).generate_magnetic_namelist(input_params) if 'magmoms' in \ + structure.get_defined_properties() else None # =================== NAMELISTS AND CARDS ======================== try: @@ -734,6 +769,14 @@ def _generate_PWCPinputdata(cls, parameters, settings, pseudos, structure, kpoin 'namelists using the NAMELISTS inside the `settings` input node' ) from exception + if magnetic_namelist is not None: + if input_params['SYSTEM'].get('nspin', 1) == 1 and not input_params['SYSTEM'].get('noncolin', False): + raise exceptions.InputValidationError( + 'The structure has magnetic moments but the inputs are not set for \ + a magnetic calculation (`nspin`, `noncolin`)' + ) + input_params['SYSTEM'].update(magnetic_namelist) + inputfile = '' for namelist_name in namelists_toprint: inputfile += f'&{namelist_name}\n' diff --git a/src/aiida_quantumespresso/calculations/pw.py b/src/aiida_quantumespresso/calculations/pw.py index 9a489e671..9df223f89 100644 --- a/src/aiida_quantumespresso/calculations/pw.py +++ b/src/aiida_quantumespresso/calculations/pw.py @@ -5,10 +5,13 @@ from aiida import orm from aiida.common.lang import classproperty +from aiida.orm import StructureData as LegacyStructureData from aiida.plugins import factories from aiida_quantumespresso.calculations import BasePwCpInputGenerator +StructureData = factories.DataFactory('atomistic.structure') + class PwCalculation(BasePwCpInputGenerator): """`CalcJob` implementation for the pw.x code of Quantum ESPRESSO.""" @@ -69,13 +72,13 @@ def define(cls, spec): 'will not fail if the XML file is missing in the retrieved folder.') spec.input('kpoints', valid_type=orm.KpointsData, help='kpoint mesh or kpoint path') - spec.input('hubbard_file', valid_type=orm.SinglefileData, required=False, + spec.input('hubbard_file', valid_type=(StructureData, LegacyStructureData), required=False, help='SinglefileData node containing the output Hubbard parameters from a HpCalculation') spec.inputs.validator = cls.validate_inputs spec.output('output_parameters', valid_type=orm.Dict, help='The `output_parameters` output node of the successful calculation.') - spec.output('output_structure', valid_type=orm.StructureData, required=False, + spec.output('output_structure', valid_type=(StructureData, LegacyStructureData), required=False, help='The `output_structure` output node of the successful calculation if present.') spec.output('output_trajectory', valid_type=orm.TrajectoryData, required=False) spec.output('output_band', valid_type=orm.BandsData, required=False, diff --git a/src/aiida_quantumespresso/workflows/protocols/utils.py b/src/aiida_quantumespresso/workflows/protocols/utils.py index 664fe17c9..ed43b797a 100644 --- a/src/aiida_quantumespresso/workflows/protocols/utils.py +++ b/src/aiida_quantumespresso/workflows/protocols/utils.py @@ -152,3 +152,70 @@ def get_starting_magnetization( starting_magnetization[kind.name] = magnetization return starting_magnetization + + +def get_starting_magnetization_noncolin( + structure: StructureData, + pseudo_family: PseudoPotentialFamily, + initial_magnetic_moments: Optional[dict] = None +) -> tuple: + """Return the dictionary with starting magnetization for each kind in the structure. + + :param structure: the structure. + :param pseudo_family: pseudopotential family. + :param initial_magnetic_moments: dictionary mapping each kind in the structure to its magnetic moment. + :returns: dictionary of starting magnetizations. + """ + # try: + # structure.mykinds + # except AttributeError: + # raise TypeError(f"structure<{structure.pk}> do not have magmom") + starting_magnetization = {} + angle1 = {} + angle2 = {} + + if initial_magnetic_moments is not None: + + nkinds = len(structure.kinds) + + if sorted(initial_magnetic_moments.keys()) != sorted(structure.get_kind_names()): + raise ValueError(f'`initial_magnetic_moments` needs one value for each of the {nkinds} kinds.') + + for kind in structure.kinds: + magmom = initial_magnetic_moments[kind.name] + if isinstance(magmom, Union[int, float]): + starting_magnetization[kind.name] = magmom / pseudo_family.get_pseudo(element=kind.symbol).z_valence + angle1[kind.name] = 0.0 + angle2[kind.name] = 0.0 + else: # tuple of 3 float (r, theta, phi) + starting_magnetization[kind.name + ] = 2 * magmom[0] / pseudo_family.get_pseudo(element=kind.symbol).z_valence + angle1[kind.name] = magmom[1] + angle2[kind.name] = magmom[2] + try: + structure.mykinds + except AttributeError: + # Normal StructureData, no magmom in structure + magnetic_parameters = get_magnetization_parameters() + + for kind in structure.kinds: + magnetic_moment = magnetic_parameters[kind.symbol]['magmom'] + + if magnetic_moment == 0: + magnetization = magnetic_parameters['default_magnetization'] + else: + z_valence = pseudo_family.get_pseudo(element=kind.symbol).z_valence + magnetization = magnetic_moment / float(z_valence) + + starting_magnetization[kind.name] = magnetization + angle1[kind.name] = 0.0 + angle2[kind.name] = 0.0 + else: + # Self defined myStructureData, read magmom from structure + for kind in structure.mykinds: + magmom = kind.get_magmom_coord() + starting_magnetization[kind.name] = 2 * magmom[0] / pseudo_family.get_pseudo(element=kind.symbol).z_valence + angle1[kind.name] = magmom[1] + angle2[kind.name] = magmom[2] + + return starting_magnetization, angle1, angle2 diff --git a/src/aiida_quantumespresso/workflows/pw/base.py b/src/aiida_quantumespresso/workflows/pw/base.py index 4a3ab0381..76445ff6f 100644 --- a/src/aiida_quantumespresso/workflows/pw/base.py +++ b/src/aiida_quantumespresso/workflows/pw/base.py @@ -4,7 +4,7 @@ from aiida.common import AttributeDict, exceptions from aiida.common.lang import type_check from aiida.engine import BaseRestartWorkChain, ExitCode, ProcessHandlerReport, process_handler, while_ -from aiida.plugins import CalculationFactory, GroupFactory +from aiida.plugins import CalculationFactory, DataFactory, GroupFactory from aiida_quantumespresso.calculations.functions.create_kpoints_from_distance import create_kpoints_from_distance from aiida_quantumespresso.common.types import ElectronicType, RestartType, SpinType @@ -17,6 +17,12 @@ PseudoDojoFamily = GroupFactory('pseudo.family.pseudo_dojo') CutoffsPseudoPotentialFamily = GroupFactory('pseudo.family.cutoffs') +try: + StructureData = DataFactory('atomistic.structure') + HAS_ATOMISTIC = True +except ImportError: + HAS_ATOMISTIC = False + class PwBaseWorkChain(ProtocolMixin, BaseRestartWorkChain): """Workchain to run a Quantum ESPRESSO pw.x calculation with automated error handling and restarts.""" @@ -131,7 +137,11 @@ def get_builder_from_protocol( the ``CalcJobs`` that are nested in this work chain. :return: a process builder instance with all inputs defined ready for launch. """ - from aiida_quantumespresso.workflows.protocols.utils import get_starting_magnetization, recursive_merge + from aiida_quantumespresso.workflows.protocols.utils import ( + get_starting_magnetization, + get_starting_magnetization_noncolin, + recursive_merge, + ) if isinstance(code, str): code = orm.load_code(code) @@ -143,7 +153,7 @@ def get_builder_from_protocol( if electronic_type not in [ElectronicType.METAL, ElectronicType.INSULATOR]: raise NotImplementedError(f'electronic type `{electronic_type}` is not supported.') - if spin_type not in [SpinType.NONE, SpinType.COLLINEAR]: + if spin_type not in [SpinType.NONE, SpinType.COLLINEAR, SpinType.NON_COLLINEAR]: raise NotImplementedError(f'spin type `{spin_type}` is not supported.') if initial_magnetic_moments is not None and spin_type is not SpinType.COLLINEAR: @@ -189,10 +199,21 @@ def get_builder_from_protocol( parameters['SYSTEM'].pop('degauss') parameters['SYSTEM'].pop('smearing') - if spin_type is SpinType.COLLINEAR: - starting_magnetization = get_starting_magnetization(structure, pseudo_family, initial_magnetic_moments) - parameters['SYSTEM']['starting_magnetization'] = starting_magnetization - parameters['SYSTEM']['nspin'] = 2 + if isinstance(structure, orm.StructureData): + if spin_type is SpinType.COLLINEAR: + starting_magnetization = get_starting_magnetization(structure, pseudo_family, initial_magnetic_moments) + parameters['SYSTEM']['starting_magnetization'] = starting_magnetization + parameters['SYSTEM']['nspin'] = 2 + + if spin_type is SpinType.NON_COLLINEAR: + starting_magnetization_noncolin, angle1, angle2 = get_starting_magnetization_noncolin( + structure=structure, pseudo_family=pseudo_family, initial_magnetic_moments=initial_magnetic_moments + ) + parameters['SYSTEM']['starting_magnetization'] = starting_magnetization_noncolin + parameters['SYSTEM']['angle1'] = angle1 + parameters['SYSTEM']['angle2'] = angle2 + parameters['SYSTEM']['noncolin'] = True + parameters['SYSTEM']['nspin'] = 4 # If overrides are provided, they are considered absolute if overrides: @@ -284,6 +305,25 @@ def validate_kpoints(self): self.ctx.inputs.kpoints = kpoints + def validate_structure(self,): + """Validate the structure input for the workflow. + + This method checks if the structure has atomistic properties and if it is supported by the PwCalculation plugin. + If the structure contains unsupported properties, a new structure is generated without those properties. + + Modifies: + self.inputs.pw.structure: Updates the structure to a new one without unsupported properties if necessary. + """ + if HAS_ATOMISTIC: + # do we want to do this, or return a warning, or except? + from aiida_atomistic.data.structure.utils import generate_striped_structure # pylint: disable=import-error + plugin_check = self.inputs.pw.structure.check_plugin_support(PwCalculation.supported_properties) + if len(plugin_check) > 0: + # Generate a new StructureData without the unsupported properties. + self.inputs.pw.structure = generate_striped_structure( + self.inputs.pw.structure, orm.List(list(plugin_check)) + ) + def set_restart_type(self, restart_type, parent_folder=None): """Set the restart type for the next iteration."""