Source code for ewoksfluo.tasks.fit.tasks

from typing import Dict
from typing import List
from typing import Optional
from typing import Union

from ewokscore import Task
from ewokscore.model import BaseOutputModel
from pydantic import Field

from ...xrffit.pymca_config import PyMcaXrfConfiguration
from ..models import FluoInputs
from ..positioner_utils import get_primary_beam_energy_value
from .execute import fit_multi
from .execute import fit_single




class FitSingleScanSingleDetectorInputs(FluoInputs):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )

    detector_name: str = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=["fx_nano_det0"],
    )

    output_root_uri: str = Field(
        description="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    config: Union[str, PyMcaXrfConfiguration] = Field(
        description="PyMca configuration",
        examples=["/config/configuration.cfg", {"peaks": {}, "fit": {}}],
    )

    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )

    process_uri_template: str = Field(default="fit/{}")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None





class FitSingleScanSingleDetectorOutputs(BaseOutputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    output_root_uri: str = Field(
        description="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    xrf_results_uri: str = Field(
        description="Output NXprocess results URI.",
        examples=["/results/dataset.h5::/1.1/fit/mca0/results"],
    )
    detector_name: str = Field(description="Name of the fitted detector.")





class FitSingleScanSingleDetector(
    Task,
    input_model=FitSingleScanSingleDetectorInputs,
    output_model=FitSingleScanSingleDetectorOutputs,
):
    """XRF fit of one scan with one detector"""



    def run(self):
        params = self.get_input_values()
        params["pymca_configuration"] = params.pop("config")
        _parse_energy(params)
        xrf_results_uri, output_root_uri = fit_single(**params)
        self.outputs.xrf_results_uri = xrf_results_uri
        self.outputs.bliss_scan_uri = params["bliss_scan_uri"]
        self.outputs.detector_name = params["detector_name"]
        self.outputs.output_root_uri = output_root_uri






class FitSingleScanMultiDetectorInputs(FluoInputs):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )

    detector_names: List[str] = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=[["fx_nano_det0", "fx_nano_det1"]],
    )

    output_root_uri: str = Field(
        description="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    configs: List[Union[str, PyMcaXrfConfiguration]] = Field(
        description="PyMca configuration",
        examples=[
            ["/config/det0.cfg", "/config/det1.cfg"],
            [{"peaks": {}, "fit": {}}, {"peaks": {}, "fit": {}}],
        ],
    )

    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )

    process_uri_template: str = Field(default="fit/{}")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None





class FitSingleScanMultiDetectorOutputs(BaseOutputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    output_root_uri: str = Field(
        description="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    xrf_results_uris: List[str] = Field(
        description="Output NXprocess results URI.",
        examples=[
            "/results/dataset.h5::/1.1/fit/fx_nano_det0/results",
            "/results/dataset.h5::/1.1/fit/fx_nano_det1/results",
        ],
    )
    detector_names: List[str] = Field(description="Names of the fitted detectors.")





class FitSingleScanMultiDetector(
    Task,
    input_model=FitSingleScanMultiDetectorInputs,
    output_model=FitSingleScanMultiDetectorOutputs,
):
    """XRF fit of one scan with multiple detectors"""



    def run(self):
        params = self.get_input_values()
        params["pymca_configurations"] = params.pop("configs")
        bliss_scan_uri = params.pop("bliss_scan_uri")
        params["bliss_scan_uris"] = [bliss_scan_uri]
        _parse_energies(params)
        xrf_results_uris, output_root_uri = fit_multi(**params)
        self.outputs.xrf_results_uris = xrf_results_uris
        self.outputs.bliss_scan_uri = bliss_scan_uri
        self.outputs.detector_names = params["detector_names"]
        self.outputs.output_root_uri = output_root_uri






class FitStackSingleDetectorInputs(FluoInputs):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )

    detector_name: str = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=["fx_nano_det0"],
    )

    output_root_uri: str = Field(
        description="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    config: Union[str, PyMcaXrfConfiguration] = Field(
        description="PyMca configuration",
        examples=["/config/configuration.cfg", {"peaks": {}, "fit": {}}],
    )

    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )

    process_uri_template: str = Field(default="fit/{}")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None





class FitStackSingleDetectorOutputs(BaseOutputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )
    output_root_uri: str = Field(
        description="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    xrf_results_uri: str = Field(
        description="Output NXprocess results URI.",
        examples=["/results/dataset.h5::/1.1/fit/fx_nano_det0/results"],
    )
    detector_name: str = Field(description="Name of the fitted detector.")





class FitStackSingleDetector(
    Task,
    input_model=FitStackSingleDetectorInputs,
    output_model=FitStackSingleDetectorOutputs,
):
    """XRF fit of a stack of identical scan with one detector"""



    def run(self):
        params = self.get_input_values()
        detector_name = params.pop("detector_name")
        params["detector_names"] = [detector_name]
        params["pymca_configurations"] = [params.pop("config")]
        _parse_energies(params)
        xrf_results_uris, output_root_uri = fit_multi(**params)
        self.outputs.xrf_results_uri = xrf_results_uris[0]
        self.outputs.bliss_scan_uris = params["bliss_scan_uris"]
        self.outputs.detector_name = detector_name
        self.outputs.output_root_uri = output_root_uri






class FitStackMultiDetectorInputs(FluoInputs):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )

    detector_names: List[str] = Field(
        description="NXdetector group name.",
        min_length=1,
        examples=[["fx_nano_det0", "fx_nano_det1"]],
    )

    output_root_uri: str = Field(
        description="Output root URI.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/fit",
        ],
    )
    configs: List[Union[str, PyMcaXrfConfiguration]] = Field(
        description="PyMca configuration",
        examples=[
            ["/config/det0.cfg", "/config/det1.cfg"],
            [{"peaks": {}, "fit": {}}, {"peaks": {}, "fit": {}}],
        ],
    )

    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )

    process_uri_template: str = Field(default="fit/{}")
    energy_name: Optional[str] = None
    energy_uri_template: str = Field(default="instrument/positioners_start/{}")
    quantification: Union[Dict[str, float], bool, None] = None
    energy_multiplier: Optional[float] = None
    fast_fitting: bool = Field(default=True)
    diagnostics: bool = Field(default=False)
    figuresofmerit: Optional[bool] = None





class FitStackMultiDetectorOutputs(BaseOutputModel):
    bliss_scan_uris: List[str] = Field(
        description="Bliss scan URI's.",
        examples=[["/data/dataset.h5::/1.1", "/data/dataset.h5::/2.1"]],
    )
    output_root_uri: str = Field(
        description="Output NXentry URI.", examples=["/results/dataset.h5::/1.1"]
    )
    xrf_results_uris: List[str] = Field(
        description="Output NXprocess results URI.",
        examples=[
            "/results/dataset.h5::/1.1/fit/fx_nano_det0/results",
            "/results/dataset.h5::/1.1/fit/fx_nano_det1/results",
        ],
    )
    detector_names: List[str] = Field(description="Names of the fitted detectors.")





class FitStackMultiDetector(
    Task,
    input_model=FitStackMultiDetectorInputs,
    output_model=FitStackMultiDetectorOutputs,
):
    """XRF fit of a stack of identical scan with multiple detectors"""



    def run(self):
        params = self.get_input_values()
        params["pymca_configurations"] = params.pop("configs")
        _parse_energies(params)
        xrf_results_uris, output_root_uri = fit_multi(**params)
        self.outputs.xrf_results_uris = xrf_results_uris
        self.outputs.bliss_scan_uris = params["bliss_scan_uris"]
        self.outputs.detector_names = params["detector_names"]
        self.outputs.output_root_uri = output_root_uri




def _parse_energy(params: dict) -> None:
    energy_name = params.pop("energy_name", None)
    energy_uri_template = params.pop("energy_uri_template", None)
    params["energy"] = get_primary_beam_energy_value(
        params["bliss_scan_uri"],
        energy_name=energy_name,
        energy_uri_template=energy_uri_template,
        search_on_units=False,
    )


def _parse_energies(params: dict) -> None:
    energy_name = params.pop("energy_name", None)
    energy_uri_template = params.pop("energy_uri_template", None)
    params["energies"] = [
        get_primary_beam_energy_value(
            bliss_scan_uri,
            energy_name=energy_name,
            energy_uri_template=energy_uri_template,
            search_on_units=False,
        )
        for bliss_scan_uri in params["bliss_scan_uris"]
    ]