[XPS] support multiple core levels for the same element

src/aiida_qe_xspec/calculations/functions/xspectra/get_xps_spectra.py

@@ -3,109 +3,109 @@
 from aiida.engine import calcfunction
 import numpy as np
 
+def spectra_broadening(points, sigma=0.1, gamma=0.1):
+    """Broadening base on the binding energy.
 
-@calcfunction
-def get_spectra_by_element(elements_list, equivalent_sites_data, voight_gamma, voight_sigma, **kwargs):  # pylint: disable=too-many-statements
-    """Generate the XPS spectra for each element.
-
-    Calculate the core level shift and binding energy for each element.
-    Generate the final spectra using the Voigt profile.
-
-    :param elements_list: a List object defining the list of elements to consider
-            when producing spectrum.
-    :param equivalent_sites_data: an Dict object containing symmetry data.
-    :param voight_gamma: a Float node for the gamma parameter of the voigt profile.
-    :param voight_sigma: a Float node for the sigma parameter of the voigt profile.
-    :param structure: the StructureData object to be analysed
-    :returns: Dict objects for all generated spectra and associated binding energy
-            and core level shift.
-
+    :param points: a Dict object containing the binding energy and multiplicity for each site.
+    :param sigma: a Float node for the sigma parameter of the voigt profile.
+    :param gamma: a Float node for the gamma parameter of the voigt profile.
     """
     from scipy.special import voigt_profile  # pylint: disable=no-name-in-module
 
-    ground_state_node = kwargs.pop('ground_state', None)
-    correction_energies = kwargs.pop('correction_energies', orm.Dict()).get_dict()
-    incoming_param_nodes = {key: value for key, value in kwargs.items() if key != 'metadata'}
-    group_state_energy = None
-    if ground_state_node is not None:
-        group_state_energy = ground_state_node.get_dict()['energy']
-    elements = elements_list.get_list()
-    sigma = voight_sigma.value
-    gamma = voight_gamma.value
-    equivalency_data = equivalent_sites_data.get_dict()
-
-    data_dict = {element: {} for element in elements}
-    for key in incoming_param_nodes:
-        xspectra_out_params = incoming_param_nodes[key].get_dict()
-        multiplicity = equivalency_data[key]['multiplicity']
-        element = equivalency_data[key]['symbol']
-        total_energy = xspectra_out_params['energy']
-        data_dict[element][key] = {'element': element, 'multiplicity': multiplicity, 'total_energy': total_energy}
-
-    result = {}
-    core_level_shifts = {}
-    binding_energies = {}
-    for element in elements:
-        spectra_list = []
-        for key in data_dict[element]:
-            site_multiplicity = data_dict[element][key]['multiplicity']
-            spectra_list.append((site_multiplicity, float(data_dict[element][key]['total_energy']), key))
-        spectra_list.sort(key=lambda entry: entry[1])
-        lowest_total_energy = spectra_list[0][1]
-        core_level_shift = [(entry[0], entry[1] - lowest_total_energy, entry[2]) for entry in spectra_list]
-        core_level_shifts[element] = core_level_shift
-        result[f'{element}_cls'] = orm.Dict(dict={entry[2]: entry[1] for entry in core_level_shift})
-
-        if group_state_energy is not None:
-            binding_energy = [(entry[0], entry[1] - group_state_energy + correction_energies[element], entry[2])
-                              for entry in spectra_list]
-            binding_energies[element] = binding_energy
-            result[f'{element}_be'] = orm.Dict(dict={entry[2]: entry[1] for entry in binding_energy})
-
     fwhm_voight = gamma / 2 + np.sqrt(gamma**2 / 4 + sigma**2)
 
-    def spectra_broadening(points, label='cls_spectra'):
-        """Broadening base on the binding energy."""
-        result_spectra = {}
-        for element in elements:
+    result_spectra = {}
+    for element, orbitals in points.items():
+        for orbital, data in orbitals.items():
             final_spectra_y_arrays = []
             final_spectra_y_labels = []
             final_spectra_y_units = []
-
-            total_multiplicity = sum(i[0] for i in points[element])
-
+            total_multiplicity = sum(d['multiplicity'] for d in data.values())
             final_spectra = orm.XyData()
-            max_core_level_shift = points[element][-1][1]
-            min_core_level_shift = points[element][0][1]
+            max_core_level_shift = max([d['energy'] for d in data.values()])
+            min_core_level_shift = min([d['energy'] for d in data.values()])
             # Energy range for the Broadening function
             x_energy_range = np.linspace(
                 min_core_level_shift - fwhm_voight - 1.5, max_core_level_shift + fwhm_voight + 1.5, 500
             )
-
-            for atoms, index in zip(points[element], range(len(points[element]))):
+            for site, d in data.items():
                 # Weight for the spectra of every atom
-                intensity = atoms[0]
-                relative_peak_position = atoms[1]
-                final_spectra_y_labels.append(f'{element}{index}_xps')
+                intensity = d['multiplicity'] / total_multiplicity
+                relative_peak_position = d['energy']
+                final_spectra_y_labels.append(f'{element}_{site}')
                 final_spectra_y_units.append('sigma')
                 final_spectra_y_arrays.append(
-                    intensity * voigt_profile(x_energy_range - relative_peak_position, sigma, gamma) /
-                    total_multiplicity
+                    intensity* voigt_profile(x_energy_range - relative_peak_position, sigma, gamma)
                 )
-
-            final_spectra_y_labels.append(f'{element}_total_xps')
+            final_spectra_y_labels.append(f'{element}_total')
             final_spectra_y_units.append('sigma')
             final_spectra_y_arrays.append(sum(final_spectra_y_arrays))
-
             final_spectra_x_label = 'energy'
             final_spectra_x_units = 'eV'
             final_spectra_x_array = x_energy_range
             final_spectra.set_x(final_spectra_x_array, final_spectra_x_label, final_spectra_x_units)
             final_spectra.set_y(final_spectra_y_arrays, final_spectra_y_labels, final_spectra_y_units)
-            result_spectra[f'{element}_{label}'] = final_spectra
-        return result_spectra
+            result_spectra[f'{element}_{orbital}'] = final_spectra
+    return result_spectra
+
+@calcfunction
+def get_spectra_by_element(core_levels, equivalent_sites_data, voight_gamma, voight_sigma, **kwargs):  # pylint: disable=too-many-statements
+    """Generate the XPS spectra for each element.
 
-    result.update(spectra_broadening(core_level_shifts))
+    Calculate the core level shift and binding energy for each element.
+    Generate the final spectra using the Voigt profile.
+
+    :param core_levels: a Dict object defining the elements and their core-levels to consider
+            when producing spectra, e.g., {"C": ["1s"], "Al": ["2s", "2p"]}.
+    :param equivalent_sites_data: an Dict object containing symmetry data.
+    :param voight_gamma: a Float node for the gamma parameter of the voigt profile.
+    :param voight_sigma: a Float node for the sigma parameter of the voigt profile.
+    :param structure: the StructureData object to be analysed
+    :returns: Dict objects for all generated spectra and associated binding energy
+            and core level shift.
+
+    """
+    from copy import deepcopy
+
+    ground_state_node = kwargs.pop('ground_state', None)
+    correction_energies = kwargs.pop('correction_energies', orm.Dict()).get_dict()
+    output_params_ch_scf = kwargs.pop('output_params_ch_scf', {})
+    group_state_energy = ground_state_node.get_dict()['energy'] if ground_state_node is not None else None
+    core_levels = core_levels.get_dict()
+    equivalency_data = equivalent_sites_data.get_dict()
+    # collect the energy and multiplicity
+    data_dict = {}
+    for element, element_data in output_params_ch_scf.items():
+        data_dict[element] = {}
+        for orbital, orbital_data in element_data.items():
+            data_dict[element][orbital] = {}
+            for site, site_data in orbital_data.items():
+                data_dict[element][orbital][site] = {
+                    'energy': site_data.get_dict()['energy'],
+                    'multiplicity': equivalency_data[site]['multiplicity']
+                }
+
+    result = {}
+    chemical_shifts = deepcopy(data_dict)
+    binding_energies = deepcopy(data_dict)
+    for element, orbitals in chemical_shifts.items():
+        for orbital in orbitals:
+            lowest_energy = min([data['energy'] for data in data_dict[element][orbital].values()])
+            for data in chemical_shifts[element][orbital].values():
+                data['energy'] -= lowest_energy
+            if group_state_energy is not None:
+                for data in binding_energies[element][orbital].values():
+                    data['energy'] += -group_state_energy + correction_energies[element][orbital]
+
+    result['chemical_shifts'] = orm.Dict(chemical_shifts)
+    spectra = spectra_broadening(chemical_shifts,
+                                     sigma=voight_sigma.value,
+                                     gamma=voight_gamma.value)
+    result['chemical_shift_spectra'] = spectra
     if ground_state_node is not None:
-        result.update(spectra_broadening(binding_energies, label='be_spectra'))
+        spectra = spectra_broadening(binding_energies,
+                                     sigma=voight_sigma.value,
+                                     gamma=voight_gamma.value)
+        result['binding_energy_spectra'] = spectra
+        result['binding_energies'] = orm.Dict(binding_energies)
     return result

src/aiida_qe_xspec/gui/xps/model.py

@@ -52,9 +52,10 @@ class XpsConfigurationSettingsModel(ConfigurationSettingsModel, HasInputStructur
     )
     core_levels = tl.Dict(
         key_trait=tl.Unicode(),  # core level
-        value_trait=tl.Bool(),  # whether the core level is included
+        value_trait=tl.List(),
         default_value={},
     )
+    atom_indices = tl.List(trait=tl.Int(), default_value=[])
 
     def update(self, specific=''):
         with self.hold_trait_notifications():
@@ -67,9 +68,9 @@ def get_supported_core_levels(self):
         for key in self.correction_energies:
             element = key.split('_')[0]
             if element not in supported_core_levels:
-                supported_core_levels[element] = [key]
+                supported_core_levels[element] = [key.split('_')[1]]
             else:
-                supported_core_levels[element].append(key)
+                supported_core_levels[element].append(key.split('_')[1])
         return supported_core_levels
 
     def get_model_state(self):
@@ -78,7 +79,8 @@ def get_model_state(self):
             'structure_type': self.structure_type,
             'pseudo_group': self.pseudo_group,
             'correction_energies': self.correction_energies,
-            'core_level_list': list(self.core_levels.keys()),
+            'core_levels': self.core_levels,
+            'atom_indices': self.atom_indices,
         }
 
     def set_model_state(self, parameters: dict):
@@ -91,10 +93,8 @@ def set_model_state(self, parameters: dict):
             self.traits()['structure_type'].default_value,
         )
 
-        core_level_list = parameters.get('core_level_list', [])
-        for orbital in self.core_levels:
-            if orbital in core_level_list:
-                self.core_levels[orbital] = True  # type: ignore
+        self.core_levels = parameters.get('core_levels', [])
+        self.atom_indices = parameters.get('atom_indices', [])
 
     def reset(self):
         with self.hold_trait_notifications():

src/aiida_qe_xspec/gui/xps/result/model.py

@@ -41,7 +41,7 @@ def update_spectrum_options(self):
             self.binding_energies,
             self.equivalent_sites_data,
         ) = export_xps_data(outputs)
-        options = [key.split('_')[0] for key in self.chemical_shifts.keys()]
+        options = [f'{element}_{orbital}' for element, data in self.chemical_shifts.items() for orbital in data.keys()]
         self.spectrum_options = options
         self.spectrum = options[0] if options else None
 

src/aiida_qe_xspec/gui/xps/result/result.py

@@ -172,6 +172,9 @@ def _render(self):
             self.plot,
             upload_container,
         ]
+        self.rendered = True
+        self._post_render()
+        self._update_plot(None)
 
     def _post_render(self):
         self._model.update_spectrum_options()

src/aiida_qe_xspec/gui/xps/result/utils.py

@@ -1,17 +1,9 @@
 def export_xps_data(outputs):
     """Export the data from the XPS workchain"""
-    chemical_shifts = {}
     symmetry_analysis_data = outputs.symmetry_analysis_data.get_dict()
     equivalent_sites_data = symmetry_analysis_data['equivalent_sites_data']
-    if 'chemical_shifts' in outputs:
-        for key, data in outputs.chemical_shifts.items():
-            ele = key[:-4]
-            chemical_shifts[ele] = data.get_dict()
-    binding_energies = {}
-    if 'binding_energies' in outputs:
-        for key, data in outputs.binding_energies.items():
-            ele = key[:-3]
-            binding_energies[ele] = data.get_dict()
+    chemical_shifts = outputs.chemical_shifts.get_dict() if 'chemical_shifts' in outputs else {}
+    binding_energies = outputs.binding_energies.get_dict() if 'binding_energies' in outputs else {}
 
     return (
         chemical_shifts,
@@ -27,29 +19,29 @@ def xps_spectra_broadening(points, equivalent_sites_data, gamma=0.3, sigma=0.3,
 
     result_spectra = {}
     fwhm_voight = gamma / 2 + np.sqrt(gamma**2 / 4 + sigma**2)
-    for element, point in points.items():
-        result_spectra[element] = {}
-        final_spectra_y_arrays = []
-        total_multiplicity = sum([equivalent_sites_data[site]['multiplicity'] for site in point])
-        max_core_level_shift = max(point.values())
-        min_core_level_shift = min(point.values())
-        # Energy range for the Broadening function
-        x_energy_range = np.linspace(
-            min_core_level_shift - fwhm_voight - 1.5,
-            max_core_level_shift + fwhm_voight + 1.5,
-            500,
-        )
-        for site in point:
-            # Weight for the spectra of every atom
-            intensity = equivalent_sites_data[site]['multiplicity'] * intensity
-            relative_core_level_position = point[site]
-            y = (
-                intensity
-                * voigt_profile(x_energy_range - relative_core_level_position, sigma, gamma)
-                / total_multiplicity
+    for element, orbitals in points.items():
+        for orbital, data in orbitals.items():
+            result_spectra[f'{element}_{orbital}'] = {}
+            final_spectra_y_arrays = []
+            total_multiplicity = sum(d['multiplicity'] for d in data.values())
+            max_core_level_shift = max([d['energy'] for d in data.values()])
+            min_core_level_shift = min([d['energy'] for d in data.values()])
+            # Energy range for the Broadening function
+            x_energy_range = np.linspace(
+                min_core_level_shift - fwhm_voight - 1.5,
+                max_core_level_shift + fwhm_voight + 1.5,
+                500,
             )
-            result_spectra[element][site] = [x_energy_range, y]
-            final_spectra_y_arrays.append(y)
-        total = sum(final_spectra_y_arrays)
-        result_spectra[element]['total'] = [x_energy_range, total]
+            for site, d in data.items():
+                # Weight for the spectra of every atom
+                relative_core_level_position = d['energy']
+                y = (
+                    intensity
+                    * voigt_profile(x_energy_range - relative_core_level_position, sigma, gamma)
+                    *d['multiplicity'] / total_multiplicity
+                )
+                result_spectra[f'{element}_{orbital}'][site] = [x_energy_range, y]
+                final_spectra_y_arrays.append(y)
+            total = sum(final_spectra_y_arrays)
+            result_spectra[f'{element}_{orbital}']['total'] = [x_energy_range, total]
     return result_spectra