Source code for ewoksxas.tasks.autobk

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

from ewoksxas.converters import larch, orange




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
        data = self.inputs.Data
        parameters: dict = self.inputs.parameters
        try:
            output_type = parameters.pop("output")
        except KeyError:
            output_type = "chik"

        converter = orange.Converter.from_table(data)

        energy, mu = converter.features
        e0_array = converter.get_meta("e0")
        edge_step_array = converter.get_meta("edge_step")
        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))]
        elif e0_array is None:
            e0_array = [None 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))]
        elif edge_step_array is None:
            edge_step_array = [None 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

        attribute_names = ["k", "chi", "chie"]
        attributes = larch.get_attribute_values(groups, attribute_names)

        if output_type in ["chi", "chik", "chik2"]:
            x_out = attributes["k"][0, :]
        else:
            x_out = energy

        if output_type == "chi":
            y_out = attributes["chi"]
        elif output_type == "chik":
            y_out = attributes["chi"] * x_out
        elif output_type == "chik2":
            y_out = attributes["chi"] * x_out * x_out
        else:
            y_out = attributes["chie"]

        converter = orange.Converter().add_features(x_out, y_out)
        self.outputs.Data = converter.to_table()
        self.outputs.Groups = groups






def main():
    pass



if __name__ == "__main__":
    main()