Source code for ewoksxas.tasks.autobk

import numpy as np
from ewokscore import Task
from ewokscore.model import BaseInputModel, BaseOutputModel
from larch import xafs
from larch.symboltable import Group
from Orange.data import Table
from pydantic import Field

from ewoksxas.converters import larch, orange
from ewoksxas.converters.resample import Alignment, resample_to_common_grid




class Inputs(BaseInputModel):
    Data: Table
    parameters: dict = Field(default_factory=dict)





class Outputs(BaseOutputModel):
    Data: Table
    Groups: list[Group]





class AutoBK(Task, input_model=Inputs, output_model=Outputs):  # type: ignore
    """Task to calculate k-space for X-ray absorption spectra using Larch."""



    def run(self):  # noqa: C901, PLR0912, PLR0915
        data = self.inputs.Data
        parameters: dict = self.inputs.parameters
        try:
            output_type = parameters.pop("output")
        except KeyError:
            output_type = "chik"

        # Alignment mode is consumed here; never forward it to xafs.autobk.
        alignment = Alignment(parameters.pop("resample", Alignment.INTERPOLATION))

        input_converter = orange.Converter.from_table(data)

        energy, mu = input_converter.features
        try:
            e0_array = input_converter.get_meta_values("e0")
        except KeyError:
            e0_array = [None] * input_converter.n_rows
        try:
            edge_step_array = input_converter.get_meta_values("edge_step")
        except KeyError:
            edge_step_array = [None] * input_converter.n_rows
        groups = larch.create_groups(energy, mu)

        # Select e0 parameter from multiple sources
        if parameters.get("e0") is not None:
            _e0 = parameters.pop("e0")
            e0_array = [_e0 for _ in range(len(groups))]

        # Select edge_step parameter from multiple sources
        if parameters.get("edge_step") is not None:
            _step = parameters.pop("edge_step", None)
            edge_step_array = [_step for _ in range(len(groups))]

        # Ensure that duplicate parameters aren't present
        if "e0" in parameters:
            parameters.pop("e0")
        if "edge_step" in parameters:
            parameters.pop("edge_step")

        for group, e0, edge_step in zip(groups, e0_array, edge_step_array, strict=True):
            xafs.autobk(group, group=group, e0=e0, edge_step=edge_step, **parameters)

        # Ensure e0 and edge_step are in group
        for group, e0, edge_step in zip(groups, e0_array, edge_step_array, strict=True):
            if "e0" not in group:
                group["e0"] = e0 if e0 is not None else group["ek0"]
            if "edge_step" not in group:
                group["edge_step"] = edge_step

        if output_type == "chie":
            # chie lives on the shared energy grid, so lengths always match.
            x_out = energy
            y_out = larch.get_attribute_values(groups, "chie")["chie"]
        else:
            k_list = [group.k for group in groups]
            chi_list = [group.chi for group in groups]

            if len({k.size for k in k_list}) == 1:
                # Every spectrum already shares one k grid: no alignment needed.
                x_out = k_list[0]
                chi = np.asarray(chi_list)
            else:
                x_out, chi = resample_to_common_grid(k_list, chi_list, alignment)

            if output_type == "chi":
                y_out = chi
            elif output_type == "chik":
                y_out = chi * x_out
            else:  # chik2
                y_out = chi * x_out * x_out

        self.outputs.Data = input_converter.with_features(x_out, y_out).to_table()
        self.outputs.Groups = groups






def main():
    pass



if __name__ == "__main__":
    main()