Source code for ewoksfluo.tasks.sum_detectors.sum_spectra_partial

from typing import Dict
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_groups: Dict[str, List[str]] = Field(
        description="Mapping of output detector name to list of detector names.",
        min_length=1,
        examples=[
            {
                "mcasum1": ["fx_nano_det0", "fx_nano_det1"],
                "mcasum2": ["fx_nano_det3", "fx_nano_det4"],
            },
        ],
    )
    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``.",
    )





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_names: List[str] = Field(
        description="Name of the detectors 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 PartialSumXrfSpectra(Task, input_model=Inputs, output_model=Outputs):
    """Add single-scan XRF spectra from multiple detector groups."""



    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"),
        )

        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)

                for (
                    sumdetector_name,
                    detector_names,
                ) in self.inputs.detector_groups.items():
                    self._sum(
                        sumdetector_name, 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_names = list(self.inputs.detector_groups)
        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)