Source code for ewoksfluo.tasks.sum_detectors.sum_spectra

from typing import List
from typing import Optional

import numpy
from ewokscore import Task
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from ewoksdata.data.hdf5.config import guess_dataset_config
from pydantic import Field

from ...io import hdf5
from ...io import output_uri
from ...io.hdf5 import split_h5uri
from .. import nexus_utils
from ..hdf5_utils import create_hdf5_link
from ..hdf5_utils import link_bliss_scan
from .sum_utils import sum_spectra_from_hdf5




class Inputs(BaseInputModel):
    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="Target HDF5 file URI with optional data path.",
        examples=[
            "/results/dataset.h5",
            "/results/dataset.h5::/1.1",
            "/results/dataset.h5::/1.1/sumspectra",
        ],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Optional group underneath ``output_root_uri``."
    )
    xrf_spectra_uri_template: str = Field(
        default="instrument/{}/data",
        description="HDF5 group URI template of an XRF detector relative to `bliss_scan_uri`.",
    )
    detector_normalization_template: Optional[str] = Field(
        default="1./<instrument/{}/live_time>",
        description="Expression to evaluate and multiply each spectrum with before summing. "
        "HDF5 URI's <...> are relative to ``bliss_scan_uris``.",
    )
    output_detector_name: str = Field(default="mcasum")





class Outputs(BaseOutputModel):
    bliss_scan_uri: str = Field(
        description="Bliss scan URI.", examples=["/data/dataset.h5::/1.1"]
    )
    output_root_uri: str = Field(
        description="Original output root URI received as input.",
        examples=["/results/dataset.h5::/1.1"],
    )
    output_root_group: Optional[str] = Field(
        default=None, description="Original output root group received as input."
    )
    detector_name: str = Field(
        description="Name of the detector containing the summed XRF spectra."
    )
    xrf_spectra_uri_template: str = Field(
        description="HDF5 group URI template of an XRF detector relative to `output_root_uri`.",
    )





class SumXrfSpectra(Task, input_model=Inputs, output_model=Outputs):
    """Add single-scan XRF spectra from multiple detectors."""



    def run(self):
        start_time = nexus_utils.now()

        _, scan_h5path = split_h5uri(self.inputs.bliss_scan_uri)
        output_root_uri = output_uri.compose_full_output_uri(
            self.inputs.output_root_uri,
            default_output_data_path=scan_h5path,
            extra_data_paths=(self.inputs.output_root_group, "sumspectra"),
        )

        sumdetector_name = self.inputs.output_detector_name

        process_config = {
            "detector_normalization_template": self.inputs.detector_normalization_template
        }

        with nexus_utils.save_in_ewoks_process(
            output_root_uri,
            start_time,
            process_config=process_config,
        ) as (process_group, already_existed):
            entry_name = hdf5.split_h5data_path(process_group.name)[0]
            out_entry = process_group.file[entry_name]
            if not already_existed:
                link_bliss_scan(out_entry, self.inputs.bliss_scan_uri, retry_timeout=0)

                self._sum(
                    sumdetector_name,
                    self.inputs.detector_names,
                    out_entry,
                    process_group,
                )

            virtual_bliss_scan_uri = f"{out_entry.file.filename}::{out_entry.name}"

        self.outputs.bliss_scan_uri = virtual_bliss_scan_uri
        self.outputs.detector_name = sumdetector_name
        self.outputs.xrf_spectra_uri_template = self.inputs.xrf_spectra_uri_template
        self.outputs.output_root_uri = virtual_bliss_scan_uri
        self.outputs.output_root_group = self.inputs.output_root_group


    def _sum(
        self,
        sumdetector_name: str,
        detector_names: List[str],
        out_entry,
        process_group: str,
    ) -> None:
        sum_spectra: numpy.ndarray = sum_spectra_from_hdf5(
            self.inputs.bliss_scan_uri,
            self.inputs.xrf_spectra_uri_template,
            self.inputs.detector_normalization_template,
            detector_names,
        )

        nxdata = nexus_utils.create_nxdata(
            process_group, sumdetector_name, signal="data"
        )

        dataset_kwargs = guess_dataset_config(
            scan_shape=(sum_spectra.shape[0],),
            detector_shape=(sum_spectra.shape[1],),
            dtype=sum_spectra.dtype,
        )
        dset = nxdata.create_dataset("data", data=sum_spectra, **dataset_kwargs)
        dset.attrs["interpretation"] = "spectrum"

        nxdetector = out_entry["instrument"].create_group(sumdetector_name)
        nxdetector.attrs["NX_class"] = "NXdetector"
        create_hdf5_link(nxdetector, "data", dset)
        create_hdf5_link(out_entry["measurement"], sumdetector_name, dset)