Source code for blissoda.fluo.parameters.mosaic_fluoxas

from typing import Dict
from typing import List
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Union

from .utils import directories
from .utils import models
from .workflows import mosaic_fluoxas_multi_detector
from .workflows import mosaic_fluoxas_multi_detector_align
from .workflows import mosaic_fluoxas_nofit_align
from .workflows import mosaic_fluoxas_single_detector
from .workflows import mosaic_fluoxas_single_detector_align




def mosaic_fluoxas_workflow_inputs(
    filenames: Sequence[Sequence[str]],
    output_root_uri: str,
    concat_output_root_uri: str,
    scan_ranges: Sequence[Sequence[Tuple[int, int]]],
    config_filenames: Sequence[str],
    instrument_name: Optional[str] = None,
    detectors: Optional[Sequence[Union[int, str]]] = None,
    detector_numbers: Optional[Sequence[int]] = None,
    detector_names: Optional[Sequence[str]] = None,
    fast_fitting: bool = True,
    quantification: bool = True,
    diagnostics: bool = True,
    livetime_ref_value: Union[str, int, float, None] = None,
    counter_ref_value: Union[str, int, float, None] = None,
    counter_name: Optional[str] = None,
    energy_name: Optional[str] = None,
    energy_multiplier: Optional[float] = 2,
    exclude_scans: Optional[Sequence[Sequence[Sequence[int]]]] = None,
    resolution: Optional[Dict[str, Tuple[Union[int, float], str]]] = None,
    do_align: bool = True,
    align_counter: Optional[str] = None,
    align_crop: Optional[bool] = True,
    fast_align_counter_selection: Optional[bool] = False,
    stack_axis: Optional[str] = None,
    real_axes: Optional[List[str]] = None,
    virtual_axes: Optional[Dict[str, str]] = None,
    ignore_axes: Optional[List[str]] = None,
    axis_units: Optional[Dict[str, str]] = None,
    group_by_index: bool = False,
) -> Tuple[str, List[dict]]:

    if len(filenames) != len(scan_ranges):
        raise ValueError(
            f"{len(filenames)} scan ranges (first, last) are needed for {len(filenames)} files"
        )

    if not exclude_scans:
        exclude_scans = [()] * len(filenames)

    parameters = models.common_parameters_model(
        detectors=detectors,
        detector_numbers=detector_numbers,
        detector_names=detector_names,
        config_filenames=config_filenames,
        energy_name=energy_name,
        energy_multiplier=energy_multiplier,
        counter_name=counter_name,
        instrument_name=instrument_name,
        fast_fitting=fast_fitting,
        quantification=quantification,
        diagnostics=diagnostics,
        livetime_ref_value=livetime_ref_value,
        counter_ref_value=counter_ref_value,
        stack=True,
        real_axes=real_axes,
        virtual_axes=virtual_axes,
        ignore_axes=ignore_axes,
        axis_units=axis_units,
    )

    parameters = models.MosaicFluoXasParameters(
        filenames=filenames,
        scan_ranges=scan_ranges,
        exclude_scans=exclude_scans,
        group_by_index=group_by_index,
        concat_output_root_uri=concat_output_root_uri,
        output_root_uri=output_root_uri,
        resolution=resolution,
        align_counter=align_counter,
        fast_align_counter_selection=fast_align_counter_selection,
        align_crop=align_crop,
        stack_axis=stack_axis,
        **parameters.model_dump(),
    )

    if parameters.no_fit:
        workflow = mosaic_fluoxas_nofit_align.WORKFLOW
        inputs = mosaic_fluoxas_nofit_align.workflow_inputs(parameters)
    elif parameters.fit_single_detector:
        if do_align:
            workflow = mosaic_fluoxas_single_detector_align.WORKFLOW
            inputs = mosaic_fluoxas_single_detector_align.workflow_inputs(parameters)
        else:
            workflow = mosaic_fluoxas_single_detector.WORKFLOW
            inputs = mosaic_fluoxas_single_detector.workflow_inputs(parameters)
    else:
        if parameters.sum_spectra:
            raise ValueError(
                f"Sum before fit is not supported yet for FluoXAS. Provide {len(parameters.detector_names)} config files."
            )
        if do_align:
            workflow = mosaic_fluoxas_multi_detector_align.WORKFLOW
            inputs = mosaic_fluoxas_multi_detector_align.workflow_inputs(parameters)
        else:
            workflow = mosaic_fluoxas_multi_detector.WORKFLOW
            inputs = mosaic_fluoxas_multi_detector.workflow_inputs(parameters)

    return workflow, inputs





def mosaic_fluoxas_paths(
    session: str,
    sample: str,
    datasets: Sequence[Sequence[str]],
    scan_ranges: Sequence[Sequence[Tuple[int, int]]],
    config_filenames: Sequence[str],
    dirname: Optional[str] = None,
    demo: bool = False,
) -> models.MosaicFluoXasPaths:
    filenames = [
        [
            directories.raw_directory(session, sample, dataset, demo=demo)
            for dataset in dataset_group
        ]
        for dataset_group in datasets
    ]

    first_dataset = datasets[0][0]
    first_scan_number = scan_ranges[0][0][0]

    concat_filename = directories.processed_path(
        session,
        sample,
        first_dataset,
        dirname,
        f"{sample}_{first_dataset}_concat.h5",
        demo=demo,
    )
    output_filename = directories.processed_path(
        session,
        sample,
        first_dataset,
        dirname,
        f"{sample}_{first_dataset}.h5",
        demo=demo,
    )
    convert_destination = directories.processed_path(
        session,
        sample,
        first_dataset,
        dirname,
        f"{sample}_{first_dataset}_scan{first_scan_number:04d}.json",
        demo=demo,
    )

    concat_output_root_uri = concat_filename
    output_root_uri = f"{output_filename}::/{first_scan_number}.1"
    config_filenames = [
        directories.pymca_config_path(session, s, demo=demo) for s in config_filenames
    ]

    workflow_path = directories.accessible_workflow_path(
        session, sample, first_dataset, dirname, demo=demo
    )

    return models.MosaicFluoXasPaths(
        filenames=filenames,
        output_root_uri=output_root_uri,
        concat_output_root_uri=concat_output_root_uri,
        convert_destination=convert_destination,
        workflow_path=workflow_path,
        config_filenames=config_filenames,
    )