providing support for aiida-atomistic in Pw CalcJob and BaseWorkChain

src/aiida_quantumespresso/calculations/__init__.py

@@ -18,13 +18,23 @@
 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
 
 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'

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,

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

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."""