XAS: Enable Correct Parsing of Hubbard and Magnetic Data

src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py

@@ -12,6 +12,8 @@
 import numpy as np
 import spglib
 
+from aiida_quantumespresso.utils.hubbard import HubbardStructureData, HubbardUtils
+
 
 @calcfunction
 def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-statements
@@ -45,6 +47,11 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                     a molecule and not a periodic solid system. Required in order to instruct
                     the CF to use Pymatgen rather than spglib to determine the symmetry. The
                     CF will assume the structure to be a periodic solid if no input is given.
+        - use_element_types: a Bool object to indicate that symmetry analysis should consider all
+                             sites of the same element to be equal and ignore any special Kind names
+                             from the parent structure. For instance, use_element_types = True would
+                             consider sites for Kinds 'Si' and 'Si1' to be equivalent if both are sites
+                             containing silicon. Defaults to False otherwise.
         - spglib_settings: an optional Dict object containing overrides for the symmetry
                             tolerance parameters used by spglib (symmprec, angle_tolerance).
         - pymatgen_settings: an optional Dict object containing overrides for the symmetry
@@ -100,7 +107,11 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                 )
     else:
         elements_defined = False
-        abs_elements_list = [Kind.symbol for Kind in structure.kinds]
+        abs_elements_list = []
+        for kind in structure.kinds:
+            if kind.symbol not in abs_elements_list:
+                abs_elements_list.append(kind.symbol)
+
     if 'is_molecule_input' in unwrapped_kwargs.keys():
         is_molecule_input = unwrapped_kwargs['is_molecule_input'].value
         # If we are working with a molecule, check for pymatgen_settings
@@ -118,6 +129,21 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     else:
         spglib_kwargs = {}
 
+    if 'use_element_types' in unwrapped_kwargs.keys():
+        use_element_types = unwrapped_kwargs['use_element_types'].value
+    else:
+        use_element_types = False
+
+    if structure.node_type == 'data.quantumespresso.hubbard_structure.HubbardStructureData.':
+        is_hubbard_structure = True
+        if standardize_structure:
+            raise ValidationError(
+                'Incoming structure set to be standardized, but hubbard data has been found. '
+                'Please set ``standardize_structure`` to false in ``**kwargs`` to preserve the hubbard data.'
+            )
+    else:
+        is_hubbard_structure = False
+
     output_params = {}
     result = {}
 
@@ -203,38 +229,83 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     # Process a periodic system
     else:
         incoming_structure_tuple = structure_to_spglib_tuple(structure)
-
-        symmetry_dataset = spglib.get_symmetry_dataset(incoming_structure_tuple[0], **spglib_kwargs)
+        spglib_tuple = incoming_structure_tuple[0]
+        types_order = spglib_tuple[-1]
+        kinds_information = incoming_structure_tuple[1]
+        kinds_list = incoming_structure_tuple[2]
+
+        # We need a way to reliably convert type number into element, so we
+        # first create a mapping of assigned number to kind name then a mapping
+        # of assigned number to ``Kind``
+
+        type_name_mapping = {str(value): key for key, value in kinds_information.items()}
+        type_mapping_dict = {}
+
+        for key, value in type_name_mapping.items():
+            for kind in kinds_list:
+                if value == kind.name:
+                    type_mapping_dict[key] = kind
+
+        # if we want to treat all sites of the same element as equal,
+        # then we must briefly operate on a "cleaned" version of the
+        # structure tuple.
+        if use_element_types:
+            clean_structure_tuple = (spglib_tuple[0], spglib_tuple[1], [])
+            for i in spglib_tuple[2]:
+                if i >= 1000:
+                    new_i = int(i / 1000)
+                else:
+                    new_i = i
+                clean_structure_tuple[2].append(new_i)
+            symmetry_dataset = spglib.get_symmetry_dataset(clean_structure_tuple, **spglib_kwargs)
+        else:  # Otherwise, proceed as usual.
+            symmetry_dataset = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)
 
         # if there is no symmetry to exploit, or no standardization is desired, then we just use
         # the input structure in the following steps. This is done to account for the case where
         # the user has submitted an improper crystal for calculation work and doesn't want it to
         # be changed.
         if symmetry_dataset['number'] in [1, 2] or not standardize_structure:
-            standardized_structure_node = spglib_tuple_to_structure(incoming_structure_tuple[0])
+            standardized_structure_node = spglib_tuple_to_structure(spglib_tuple, kinds_information, kinds_list)
             structure_is_standardized = False
         else:  # otherwise, we proceed with the standardized structure.
-            standardized_structure_tuple = spglib.standardize_cell(incoming_structure_tuple[0], **spglib_kwargs)
-            standardized_structure_node = spglib_tuple_to_structure(standardized_structure_tuple)
+            standardized_structure_tuple = spglib.standardize_cell(spglib_tuple, **spglib_kwargs)
+            standardized_structure_node = spglib_tuple_to_structure(
+                standardized_structure_tuple, kinds_information, kinds_list
+            )
             # if we are standardizing the structure, then we need to update the symmetry
             # information for the standardized structure
             symmetry_dataset = spglib.get_symmetry_dataset(standardized_structure_tuple, **spglib_kwargs)
             structure_is_standardized = True
 
         equivalent_atoms_array = symmetry_dataset['equivalent_atoms']
-        element_types = symmetry_dataset['std_types']
+        if structure_is_standardized:
+            element_types = symmetry_dataset['std_types']
+        else:  # convert the 'std_types' from standardized to primitive cell
+            spglib_std_types = symmetry_dataset['std_types']
+            spglib_std_map_to_prim = symmetry_dataset['std_mapping_to_primitive']
+
+            map_std_pos_to_type = {}
+            for position, atom_type in zip(spglib_std_map_to_prim, spglib_std_types):
+                map_std_pos_to_type[str(position)] = atom_type
+            primitive_types = []
+            for position in symmetry_dataset['mapping_to_primitive']:
+                atom_type = map_std_pos_to_type[str(position)]
+                primitive_types.append(atom_type)
+            element_types = primitive_types
 
         equivalency_dict = {}
         index_counter = 0
         for symmetry_value, element_type in zip(equivalent_atoms_array, element_types):
             if elements_defined:  # only process the elements given in the list
                 if f'site_{symmetry_value}' in equivalency_dict:
                     equivalency_dict[f'site_{symmetry_value}']['equivalent_sites_list'].append(index_counter)
-                elif elements[element_type]['symbol'] not in abs_elements_list:
+                elif type_mapping_dict[str(element_type)].symbol not in abs_elements_list:
                     pass
                 else:
                     equivalency_dict[f'site_{symmetry_value}'] = {
-                        'symbol': elements[element_type]['symbol'],
+                        'kind_name': type_mapping_dict[str(element_type)].name,
+                        'symbol': type_mapping_dict[str(element_type)].symbol,
                         'site_index': symmetry_value,
                         'equivalent_sites_list': [symmetry_value]
                     }
@@ -243,7 +314,8 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                     equivalency_dict[f'site_{symmetry_value}']['equivalent_sites_list'].append(index_counter)
                 else:
                     equivalency_dict[f'site_{symmetry_value}'] = {
-                        'symbol': elements[element_type]['symbol'],
+                        'kind_name': type_mapping_dict[str(element_type)].name,
+                        'symbol': type_mapping_dict[str(element_type)].symbol,
                         'site_index': symmetry_value,
                         'equivalent_sites_list': [symmetry_value]
                     }
@@ -274,9 +346,45 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
 
         ase_structure = standardized_structure_node.get_ase()
         ase_supercell = ase_structure * multiples
-        new_supercell = StructureData(ase=ase_supercell)
 
-        result['supercell'] = new_supercell
+        # if there are hubbard data to apply, we re-construct
+        # the supercell site-by-site to keep the correct ordering
+        if is_hubbard_structure:
+            blank_supercell = StructureData(ase=ase_supercell)
+            new_supercell = StructureData()
+            new_supercell.set_cell(blank_supercell.cell)
+            num_extensions = np.product(multiples)
+            supercell_types_order = []
+            # For each supercell extension, loop over each site.
+            # This way, the original pattern-ordering of sites is
+            # preserved.
+            for i in range(0, num_extensions):  # pylint: disable=unused-variable
+                for type_number in types_order:
+                    supercell_types_order.append(type_number)
+
+            for site, type_number in zip(blank_supercell.sites, supercell_types_order):
+                kind_present = type_mapping_dict[str(type_number)]
+                if kind_present.name not in [kind.name for kind in new_supercell.kinds]:
+                    new_supercell.append_kind(kind_present)
+                new_site = Site(kind_name=kind_present.name, position=site.position)
+                new_supercell.append_site(new_site)
+        else:  # otherwise, simply re-construct the supercell with ASE
+            new_supercell = StructureData(ase=ase_supercell)
+
+        if is_hubbard_structure:  # Scale up the hubbard parameters to match and return the HubbardStructureData
+            if multiples == [1, 1, 1]:  # If no new supercell was made, just re-apply the incoming parameters
+                new_hubbard_supercell = HubbardStructureData.from_structure(new_supercell)
+                new_hubbard_supercell.hubbard = structure.hubbard
+                new_supercell = new_hubbard_supercell
+                supercell_hubbard_params = new_supercell.hubbard
+            else:
+                hubbard_manip = HubbardUtils(structure)
+                new_hubbard_supercell = hubbard_manip.get_hubbard_for_supercell(new_supercell)
+                new_supercell = new_hubbard_supercell
+                supercell_hubbard_params = new_supercell.hubbard
+            result['supercell'] = new_supercell
+        else:
+            result['supercell'] = new_supercell
         output_params['supercell_factors'] = multiples
         output_params['supercell_num_sites'] = len(new_supercell.sites)
         output_params['supercell_cell_matrix'] = new_supercell.cell
@@ -290,7 +398,12 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
 
         for index, site in enumerate(new_supercell.sites):
             if index == target_site:
-                absorbing_kind = Kind(name=abs_atom_marker, symbols=site.kind_name)
+                kind_name_at_position = site.kind_name
+                for kind in new_supercell.kinds:
+                    if kind_name_at_position == kind.name:
+                        kind_at_position = kind
+                        break
+                absorbing_kind = Kind(name=abs_atom_marker, symbols=kind_at_position.symbol)
                 absorbing_site = Site(kind_name=absorbing_kind.name, position=site.position)
                 marked_structure.append_kind(absorbing_kind)
                 marked_structure.append_site(absorbing_site)
@@ -300,7 +413,12 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                 new_site = Site(kind_name=site.kind_name, position=site.position)
                 marked_structure.append_site(new_site)
 
-        result[f'site_{target_site}_{value["symbol"]}'] = marked_structure
+        if is_hubbard_structure:
+            marked_hubbard_structure = HubbardStructureData.from_structure(marked_structure)
+            marked_hubbard_structure.hubbard = supercell_hubbard_params
+            result[f'site_{target_site}_{value["symbol"]}'] = marked_hubbard_structure
+        else:
+            result[f'site_{target_site}_{value["symbol"]}'] = marked_structure
 
     output_params['is_molecule_input'] = is_molecule_input
     result['output_parameters'] = orm.Dict(dict=output_params)

src/aiida_quantumespresso/workflows/xspectra/core.py

@@ -10,7 +10,7 @@
 from aiida.common import AttributeDict
 from aiida.engine import ToContext, WorkChain, append_, if_
 from aiida.orm.nodes.data.base import to_aiida_type
-from aiida.plugins import CalculationFactory, WorkflowFactory
+from aiida.plugins import CalculationFactory, DataFactory, WorkflowFactory
 import yaml
 
 from aiida_quantumespresso.calculations.functions.xspectra.get_powder_spectrum import get_powder_spectrum
@@ -21,6 +21,7 @@
 PwCalculation = CalculationFactory('quantumespresso.pw')
 PwBaseWorkChain = WorkflowFactory('quantumespresso.pw.base')
 XspectraBaseWorkChain = WorkflowFactory('quantumespresso.xspectra.base')
+HubbardStructureData = DataFactory('quantumespresso.hubbard_structure')
 
 
 class XspectraCoreWorkChain(ProtocolMixin, WorkChain):
@@ -101,7 +102,7 @@ def define(cls, spec):
         spec.inputs.validator = cls.validate_inputs
         spec.input(
             'structure',
-            valid_type=orm.StructureData,
+            valid_type=(orm.StructureData, HubbardStructureData),
             help=(
                 'Structure to be used for calculation, with at least one site containing the `abs_atom_marker` '
                 'as the kind label.'
@@ -294,6 +295,7 @@ def get_builder_from_protocol(
         pw_code,
         xs_code,
         structure,
+        initial_magnetic_moments=None,
         upf2plotcore_code=None,
         core_wfc_data=None,
         core_hole_pseudos=None,
@@ -334,6 +336,7 @@ def get_builder_from_protocol(
         :return: a process builder instance with all inputs defined ready for launch.
         """
 
+        from aiida_quantumespresso.common.types import SpinType
         inputs = cls.get_protocol_inputs(protocol, overrides)
         pw_inputs = PwBaseWorkChain.get_protocol_inputs(protocol=protocol, overrides=inputs.get('scf', {}))
         pw_params = pw_inputs['pw']['parameters']
@@ -348,9 +351,17 @@ def get_builder_from_protocol(
         )
 
         pw_inputs['pw']['parameters'] = pw_params
-
         pw_args = (pw_code, structure, protocol)
-        scf = PwBaseWorkChain.get_builder_from_protocol(*pw_args, overrides=pw_inputs, options=options, **kwargs)
+
+        if initial_magnetic_moments:
+            spin_type = SpinType.COLLINEAR
+            pw_kwargs = {'initial_magnetic_moments': initial_magnetic_moments, 'spin_type': spin_type}
+        else:
+            pw_kwargs = {}
+
+        scf = PwBaseWorkChain.get_builder_from_protocol(
+            *pw_args, overrides=pw_inputs, options=options, **pw_kwargs, **kwargs
+        )
 
         scf.pop('clean_workdir', None)
         scf['pw'].pop('structure', None)
@@ -368,12 +379,16 @@ def get_builder_from_protocol(
         abs_atom_marker = inputs['abs_atom_marker']
         xs_prod_parameters['INPUT_XSPECTRA']['xiabs'] = kinds_present.index(abs_atom_marker) + 1
         if core_hole_pseudos:
+            abs_element_kinds = []
             for kind in structure.kinds:
                 if kind.name == abs_atom_marker:
                     abs_element = kind.symbol
-
+            for kind in structure.kinds:  # run a second pass to check for multiple kinds of the same absorbing element
+                if kind.symbol == abs_element and kind.name != abs_atom_marker:
+                    abs_element_kinds.append(kind.name)
             builder.scf.pw.pseudos[abs_atom_marker] = core_hole_pseudos[abs_atom_marker]
-            builder.scf.pw.pseudos[abs_element] = core_hole_pseudos[abs_element]
+            for kind_name in abs_element_kinds:
+                builder.scf.pw.pseudos[kind_name] = core_hole_pseudos[abs_element]
 
         builder.xs_prod.xspectra.code = xs_code
         builder.xs_prod.xspectra.parameters = orm.Dict(xs_prod_parameters)