From 4b1bae1f2f998e0385bd91b95f91d838fabd2a41 Mon Sep 17 00:00:00 2001
From: Peter N O Gillespie <peter.no.gillespie@gmail.com>
Date: Wed, 30 Aug 2023 12:50:02 +0000
Subject: [PATCH 01/11] Enable Hubbard and Magnetic Moments for
 `XspectraCoreWorkChain` and `get_xspectra_structures`

This commit enables `XspectraCoreWorkChain` and
`get_xspectra_structures` to correctly parse `HubbardStructureData`
and magnetic moments.

Currently enabled:
* `get_xspectra_structures` will preserve Hubbard data if the input
  structure is `HubbardStructureData` node type and scale parameters to
match those required for the final supercell.
* the `get_builder_from_protocol` function of `XspectraCoreWorkChain`
 will accept `initial_magnetic_moments` as an optional input and pass
this information to the `scf` namespace.

To do:
* Enable `get_xspectra_structures` to process incoming magnetic moments
  data and return correct magnetic moments for the structures returned
by the function.
* Update the `run_upf2plotcore` step of `XspectraCoreWorkChain` to
  reliably pick the correct pseudo for the ShellJob step when more than
one Kind of the absorbing element is present.
* Implement steps in the `XspectraCrystalWorkChain` to correctly parse
  both HubbardStructureData and magnetic moments.
---
 .../functions/get_xspectra_structures.py      | 107 +++++++++++++++---
 .../workflows/xspectra/core.py                |  25 +++-
 2 files changed, 113 insertions(+), 19 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 407e62916..6bba81483 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -12,9 +12,11 @@
 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
+def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):  # pylint: disable=too-many-statements
     """Read a StructureData object using spglib for its symmetry data and return structures for XSpectra calculations.
 
     Takes an incoming StructureData node and prepares structures suitable for calculation with
@@ -52,6 +54,8 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                              systems (tolerance, eigen_tolerance, matrix_tolerance).
 
     :param structure: the StructureData object to be analysed
+    :param initial_magnetic_moments: an optional Dict node containing the magnetic moment for
+                                     each Kind in the structure
     :returns: StructureData objects for the standardized crystal structure, the supercell, and
               all generated structure and associated symmetry data
     """
@@ -100,7 +104,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 +126,22 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     else:
         spglib_kwargs = {}
 
+    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
+
+    if initial_magnetic_moments and standardize_structure:
+        raise ValidationError(
+            'Incoming structure set to be standardized, but magnetic moments data has been found. '
+            'Please set ``standardize_structure`` to false in ``**kwargs`` to preserve the magnetic structure.'
+        )
+
     output_params = {}
     result = {}
 
@@ -203,19 +227,37 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     # Process a periodic system
     else:
         incoming_structure_tuple = structure_to_spglib_tuple(structure)
+        spglib_tuple = incoming_structure_tuple[0]
+        types_order = spglib_tuple[-1]
+        kinds_information = incoming_structure_tuple[1]
+        kinds_list = incoming_structure_tuple[2]
 
-        symmetry_dataset = spglib.get_symmetry_dataset(incoming_structure_tuple[0], **spglib_kwargs)
+        # 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
+
+        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)
@@ -230,11 +272,12 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
             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 +286,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 +318,36 @@ 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 is magnetic domain or hubbard data to apply, we re-construct
+        # the supercell to keep the correct ordering
+        if is_hubbard_structure or initial_magnetic_moments is not None:
+            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 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:  # If there is no special information, we simply re-construct the supercell
+            new_supercell = StructureData(ase=ase_supercell)
+
+        if is_hubbard_structure:  # Scale up the hubbard parameters to match and return the HubbardStructureData
+            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 +361,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 +376,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)
diff --git a/src/aiida_quantumespresso/workflows/xspectra/core.py b/src/aiida_quantumespresso/workflows/xspectra/core.py
index 276f2e7f4..6cd76dc7d 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/core.py
+++ b/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,15 @@ 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}
+
+        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 +377,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)

From 1d17eaa30815ac0b16e0dee580eacbf7593e722c Mon Sep 17 00:00:00 2001
From: Peter N O Gillespie <peter.no.gillespie@gmail.com>
Date: Wed, 30 Aug 2023 14:21:12 +0000
Subject: [PATCH 02/11] Tests: Fix minor errors in tests from previous commit

---
 .../workflows/functions/get_xspectra_structures.py              | 2 +-
 src/aiida_quantumespresso/workflows/xspectra/core.py            | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 6bba81483..4f8041134 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -272,7 +272,7 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):
             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 type_mapping_dict[str(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}'] = {
diff --git a/src/aiida_quantumespresso/workflows/xspectra/core.py b/src/aiida_quantumespresso/workflows/xspectra/core.py
index 6cd76dc7d..2086d2de4 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/core.py
+++ b/src/aiida_quantumespresso/workflows/xspectra/core.py
@@ -356,6 +356,8 @@ def get_builder_from_protocol(
         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

From 7677a60da45d30f56f015987fb744d10bada718b Mon Sep 17 00:00:00 2001
From: "Peter N. O. Gillespie" <peter.no.gillespie@gmail.com>
Date: Thu, 7 Sep 2023 14:14:57 +0200
Subject: [PATCH 03/11] `get_xspectra_structures()`: Fix handling of marked
 structures

Updates `get_xspectra_structures()` to retrieve and re-apply Kind
data to output structures so that Kind names are kept. The function will
now also correctly determine the list of inequivalent atom sites if
`standardize_structure = False` and return the correct set of marked
structures.
---
 .../functions/get_xspectra_structures.py      | 29 +++++++++++++++----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 4f8041134..f7d525689 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -264,7 +264,20 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **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
@@ -341,10 +354,16 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):
             new_supercell = StructureData(ase=ase_supercell)
 
         if is_hubbard_structure:  # Scale up the hubbard parameters to match and return the HubbardStructureData
-            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
+            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

From d685dcde07b731905604074eb6855395f39059d6 Mon Sep 17 00:00:00 2001
From: "Peter N. O. Gillespie" <peter.no.gillespie@gmail.com>
Date: Thu, 14 Sep 2023 16:03:20 +0200
Subject: [PATCH 04/11] `get_xspectra_structures()`: remove options for
 `initial_magnetization`

Removes options for `initial_magnetization` from the function, since
this has been shown to be superfluous in practice.
---
 .../functions/get_xspectra_structures.py      | 21 +++++++------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index f7d525689..3f63c6773 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -16,7 +16,7 @@
 
 
 @calcfunction
-def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):  # pylint: disable=too-many-statements
+def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-statements
     """Read a StructureData object using spglib for its symmetry data and return structures for XSpectra calculations.
 
     Takes an incoming StructureData node and prepares structures suitable for calculation with
@@ -54,8 +54,6 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):
                              systems (tolerance, eigen_tolerance, matrix_tolerance).
 
     :param structure: the StructureData object to be analysed
-    :param initial_magnetic_moments: an optional Dict node containing the magnetic moment for
-                                     each Kind in the structure
     :returns: StructureData objects for the standardized crystal structure, the supercell, and
               all generated structure and associated symmetry data
     """
@@ -136,12 +134,6 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):
     else:
         is_hubbard_structure = False
 
-    if initial_magnetic_moments and standardize_structure:
-        raise ValidationError(
-            'Incoming structure set to be standardized, but magnetic moments data has been found. '
-            'Please set ``standardize_structure`` to false in ``**kwargs`` to preserve the magnetic structure.'
-        )
-
     output_params = {}
     result = {}
 
@@ -332,14 +324,17 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):
         ase_structure = standardized_structure_node.get_ase()
         ase_supercell = ase_structure * multiples
 
-        # if there is magnetic domain or hubbard data to apply, we re-construct
-        # the supercell to keep the correct ordering
-        if is_hubbard_structure or initial_magnetic_moments is not None:
+        # 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)
@@ -350,7 +345,7 @@ def get_xspectra_structures(structure, initial_magnetic_moments=None, **kwargs):
                     new_supercell.append_kind(kind_present)
                 new_site = Site(kind_name=kind_present.name, position=site.position)
                 new_supercell.append_site(new_site)
-        else:  # If there is no special information, we simply re-construct the supercell
+        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

From 7f82c72f3301c8eec6f91810d6b906d10cc97d90 Mon Sep 17 00:00:00 2001
From: Peter N O Gillespie <peter.no.gillespie@gmail.com>
Date: Thu, 21 Sep 2023 08:57:08 +0000
Subject: [PATCH 05/11] `XspectraCrystalWorkChain`: Enable Hubbard and Magnetic
 Data

This commit fully enables the `XspectraCrystalWorkChain` to work with
`HubbardStructureData` and to apply `initial_magnetic_moments` to all
sub-processes. Automatically sets `standardize_structure` to `False` in
the `get_xspectra_structures` step of the `WorkChain` if the input
structure is `HubbardStructureData` as required by the `CalcFunction`.

Also updates `get_xspectra_structures` with a new optional parameter
`use_element_types` which instructs the `CalcFunction` to treat all
`Kinds` of the same element as the same for the purposes of symmetry
analysis even if the `kind_name`s are different. Defaults to `False`.
---
 .../functions/get_xspectra_structures.py      | 25 ++++++++++++++++++-
 .../workflows/xspectra/crystal.py             | 15 +++++++++--
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 3f63c6773..219e2b5c6 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -47,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
@@ -124,6 +129,11 @@ 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:
@@ -236,7 +246,20 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                 if value == kind.name:
                     type_mapping_dict[key] = kind
 
-        symmetry_dataset = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)
+        # 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
diff --git a/src/aiida_quantumespresso/workflows/xspectra/crystal.py b/src/aiida_quantumespresso/workflows/xspectra/crystal.py
index d2b851af4..2fe76495f 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/crystal.py
+++ b/src/aiida_quantumespresso/workflows/xspectra/crystal.py
@@ -11,6 +11,7 @@
 from aiida_pseudo.data.pseudo import UpfData as aiida_pseudo_upf
 
 from aiida_quantumespresso.calculations.functions.xspectra.get_spectra_by_element import get_spectra_by_element
+from aiida_quantumespresso.utils.hubbard import HubbardStructureData
 from aiida_quantumespresso.utils.mapping import prepare_process_inputs
 from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin, recursive_merge
 
@@ -481,6 +482,10 @@ def get_xspectra_structures(self):
             for key, node in optional_cell_prep.items():
                 inputs[key] = node
 
+        if isinstance(self.inputs.structure, HubbardStructureData):
+            # This must be False in the case of HubbardStructureData, otherwise get_xspectra_structures will except
+            inputs['standardize_structure'] = orm.Bool(False)
+
         if 'spglib_settings' in self.inputs:
             inputs['spglib_settings'] = self.inputs.spglib_settings
 
@@ -527,8 +532,10 @@ def run_upf2plotcore(self):
 
             shell_inputs['code'] = self.inputs.upf2plotcore_code
             shell_inputs['nodes'] = {'upf': upf}
-            shell_inputs['arguments'] = ['upf']
-            shell_inputs['metadata'] = {'call_link_label': f'upf2plotcore_{element}'}
+            shell_inputs['metadata'] = {
+                'call_link_label': f'upf2plotcore_{element}',
+                'options' : {'filename_stdin' : upf.filename}
+                }
 
             future_shelljob = self.submit(ShellJob, **shell_inputs)
             self.report(f'Launching upf2plotcore.sh for {element}<{future_shelljob.pk}>')
@@ -566,6 +573,7 @@ def run_all_xspectra_core(self):
             structure = structures_to_process[site]
             inputs.structure = structure
             abs_element = equivalent_sites_data[site]['symbol']
+            abs_atom_kind = equivalent_sites_data[site]['kind_name']
 
             if 'core_hole_treatments' in self.inputs:
                 ch_treatments = self.inputs.core_hole_treatments.get_dict()
@@ -599,6 +607,9 @@ def run_all_xspectra_core(self):
             # to avoid the need for fudges like these.
             if ch_treatment == 'xch_smear':
                 new_scf_params['SYSTEM'][f'starting_magnetization({abs_species_index})'] = 1
+            elif 'starting_magnetization' in new_scf_params['SYSTEM']:
+                inherited_mag =  new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_kind]
+                new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_marker] = inherited_mag
 
             core_hole_pseudo = self.inputs.core_hole_pseudos[abs_element]
             inputs.scf.pw.pseudos[abs_atom_marker] = core_hole_pseudo

From 1d0782ca72ebe5b6d79011f00aaec4fed8603e72 Mon Sep 17 00:00:00 2001
From: "Peter N. O. Gillespie" <peter.no.gillespie@gmail.com>
Date: Fri, 22 Sep 2023 12:29:10 +0200
Subject: [PATCH 06/11] `XspectraCrystalWorkChain`: Fix handling of GIPAW
 pseudos

Updates the `run_all_xspectra_core` step in the `XspectraCrystalWorkChain`
to correctly assign the GIPAW pseudo provided for the element to all
`Kind`s corresponding to the element, retaining the step to remove the
GIPAW pseudo if there is only one atom of the absorbing element present
in the structure (and thus, to avoid a crash due to incorrect pseudo
allocation).
---
 .../workflows/xspectra/crystal.py             | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/xspectra/crystal.py b/src/aiida_quantumespresso/workflows/xspectra/crystal.py
index 2fe76495f..882dee4ae 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/crystal.py
+++ b/src/aiida_quantumespresso/workflows/xspectra/crystal.py
@@ -560,7 +560,7 @@ def inspect_upf2plotcore(self):
             if num_core_states == 0:
                 return self.exit_codes.ERROR_NO_GIPAW_INFO_FOUND
 
-    def run_all_xspectra_core(self):
+    def run_all_xspectra_core(self): # pylint: disable=too-many-statements
         """Call all XspectraCoreWorkChains required to compute all requested spectra."""
 
         structures_to_process = self.ctx.structures_to_process
@@ -612,13 +612,18 @@ def run_all_xspectra_core(self):
                 new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_marker] = inherited_mag
 
             core_hole_pseudo = self.inputs.core_hole_pseudos[abs_element]
+            gipaw_pseudo = self.inputs.gipaw_pseudos[abs_element]
             inputs.scf.pw.pseudos[abs_atom_marker] = core_hole_pseudo
-            # In the case where the absorbing atom is the only one of its element in the
-            # structure, we avoid setting the GIPAW pseudo for it and remove the one .
-            if abs_element in kinds_present:
-                gipaw_pseudo = self.inputs.gipaw_pseudos[abs_element]
-                inputs.scf.pw.pseudos[abs_element] = gipaw_pseudo
-            else:
+            # Check how many instances of the absorbing element are present and assign each the GIPAW
+            # pseudo if they are not the absorbing atom itself.
+            abs_element_kinds = []
+            for kind in structure.kinds:
+                if kind.symbol == abs_element and kind.name != abs_atom_marker:
+                    abs_element_kinds.append(kind.name)
+            if len(abs_element_kinds) > 0:
+                for kind_name in abs_element_kinds:
+                    scf_inputs['pseudos'][kind_name] = gipaw_pseudo
+            else: # if there is only one atom of the absorbing element, pop the GIPAW pseudo to avoid a crash
                 scf_inputs['pseudos'].pop(abs_element, None)
 
             scf_inputs.parameters = orm.Dict(new_scf_params)
@@ -631,7 +636,7 @@ def run_all_xspectra_core(self):
             xspectra_core_workchains[site] = future
             self.report(f'launched XspectraCoreWorkChain for {site}<{future.pk}>')
 
-        return ToContext(**xspectra_core_workchains)
+        return ToContext(**xspectra_core_workchains) # pylint: enable=too-many-statements
 
     def inspect_all_xspectra_core(self):
         """Check that all the XspectraCoreWorkChain sub-processes finished sucessfully."""

From 2448f63630fdcf1b7c3c2431246811017539921c Mon Sep 17 00:00:00 2001
From: "Peter N. O. Gillespie" <peter.no.gillespie@gmail.com>
Date: Wed, 29 Nov 2023 11:26:13 +0100
Subject: [PATCH 07/11] XSpectra `WorkChain`s: Refinements and Improvements

Addresses requested changes in PR #969 and refines the logic for
assigning magnetic states to the absorbing atom depending on chosen
core-hole treatment type, as some oversights were made in the case of
magnetic systems. The absorbing atom is now always given a spin of 1
if it was 0 in the neutral system and we are using an XCH-type treatment.
Otherwise, the spin is set to the value of its inherited `Kind`.

Other refinements:
* Fixed a typo where the core-hole treatments would override the
overrides dictionary (the opposite to intended behaviour).
* Removed unnecessary options for SpinType and initial_magnetic_moments
in the `CoreWorkChain`.
* Added tests for `get_xspectra_structures` to cover basic behaviour and
correct handling of HubbardStructureData.
---
 .../functions/get_xspectra_structures.py      | 17 +++-
 .../workflows/xspectra/core.py                | 12 +--
 .../workflows/xspectra/crystal.py             | 46 ++++++---
 .../functions/test_get_xspectra_structures.py | 96 +++++++++++++++++++
 4 files changed, 143 insertions(+), 28 deletions(-)
 create mode 100644 tests/workflows/functions/test_get_xspectra_structures.py

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 219e2b5c6..52103a766 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -51,7 +51,8 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                              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.
+                             containing silicon. Defaults to False otherwise. Only meaningful in
+                             conjunction with ``standardize_structure = False``.
         - 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
@@ -134,7 +135,7 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     else:
         use_element_types = False
 
-    if structure.node_type == 'data.quantumespresso.hubbard_structure.HubbardStructureData.':
+    if isinstance(structure, HubbardStructureData):
         is_hubbard_structure = True
         if standardize_structure:
             raise ValidationError(
@@ -246,9 +247,15 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                 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.
+        # By default, `structure_to_spglib_tuple` gives different
+        # ``Kinds`` of the same element a distinct atomic number by
+        # multiplying the normal atomic number by 1000, then adding
+        # 100 for each distinct duplicate. if we want to treat all sites
+        # of the same element as equal, then we must therefore briefly
+        # operate on a "cleaned" version of the structure tuple where this
+        # new label is reduced to its normal element number. This relies on
+        # the fact that int(i) will essentially ignore the value of the decimal
+        # when converting floats to integers.
         if use_element_types:
             clean_structure_tuple = (spglib_tuple[0], spglib_tuple[1], [])
             for i in spglib_tuple[2]:
diff --git a/src/aiida_quantumespresso/workflows/xspectra/core.py b/src/aiida_quantumespresso/workflows/xspectra/core.py
index 2086d2de4..103108485 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/core.py
+++ b/src/aiida_quantumespresso/workflows/xspectra/core.py
@@ -295,7 +295,6 @@ 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,
@@ -336,7 +335,6 @@ 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']
@@ -353,15 +351,7 @@ def get_builder_from_protocol(
         pw_inputs['pw']['parameters'] = pw_params
         pw_args = (pw_code, structure, protocol)
 
-        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 = PwBaseWorkChain.get_builder_from_protocol(*pw_args, overrides=pw_inputs, options=options, **kwargs)
 
         scf.pop('clean_workdir', None)
         scf['pw'].pop('structure', None)
diff --git a/src/aiida_quantumespresso/workflows/xspectra/crystal.py b/src/aiida_quantumespresso/workflows/xspectra/crystal.py
index 882dee4ae..32737dbec 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/crystal.py
+++ b/src/aiida_quantumespresso/workflows/xspectra/crystal.py
@@ -534,8 +534,14 @@ def run_upf2plotcore(self):
             shell_inputs['nodes'] = {'upf': upf}
             shell_inputs['metadata'] = {
                 'call_link_label': f'upf2plotcore_{element}',
-                'options' : {'filename_stdin' : upf.filename}
+                'options' : {
+                    'filename_stdin' : upf.filename,
+                    'resources' : {
+                        'num_machines' : 1,
+                        'num_mpiprocs_per_machine' : 1
+                    }
                 }
+            }
 
             future_shelljob = self.submit(ShellJob, **shell_inputs)
             self.report(f'Launching upf2plotcore.sh for {element}<{future_shelljob.pk}>')
@@ -593,7 +599,7 @@ def run_all_xspectra_core(self): # pylint: disable=too-many-statements
             scf_inputs = inputs.scf.pw
             scf_params = scf_inputs.parameters.get_dict()
             ch_inputs = XspectraCoreWorkChain.get_treatment_inputs(treatment=ch_treatment)
-            new_scf_params = recursive_merge(left=scf_params, right=ch_inputs)
+            new_scf_params = recursive_merge(left=ch_inputs, right=scf_params)
 
             # Set the absorbing species index (`xiabs`) for the xspectra.x input.
             new_xs_params = inputs.xs_prod.xspectra.parameters.get_dict()
@@ -601,21 +607,37 @@ def run_all_xspectra_core(self): # pylint: disable=too-many-statements
             abs_species_index = kinds_present.index(abs_atom_marker) + 1
             new_xs_params['INPUT_XSPECTRA']['xiabs'] = abs_species_index
 
-            # Set `starting_magnetization` if we are using an XCH approximation, using the
-            # absorbing species as a reasonable place for the unpaired electron.
-            # As a future note, we need to re-visit the core-hole treatment settings, in order
-            # to avoid the need for fudges like these.
-            if ch_treatment == 'xch_smear':
-                new_scf_params['SYSTEM'][f'starting_magnetization({abs_species_index})'] = 1
-            elif 'starting_magnetization' in new_scf_params['SYSTEM']:
+            # Set `starting_magnetization` if we are using an XCH approximation, using
+            # the absorbing species as a reasonable place for the unpaired electron.
+            # Alternatively, ensure the starting magnetic moment is a reasonable guess
+            # given the input parameters. (e.g. it conforms to an existing magnetic
+            # structure already defined for the system)
+
+            # TODO: we need to re-visit the core-hole treatment settings,
+            # in order to avoid the need for fudges like these and set these at
+            # submission rather than inside the WorkChain itself.
+            if 'starting_magnetization' in new_scf_params['SYSTEM']:
                 inherited_mag =  new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_kind]
-                new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_marker] = inherited_mag
+                if ch_treatment not in ['xch_smear', 'xch_fixed']:
+                    new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_marker] = inherited_mag
+                else: # if there is meant to be an unpaired electron, give it to the absorbing atom.
+                    if inherited_mag == 0: # set it to 1, if it would be neutral in the ground-state.
+                        new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_marker] =  1
+                    else: # assume that it takes the same magnetic configuration as the kind that it replaces.
+                        new_scf_params['SYSTEM']['starting_magnetization'][abs_atom_marker] =  inherited_mag
+            elif ch_treatment in ['xch_smear', 'xch_fixed']:
+                new_scf_params['SYSTEM']['starting_magnetization'] = {abs_atom_marker : 1}
+
+            # remove any duplicates created from the "core_hole_treatments.yaml" defaults
+            for key in new_scf_params['SYSTEM'].keys():
+                if 'starting_magnetization(' in key:
+                    new_scf_params['SYSTEM'].pop(key, None)
 
             core_hole_pseudo = self.inputs.core_hole_pseudos[abs_element]
             gipaw_pseudo = self.inputs.gipaw_pseudos[abs_element]
             inputs.scf.pw.pseudos[abs_atom_marker] = core_hole_pseudo
-            # Check how many instances of the absorbing element are present and assign each the GIPAW
-            # pseudo if they are not the absorbing atom itself.
+            # Check how many instances of the absorbing element are present and assign
+            # each the GIPAW pseudo if they are not the absorbing atom itself.
             abs_element_kinds = []
             for kind in structure.kinds:
                 if kind.symbol == abs_element and kind.name != abs_atom_marker:
diff --git a/tests/workflows/functions/test_get_xspectra_structures.py b/tests/workflows/functions/test_get_xspectra_structures.py
new file mode 100644
index 000000000..365ce542f
--- /dev/null
+++ b/tests/workflows/functions/test_get_xspectra_structures.py
@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""Tests for the `get_marked_structure` class."""
+from aiida.orm import Bool, Dict
+import pytest
+
+from aiida_quantumespresso.utils.hubbard import HubbardStructureData, HubbardUtils
+from aiida_quantumespresso.workflows.functions.get_xspectra_structures import get_xspectra_structures
+
+
+@pytest.fixture
+def generate_hubbard():
+    """Return a `Hubbard` instance."""
+
+    def _generate_hubbard():
+        from aiida_quantumespresso.common.hubbard import Hubbard
+        return Hubbard.from_list([(0, '1s', 0, '1s', 5.0, (0, 0, 0), 'Ueff')])
+
+    return _generate_hubbard
+
+
+@pytest.fixture
+def generate_hubbard_structure(generate_structure):
+    """Return a `HubbardStructureData` instance."""
+
+    def _generate_hubbard_structure():
+        from aiida_quantumespresso.common.hubbard import Hubbard
+        structure = generate_structure('silicon-kinds')
+        hp_list = [(0, '1s', 0, '1s', 5.0, (0, 0, 0), 'Ueff')]
+        hubbard = Hubbard.from_list(hp_list)
+        return HubbardStructureData.from_structure(structure=structure, hubbard=hubbard)
+
+    return _generate_hubbard_structure
+
+
+def test_base(generate_structure):
+    """Test the basic operation of get_xspectra_structures."""
+
+    c_si = generate_structure('silicon')
+    spglib_options = Dict({'symprec': 1.0e-3})
+    inputs = {'structure': c_si, 'spglib_options': spglib_options}
+    result = get_xspectra_structures(**inputs)
+    assert len(result) == 4
+    out_params = result['output_parameters'].get_dict()
+    assert out_params['spacegroup_number'] == 227
+    assert out_params['supercell_num_sites'] == 64
+    assert len(out_params['equivalent_sites_data']) == 1
+
+
+def test_use_element_types(generate_structure):
+    """Test the CF's `use_element_types` flag."""
+
+    c_si = generate_structure('silicon')
+    c_si_kinds = generate_structure('silicon-kinds')
+    spglib_options = Dict({'symprec': 1.0e-3})
+    inputs_bare = {'structure': c_si, 'spglib_options': spglib_options}
+    inputs_kinds = {'structure': c_si_kinds, 'spglib_options': spglib_options, 'standardize_structure': Bool(False)}
+
+    result_bare = get_xspectra_structures(**inputs_bare)
+    result_kinds = get_xspectra_structures(**inputs_kinds)
+
+    inputs_element_types = {
+        'structure': c_si_kinds,
+        'spglib_options': spglib_options,
+        'standardize_structure': Bool(False),
+        'use_element_types': Bool(True),
+    }
+    result_element_types = get_xspectra_structures(**inputs_element_types)
+
+    assert 'site_1_Si' in result_kinds
+    assert 'site_1_Si' not in result_element_types
+    assert 'site_1_Si' not in result_bare
+
+
+def test_hubbard(generate_structure):
+    """Test that the CalcFunction will pass Hubbard parameters to the output structures.
+
+    The intent here is to confirm that simply using the `initialize_` methods to get
+    hubbard parameters will propogate to the resulting supercells and (crucially)
+    generate the correct hubbard card.
+    """
+
+    c_si = generate_structure()
+    c_si_hub = HubbardStructureData.from_structure(c_si)
+    c_si_hub.initialize_onsites_hubbard('Si', '1s', 0.0, 'Ueff', False)
+
+    inputs = {'structure': c_si_hub, 'standardize_structure': Bool(False)}
+    result = get_xspectra_structures(**inputs)
+
+    marked = result['site_0_Si']
+    utils_marked = HubbardUtils(marked)
+    hub_card_lines = [i.strip() for i in utils_marked.get_hubbard_card().splitlines()]
+    out_params = result['output_parameters'].get_dict()
+
+    assert out_params['supercell_num_sites'] == 54
+    assert 'U\tX-1s\t0.0' in hub_card_lines
+    assert 'U\tSi-1s\t0.0' in hub_card_lines

From c4701b3c4ba40978a9f40725cab5ca89b94ed74d Mon Sep 17 00:00:00 2001
From: "Peter N. O. Gillespie" <peter.no.gillespie@gmail.com>
Date: Mon, 4 Mar 2024 10:16:28 +0100
Subject: [PATCH 08/11] `get_xspectra_structures`: Bugfixes and Improvements

Changes:
* Fixes a bug where, if `use_element_types = True`, the symmetry data of
the "non-cleaned" structure would be erroneously returned in
`output_parameters`
* Fixes a bug where, if `use_element_types = True`, the `kind_name` for
each entry in `equivalent_sites_data` would be reported incorrectly.
* Changes the default setting of `use_element_types` to `True`.
* Changes the logic for converting the atomic number of each `type` in
`spglib_tuple` used to generate the `cleaned_structure_tuple` to use
`np.trunc(int()) instead of just `int()`.
---
 .../functions/get_xspectra_structures.py      | 45 ++++++++++---------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 52103a766..4ef20b906 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -51,8 +51,7 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                              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. Only meaningful in
-                             conjunction with ``standardize_structure = False``.
+                             containing silicon. Defaults to True.
         - 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
@@ -133,7 +132,7 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     if 'use_element_types' in unwrapped_kwargs.keys():
         use_element_types = unwrapped_kwargs['use_element_types'].value
     else:
-        use_element_types = False
+        use_element_types = True
 
     if isinstance(structure, HubbardStructureData):
         is_hubbard_structure = True
@@ -253,20 +252,18 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
         # 100 for each distinct duplicate. if we want to treat all sites
         # of the same element as equal, then we must therefore briefly
         # operate on a "cleaned" version of the structure tuple where this
-        # new label is reduced to its normal element number. This relies on
-        # the fact that int(i) will essentially ignore the value of the decimal
-        # when converting floats to integers.
+        # new label is reduced to its normal element number.
         if use_element_types:
-            clean_structure_tuple = (spglib_tuple[0], spglib_tuple[1], [])
+            cleaned_structure_tuple = (spglib_tuple[0], spglib_tuple[1], [])
             for i in spglib_tuple[2]:
                 if i >= 1000:
-                    new_i = int(i / 1000)
+                    new_i = int(np.trunc(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)
+                cleaned_structure_tuple[2].append(new_i)
+            cleaned_symmetry_dataset = spglib.get_symmetry_dataset(cleaned_structure_tuple, **spglib_kwargs)
+
+        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
@@ -285,7 +282,11 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
             symmetry_dataset = spglib.get_symmetry_dataset(standardized_structure_tuple, **spglib_kwargs)
             structure_is_standardized = True
 
-        equivalent_atoms_array = symmetry_dataset['equivalent_atoms']
+        if use_element_types:
+            equivalent_atoms_array = cleaned_symmetry_dataset['equivalent_atoms']
+        else:
+            equivalent_atoms_array = symmetry_dataset['equivalent_atoms']
+
         if structure_is_standardized:
             element_types = symmetry_dataset['std_types']
         else:  # convert the 'std_types' from standardized to primitive cell
@@ -302,11 +303,11 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
             element_types = primitive_types
 
         equivalency_dict = {}
-        index_counter = 0
-        for symmetry_value, element_type in zip(equivalent_atoms_array, element_types):
+        for index, symmetry_types in enumerate(zip(equivalent_atoms_array, element_types)):
+            symmetry_value, element_type = symmetry_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)
+                    equivalency_dict[f'site_{symmetry_value}']['equivalent_sites_list'].append(index)
                 elif type_mapping_dict[str(element_type)].symbol not in abs_elements_list:
                     pass
                 else:
@@ -318,7 +319,7 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                     }
             else:  # process everything in the system
                 if f'site_{symmetry_value}' in equivalency_dict:
-                    equivalency_dict[f'site_{symmetry_value}']['equivalent_sites_list'].append(index_counter)
+                    equivalency_dict[f'site_{symmetry_value}']['equivalent_sites_list'].append(index)
                 else:
                     equivalency_dict[f'site_{symmetry_value}'] = {
                         'kind_name': type_mapping_dict[str(element_type)].name,
@@ -326,14 +327,18 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                         'site_index': symmetry_value,
                         'equivalent_sites_list': [symmetry_value]
                     }
-            index_counter += 1
 
         for value in equivalency_dict.values():
             value['multiplicity'] = len(value['equivalent_sites_list'])
 
         output_params['equivalent_sites_data'] = equivalency_dict
-        output_params['spacegroup_number'] = symmetry_dataset['number']
-        output_params['international_symbol'] = symmetry_dataset['international']
+
+        if use_element_types:
+            output_params['spacegroup_number'] = cleaned_symmetry_dataset['number']
+            output_params['international_symbol'] = cleaned_symmetry_dataset['international']
+        else:
+            output_params['spacegroup_number'] = symmetry_dataset['number']
+            output_params['international_symbol'] = symmetry_dataset['international']
 
         output_params['structure_is_standardized'] = structure_is_standardized
         if structure_is_standardized:

From 1f1d8f5309c54205b707c5c56d64c25c2dec12ec Mon Sep 17 00:00:00 2001
From: "Peter N. O. Gillespie" <peter.no.gillespie@gmail.com>
Date: Mon, 4 Mar 2024 11:32:52 +0100
Subject: [PATCH 09/11] Tests: Update `test_get_xspectra_structures`

Updates `test_use_element_types` following changes to default function
behaviour introduced in previous commit
(c4701b3c4ba40978a9f40725cab5ca89b94ed74d)
---
 tests/workflows/functions/test_get_xspectra_structures.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/workflows/functions/test_get_xspectra_structures.py b/tests/workflows/functions/test_get_xspectra_structures.py
index 365ce542f..e192bbc44 100644
--- a/tests/workflows/functions/test_get_xspectra_structures.py
+++ b/tests/workflows/functions/test_get_xspectra_structures.py
@@ -53,7 +53,12 @@ def test_use_element_types(generate_structure):
     c_si_kinds = generate_structure('silicon-kinds')
     spglib_options = Dict({'symprec': 1.0e-3})
     inputs_bare = {'structure': c_si, 'spglib_options': spglib_options}
-    inputs_kinds = {'structure': c_si_kinds, 'spglib_options': spglib_options, 'standardize_structure': Bool(False)}
+    inputs_kinds = {
+        'structure': c_si_kinds,
+        'use_element_types': Bool(False),
+        'spglib_options': spglib_options,
+        'standardize_structure': Bool(False)
+    }
 
     result_bare = get_xspectra_structures(**inputs_bare)
     result_kinds = get_xspectra_structures(**inputs_kinds)
@@ -62,7 +67,6 @@ def test_use_element_types(generate_structure):
         'structure': c_si_kinds,
         'spglib_options': spglib_options,
         'standardize_structure': Bool(False),
-        'use_element_types': Bool(True),
     }
     result_element_types = get_xspectra_structures(**inputs_element_types)
 

From 7f3842fbfdc6ce45486499a1611a02e842e4aebc Mon Sep 17 00:00:00 2001
From: Peter N O Gillespie <peter.no.gillespie@gmail.com>
Date: Thu, 21 Mar 2024 18:03:15 +0000
Subject: [PATCH 10/11] `get_xspectra_structures`: Consolidate Logic

Removes various unneeded `if` statements and re-uses information
generated during the process flow to simplify steps.
---
 .../functions/get_xspectra_structures.py      | 75 +++++++------------
 1 file changed, 26 insertions(+), 49 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 4ef20b906..591ae175d 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -96,7 +96,6 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     else:
         standardize_structure = True
     if 'absorbing_elements_list' in unwrapped_kwargs.keys():
-        elements_defined = True
         abs_elements_list = unwrapped_kwargs['absorbing_elements_list'].get_list()
         # confirm that the elements requested are actually in the input structure
         for req_element in abs_elements_list:
@@ -106,7 +105,6 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                     f' {elements_present}'
                 )
     else:
-        elements_defined = False
         abs_elements_list = []
         for kind in structure.kinds:
             if kind.symbol not in abs_elements_list:
@@ -261,9 +259,9 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
                 else:
                     new_i = i
                 cleaned_structure_tuple[2].append(new_i)
-            cleaned_symmetry_dataset = spglib.get_symmetry_dataset(cleaned_structure_tuple, **spglib_kwargs)
-
-        symmetry_dataset = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)
+            symmetry_dataset = spglib.get_symmetry_dataset(cleaned_structure_tuple, **spglib_kwargs)
+        else:
+            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
@@ -282,22 +280,24 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
             symmetry_dataset = spglib.get_symmetry_dataset(standardized_structure_tuple, **spglib_kwargs)
             structure_is_standardized = True
 
-        if use_element_types:
-            equivalent_atoms_array = cleaned_symmetry_dataset['equivalent_atoms']
-        else:
-            equivalent_atoms_array = symmetry_dataset['equivalent_atoms']
+        equivalent_atoms_array = symmetry_dataset['equivalent_atoms']
 
         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']
+            # we generate the type-specific data on-the-fly since we need to
+            # know which type (and thus kind) *should* be at each site
+            # even if we "cleaned" the structure previously
+            spglib_std_types = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)['std_types']
+            spglib_map_to_prim = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)['mapping_to_primitive']
+            spglib_std_map_to_prim = spglib.get_symmetry_dataset(spglib_tuple,
+                                                                 **spglib_kwargs)['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']:
+            for position in spglib_map_to_prim:
                 atom_type = map_std_pos_to_type[str(position)]
                 primitive_types.append(atom_type)
             element_types = primitive_types
@@ -305,19 +305,7 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
         equivalency_dict = {}
         for index, symmetry_types in enumerate(zip(equivalent_atoms_array, element_types)):
             symmetry_value, element_type = symmetry_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)
-                elif type_mapping_dict[str(element_type)].symbol not in abs_elements_list:
-                    pass
-                else:
-                    equivalency_dict[f'site_{symmetry_value}'] = {
-                        '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]
-                    }
-            else:  # process everything in the system
+            if type_mapping_dict[str(element_type)].symbol in abs_elements_list:
                 if f'site_{symmetry_value}' in equivalency_dict:
                     equivalency_dict[f'site_{symmetry_value}']['equivalent_sites_list'].append(index)
                 else:
@@ -333,12 +321,8 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
 
         output_params['equivalent_sites_data'] = equivalency_dict
 
-        if use_element_types:
-            output_params['spacegroup_number'] = cleaned_symmetry_dataset['number']
-            output_params['international_symbol'] = cleaned_symmetry_dataset['international']
-        else:
-            output_params['spacegroup_number'] = symmetry_dataset['number']
-            output_params['international_symbol'] = symmetry_dataset['international']
+        output_params['spacegroup_number'] = symmetry_dataset['number']
+        output_params['international_symbol'] = symmetry_dataset['international']
 
         output_params['structure_is_standardized'] = structure_is_standardized
         if structure_is_standardized:
@@ -384,16 +368,13 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
             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
+            # we can exploit the fact that `get_hubbard_for_supercell` will return a HubbardStructureData node
+            # with the same hubbard parameters in the case where the input structure and the supercell are the
+            # same (i.e. multiples == [1, 1, 1])
+            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
@@ -405,23 +386,19 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
     for value in equivalency_dict.values():
         target_site = value['site_index']
         marked_structure = StructureData()
-        supercell_kinds = {kind.name: kind for kind in new_supercell.kinds}
         marked_structure.set_cell(new_supercell.cell)
+        new_kind_names = [kind.name for kind in new_supercell.kinds]
 
         for index, site in enumerate(new_supercell.sites):
+            kind_at_position = new_supercell.kinds[new_kind_names.index(site.kind_name)]
             if index == target_site:
-                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)
             else:
-                if site.kind_name not in [kind.name for kind in marked_structure.kinds]:
-                    marked_structure.append_kind(supercell_kinds[site.kind_name])
+                if kind_at_position.name not in [kind.name for kind in marked_structure.kinds]:
+                    marked_structure.append_kind(kind_at_position)
                 new_site = Site(kind_name=site.kind_name, position=site.position)
                 marked_structure.append_site(new_site)
 

From d34f16f51bad5d8b101f1d247a94fc1e84b47d3d Mon Sep 17 00:00:00 2001
From: Peter N O Gillespie <peter.no.gillespie@gmail.com>
Date: Thu, 28 Mar 2024 15:05:35 +0000
Subject: [PATCH 11/11] `get_xspectra_structures`: Reduce Use of
 `get_symmetry_dataset`

Changes requested for PR #969

Replaces some uses of `spglib.get_symmetry_dataset` with a single
function call.
---
 .../workflows/functions/get_xspectra_structures.py        | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
index 591ae175d..7ae29d2ff 100644
--- a/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
+++ b/src/aiida_quantumespresso/workflows/functions/get_xspectra_structures.py
@@ -288,10 +288,10 @@ def get_xspectra_structures(structure, **kwargs):  # pylint: disable=too-many-st
             # we generate the type-specific data on-the-fly since we need to
             # know which type (and thus kind) *should* be at each site
             # even if we "cleaned" the structure previously
-            spglib_std_types = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)['std_types']
-            spglib_map_to_prim = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)['mapping_to_primitive']
-            spglib_std_map_to_prim = spglib.get_symmetry_dataset(spglib_tuple,
-                                                                 **spglib_kwargs)['std_mapping_to_primitive']
+            non_cleaned_dataset = spglib.get_symmetry_dataset(spglib_tuple, **spglib_kwargs)
+            spglib_std_types = non_cleaned_dataset['std_types']
+            spglib_map_to_prim = non_cleaned_dataset['mapping_to_primitive']
+            spglib_std_map_to_prim = non_cleaned_dataset['std_mapping_to_primitive']
 
             map_std_pos_to_type = {}
             for position, atom_type in zip(spglib_std_map_to_prim, spglib_std_types):