Source code for ewoksfluo.tasks.example_data.tasks

from typing import List
from typing import Tuple

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

from ..models import FluoInputs
from .scan_data import save_2d_xrf_scans

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




class MeshSingleScanSingleDetectorInputs(FluoInputs):
    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(FluoInputs):
    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 MosaicMeshSingleDetectorInputs(FluoInputs):
    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(FluoInputs):
    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 MeshStackSingleDetectorInputs(FluoInputs):
    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(FluoInputs):
    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>"
        )