Source code for ewoksfluo.tasks.example_data.tasks

from typing import Dict
from typing import List
from typing import Tuple

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

from .scan_data import save_2d_xrf_scans

SEED = 100
MAX_DEVIATION = 0.05  # Fraction of a single step




class MeshSingleScanSingleDetectorInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    expo_time: float = Field(default=0.1)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MeshSingleScanSingleDetectorOutputs(BaseOutputModel):
    filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    scan_number: int = Field(description="Scan number.", examples=[1, 2])
    config: str
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_name: str
    detector_normalization_template: str





class MeshSingleScanSingleDetector(
    Task,
    input_model=MeshSingleScanSingleDetectorInputs,
    output_model=MeshSingleScanSingleDetectorOutputs,
):
    """XRF test data of one scan with one detector"""



    def run(self):
        scan_number = 1
        filename = self.inputs.output_filename

        _ = save_2d_xrf_scans(
            filename=filename,
            emission_line_groups=self.inputs.emission_line_groups,
            first_scan_number=scan_number,
            shape=self.inputs.shape,
            mosaic=(1, 1),
            energy=self.inputs.energy,
            flux=self.inputs.flux,
            expo_time=self.inputs.expo_time,
            counting_noise=self.inputs.counting_noise,
            integral_type=self.inputs.integral_type,
            rois=self.inputs.rois,
            nmcas=1,
            max_deviation=MAX_DEVIATION,
            seed=SEED,
        )

        self.outputs.filename = filename
        self.outputs.scan_number = scan_number
        self.outputs.config = f"{filename}::/{scan_number}.1/theory/configuration/data"
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.detector_name = "mca0"
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MeshSingleScanMultiDetectorInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    expo_time: float = Field(default=0.1)
    ndetectors: int = Field(default=2)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MeshSingleScanMultiDetectorOutputs(BaseOutputModel):
    filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    scan_number: int = Field(description="Scan number.", examples=[1, 2])
    config: str
    configs: List[str]
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_names: List[str]
    detector_normalization_template: str





class MeshSingleScanMultiDetector(
    Task,
    input_model=MeshSingleScanMultiDetectorInputs,
    output_model=MeshSingleScanMultiDetectorOutputs,
):
    """XRF test data of one scan with multiple detectors"""



    def run(self):
        scan_number = 1
        filename = self.inputs.output_filename

        _ = save_2d_xrf_scans(
            filename=filename,
            emission_line_groups=self.inputs.emission_line_groups,
            first_scan_number=scan_number,
            shape=self.inputs.shape,
            mosaic=(1, 1),
            energy=self.inputs.energy,
            flux=self.inputs.flux,
            expo_time=self.inputs.expo_time,
            counting_noise=self.inputs.counting_noise,
            integral_type=self.inputs.integral_type,
            rois=self.inputs.rois,
            nmcas=self.inputs.ndetectors,
            max_deviation=MAX_DEVIATION,
            seed=SEED,
        )

        self.outputs.filename = filename
        self.outputs.scan_number = scan_number
        self.outputs.configs = [
            f"{filename}::/{scan_number}.1/theory/configuration/data"
        ] * self.inputs.ndetectors
        self.outputs.config = f"{filename}::/{scan_number}.1/theory/configuration/data"
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.detector_names = [f"mca{i}" for i in range(self.inputs.ndetectors)]
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MeshSingleScanMultiDetectorGroupsInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    expo_time: float = Field(default=0.1)
    ngroups: int = Field(default=2)
    ndetectors: int = Field(default=2)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MeshSingleScanMultiDetectorGroupsOutputs(BaseOutputModel):
    filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    scan_number: int = Field(description="Scan number.", examples=[1, 2])
    config: str
    configs: List[str]
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_groups: Dict[str, List[str]]
    detector_normalization_template: str





class MeshSingleScanMultiDetectorGroups(
    Task,
    input_model=MeshSingleScanMultiDetectorGroupsInputs,
    output_model=MeshSingleScanMultiDetectorGroupsOutputs,
):
    """XRF test data of one scan with multiple detector groups"""



    def run(self):
        scan_number = 1
        filename = self.inputs.output_filename

        _ = save_2d_xrf_scans(
            filename=filename,
            emission_line_groups=self.inputs.emission_line_groups,
            first_scan_number=scan_number,
            shape=self.inputs.shape,
            mosaic=(1, 1),
            energy=self.inputs.energy,
            flux=self.inputs.flux,
            expo_time=self.inputs.expo_time,
            counting_noise=self.inputs.counting_noise,
            integral_type=self.inputs.integral_type,
            rois=self.inputs.rois,
            nmcas=self.inputs.ndetectors * self.inputs.ngroups,
            max_deviation=MAX_DEVIATION,
            seed=SEED,
        )

        self.outputs.filename = filename
        self.outputs.scan_number = scan_number
        self.outputs.configs = [
            f"{filename}::/{scan_number}.1/theory/configuration/data"
        ] * self.inputs.ngroups
        self.outputs.config = f"{filename}::/{scan_number}.1/theory/configuration/data"
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.detector_groups = {
            f"mcasum{i}": [
                f"mca{i*self.inputs.ndetectors + j}"
                for j in range(self.inputs.ndetectors)
            ]
            for i in range(self.inputs.ngroups)
        }
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MosaicMeshSingleDetectorInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    mosaic: Tuple[int, int] = Field(default=(2, 3))
    expo_time: float = Field(default=0.1)
    ndetectors: int = Field(default=1)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MosaicMeshSingleDetectorOutputs(BaseOutputModel):
    filenames: List[str]
    scan_ranges: List[Tuple[int, int]]
    config: str
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_name: str
    detector_normalization_template: str





class MosaicMeshSingleDetector(
    Task,
    input_model=MosaicMeshSingleDetectorInputs,
    output_model=MosaicMeshSingleDetectorOutputs,
):
    """XRF test data of a mosaic scan with one detector"""



    def run(self):
        filename = self.inputs.output_filename

        scan_numbers = save_2d_xrf_scans(
            filename=filename,
            emission_line_groups=self.inputs.emission_line_groups,
            first_scan_number=1,
            shape=self.inputs.shape,
            mosaic=self.inputs.mosaic,
            energy=self.inputs.energy,
            flux=self.inputs.flux,
            expo_time=self.inputs.expo_time,
            counting_noise=self.inputs.counting_noise,
            integral_type=self.inputs.integral_type,
            rois=self.inputs.rois,
            nmcas=1,
            max_deviation=MAX_DEVIATION,
            seed=SEED,
        )

        self.outputs.filenames = [filename]
        self.outputs.scan_ranges = [[scan_numbers[0], scan_numbers[-1]]]
        self.outputs.config = (
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        )
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.detector_name = "mca0"
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MosaicMeshMultiDetectorInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    mosaic: Tuple[int, int] = Field(default=(2, 3))
    expo_time: float = Field(default=0.1)
    ndetectors: int = Field(default=2)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MosaicMeshMultiDetectorOutputs(BaseOutputModel):
    filenames: List[str]
    scan_ranges: List[Tuple[int, int]]
    configs: List[str]
    config: str
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_names: List[str]
    detector_normalization_template: str





class MosaicMeshMultiDetector(
    Task,
    input_model=MosaicMeshMultiDetectorInputs,
    output_model=MosaicMeshMultiDetectorOutputs,
):
    """XRF test data of a mosaic scan with multiple detectors"""



    def run(self):
        filename = self.inputs.output_filename

        scan_numbers = save_2d_xrf_scans(
            filename=filename,
            emission_line_groups=self.inputs.emission_line_groups,
            first_scan_number=1,
            shape=self.inputs.shape,
            mosaic=self.inputs.mosaic,
            energy=self.inputs.energy,
            flux=self.inputs.flux,
            expo_time=self.inputs.expo_time,
            counting_noise=self.inputs.counting_noise,
            integral_type=self.inputs.integral_type,
            rois=self.inputs.rois,
            nmcas=self.inputs.ndetectors,
            max_deviation=MAX_DEVIATION,
            seed=SEED,
        )

        self.outputs.filenames = [filename]
        self.outputs.scan_ranges = [[scan_numbers[0], scan_numbers[-1]]]
        self.outputs.configs = [
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        ] * self.inputs.ndetectors
        self.outputs.config = (
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        )
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.detector_names = [f"mca{i}" for i in range(self.inputs.ndetectors)]
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MosaicMeshMultiDetectorGroupsInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    mosaic: Tuple[int, int] = Field(default=(2, 3))
    expo_time: float = Field(default=0.1)
    ngroups: int = Field(default=2)
    ndetectors: int = Field(default=2)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MosaicMeshMultiDetectorGroupsOutputs(BaseOutputModel):
    filenames: List[str]
    scan_ranges: List[Tuple[int, int]]
    configs: List[str]
    config: str
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_groups: Dict[str, List[str]]
    detector_normalization_template: str





class MosaicMeshMultiDetectorGroups(
    Task,
    input_model=MosaicMeshMultiDetectorGroupsInputs,
    output_model=MosaicMeshMultiDetectorGroupsOutputs,
):
    """XRF test data of a mosaic scan with multiple detectors"""



    def run(self):
        filename = self.inputs.output_filename

        scan_numbers = save_2d_xrf_scans(
            filename=filename,
            emission_line_groups=self.inputs.emission_line_groups,
            first_scan_number=1,
            shape=self.inputs.shape,
            mosaic=self.inputs.mosaic,
            energy=self.inputs.energy,
            flux=self.inputs.flux,
            expo_time=self.inputs.expo_time,
            counting_noise=self.inputs.counting_noise,
            integral_type=self.inputs.integral_type,
            rois=self.inputs.rois,
            nmcas=self.inputs.ndetectors * self.inputs.ngroups,
            max_deviation=MAX_DEVIATION,
            seed=SEED,
        )

        self.outputs.filenames = [filename]
        self.outputs.scan_ranges = [[scan_numbers[0], scan_numbers[-1]]]
        self.outputs.configs = [
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        ] * self.inputs.ndetectors
        self.outputs.config = (
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        )
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.detector_groups = {
            f"mcasum{i}": [
                f"mca{i*self.inputs.ndetectors + j}"
                for j in range(self.inputs.ndetectors)
            ]
            for i in range(self.inputs.ngroups)
        }
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MeshStackSingleDetectorInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    nscans: int = Field(default=3)
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    expo_time: float = Field(default=0.1)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MeshStackSingleDetectorOutputs(BaseOutputModel):
    filenames: List[str] = Field(
        description="Bliss dataset HDF5 file name.",
        examples=[["/data/dataset1.h5", "/data/dataset2.h5"]],
    )
    scan_ranges: List[Tuple[int, int]]
    config: str
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_name: str
    detector_normalization_template: str





class MeshStackSingleDetector(
    Task,
    input_model=MeshStackSingleDetectorInputs,
    output_model=MeshStackSingleDetectorOutputs,
):
    """XRF test data of a stack of identical scans with one detector"""



    def run(self):
        filename = self.inputs.output_filename

        bliss_scan_uris = list()
        scan_numbers = list(range(1, self.inputs.nscans + 1))
        energy = self.inputs.energy
        for scan_number in scan_numbers:
            _ = save_2d_xrf_scans(
                filename=filename,
                emission_line_groups=self.inputs.emission_line_groups,
                first_scan_number=scan_number,
                shape=self.inputs.shape,
                mosaic=(1, 1),
                energy=energy,
                flux=self.inputs.flux,
                expo_time=self.inputs.expo_time,
                counting_noise=self.inputs.counting_noise,
                integral_type=self.inputs.integral_type,
                rois=self.inputs.rois,
                nmcas=1,
                max_deviation=MAX_DEVIATION,
                seed=SEED,
            )

            bliss_scan_uris.append(f"{filename}::/{scan_number}.1")
            energy += 0.010

        self.outputs.filenames = [filename]
        self.outputs.scan_ranges = [[scan_numbers[0], scan_numbers[-1]]]
        self.outputs.config = (
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        )
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.detector_name = "mca0"
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )






class MeshStackMultiDetectorInputs(BaseInputModel):
    output_filename: str = Field(
        description="Bliss dataset HDF5 file name.", examples=["/data/dataset.h5"]
    )
    nscans: int = Field(default=3)
    emission_line_groups: List[str] = Field(
        default_factory=lambda: ["Si-K", "Ca-K", "Ce-L", "Fe-K"]
    )
    energy: float = Field(default=12.0)
    shape: Tuple[int, int] = Field(default=(50, 60))
    expo_time: float = Field(default=0.1)
    ndetectors: int = Field(default=2)
    flux: float = Field(default=1e7)
    counting_noise: bool = Field(default=True)
    rois: List[Tuple[int, int]] = Field(
        default_factory=lambda: [(100, 200), (300, 600)]
    )
    integral_type: bool = Field(default=True)





class MeshStackMultiDetectorOutputs(BaseOutputModel):
    filenames: List[str] = Field(
        description="Bliss dataset HDF5 file name.",
        examples=[["/data/dataset1.h5", "/data/dataset2.h5"]],
    )
    scan_ranges: List[Tuple[int, int]]
    configs: List[str]
    config: str
    expo_time: float
    monitor_name: str
    monitor_normalization_template: str
    detector_names: List[str]
    detector_normalization_template: str





class MeshStackMultiDetector(
    Task,
    input_model=MeshStackMultiDetectorInputs,
    output_model=MeshStackMultiDetectorOutputs,
):
    """XRF test data of a stack of identical scans with multiple detectors"""



    def run(self):
        filename = self.inputs.output_filename

        bliss_scan_uris = list()
        scan_numbers = list(range(1, self.inputs.nscans + 1))
        energy = self.inputs.energy
        for scan_number in scan_numbers:
            _ = save_2d_xrf_scans(
                filename=filename,
                emission_line_groups=self.inputs.emission_line_groups,
                first_scan_number=scan_number,
                shape=self.inputs.shape,
                mosaic=(1, 1),
                energy=energy,
                flux=self.inputs.flux,
                expo_time=self.inputs.expo_time,
                counting_noise=self.inputs.counting_noise,
                integral_type=self.inputs.integral_type,
                rois=self.inputs.rois,
                nmcas=self.inputs.ndetectors,
                max_deviation=MAX_DEVIATION,
                seed=SEED,
            )

            bliss_scan_uris.append(f"{filename}::/{scan_number}.1")
            energy += 0.010

        self.outputs.filenames = [filename]
        self.outputs.scan_ranges = [[scan_numbers[0], scan_numbers[-1]]]
        self.outputs.configs = [
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        ] * self.inputs.ndetectors
        self.outputs.config = (
            f"{filename}::/{scan_numbers[0]}.1/theory/configuration/data"
        )
        self.outputs.monitor_name = "I0"
        self.outputs.monitor_normalization_template = (
            f"{int(self.inputs.flux * self.inputs.expo_time)}/<instrument/{{}}/data>"
        )
        self.outputs.expo_time = self.inputs.expo_time
        self.outputs.detector_names = [f"mca{i}" for i in range(self.inputs.ndetectors)]
        self.outputs.detector_normalization_template = (
            f"{self.inputs.expo_time}/<instrument/{{}}/live_time>"
        )