Source code for ewoksfluo.xrffit.legacy.fit

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 PyMca5.PyMcaIO import HDF5Stack1D
from PyMca5.PyMcaPhysics.xrf import FastXRFLinearFit
from PyMca5.PyMcaPhysics.xrf import McaAdvancedFitBatch

from ...io.hdf5 import split_h5uri
from ..pymca_config import PyMcaXrfConfiguration
from ..pymca_config import modify_pymca_configuration
from ..pymca_config import pymca_configdict_from_file
from ..pymca_config import pymca_configdict_from_model
from ..pymca_config import pymca_configdict_to_model
from ..pymca_config import temporary_pymca_configdict_to_file




def perform_batch_fit(
    xrf_spectra_uris: Sequence[str],
    pymca_configuration: Union[str, PyMcaXrfConfiguration],
    output_buffer,
    energy: Optional[float] = None,
    energy_multiplier: Optional[float] = None,
    mlines: Optional[dict] = None,
    quantification: Optional[dict] = None,
    fast_fitting: bool = False,
) -> None:
    """Fit XRF spectra in batch with one primary beam energy. Least-square fitting.
    If you intend a linear fit, modify the configuration:

        - Get current energy calibration with "Load From Fit"
        - Enable: Perform a Linear Fit
        - Disable: Stripping
        - Strip iterations = 0

    Fast linear least squares:

          - Use SNIP instead of STRIP

    :param xrf_spectra_uris: spectra to fit.
    :param pymca_configuration: configuration file to use.
    :param output_buffer: object that receives the fit results.
    :param energy: primary beam energy. Defaults to None.
    :param mlines: elements (keys) which M line group must be replaced by some M subgroups (values). Defaults to None.
    :param quantification: save massfractions. Defaults to False.
    :param fast_fitting: use SNIP instead of STRIP. Defaults to False.
    :param diagnostics: fit model and residuals. Defaults to False.
    :param figuresofmerit: chi-square value and other figures of merit. Defaults to False.
    :param energy_multiplier: adds a higher energy bound equal to energy*energy_multiplier to include high-energy peaks. Default: no bound is added.
    :raises ValueError: if all paths in filelist do not have the same HDF5 group as parent.
    """
    if not pymca_configuration:
        raise RuntimeError("PyMca configuration is empty")
    if isinstance(pymca_configuration, str):
        pymca_configuration = pymca_configdict_to_model(
            pymca_configdict_from_file(pymca_configuration)
        )

    modify_pymca_configuration(
        pymca_configuration,
        energy=energy,
        mlines=mlines,
        quantification=quantification,
        fast_fitting=fast_fitting,
        energy_multiplier=energy_multiplier,
    )

    filelist, selection, scanlist = _parse_hdf5_uris(xrf_spectra_uris)
    try:
        if fast_fitting:
            _fast_fit_xrf_spectra_uris(
                filelist,
                selection,
                scanlist,
                pymca_configuration,
                output_buffer,
                quantification=quantification,
            )
        else:
            _fit_xrf_spectra_uris(
                filelist,
                selection,
                scanlist,
                pymca_configuration,
                output_buffer,
                quantification=quantification,
            )
    except Exception as e:
        raise RuntimeError(f"PyMca fitting failed for {xrf_spectra_uris}") from e



def _fast_fit_xrf_spectra_uris(
    filelist: List[str],
    selection: Dict[str, List[str]],
    scanlist: List[str],
    pymca_configuration: PyMcaXrfConfiguration,
    output_buffer,
    quantification: Optional[dict] = None,
):
    batch = FastXRFLinearFit.FastXRFLinearFit()
    stack = HDF5Stack1D.HDF5Stack1D(filelist, selection, scanlist=scanlist)
    kwargs = {
        "y": stack,
        "configuration": pymca_configdict_from_model(pymca_configuration),
        "concentrations": bool(quantification),
        "refit": 1,  # needed to avoid negative peak areas
        "outbuffer": output_buffer,
    }
    with output_buffer.saveContext():
        batch.fitMultipleSpectra(**kwargs)


def _fit_xrf_spectra_uris(
    filelist: List[str],
    selection: Dict[str, List[str]],
    scanlist: List[str],
    pymca_configuration: PyMcaXrfConfiguration,
    output_buffer,
    quantification: Optional[dict] = None,
):
    selection["entry"] = scanlist
    kwargs = {
        "filelist": filelist,
        "selection": selection,
        "concentrations": bool(quantification),
        "fitfiles": 0,
        "fitconcfile": 0,
        "outbuffer": output_buffer,
    }
    with temporary_pymca_configdict_to_file(
        pymca_configdict_from_model(pymca_configuration)
    ) as cfgfilename:
        batch = McaAdvancedFitBatch.McaAdvancedFitBatch(cfgfilename, **kwargs)
        with output_buffer.saveContext():
            batch.processList()


def _parse_hdf5_uris(
    uris: Sequence[str],
) -> Tuple[List[str], Dict[str, List[str]], List[str]]:
    split_results = list(zip(*(split_h5uri(uri) for uri in uris)))
    filelist, datasetlist = split_results
    filelist = list(set(filelist))
    if len(filelist) != 1:
        raise ValueError("cannot handle scans in different files")
    entries, datasets = zip(*(_parse_hdf5_dataset(dataset) for dataset in datasetlist))
    datasets = list(set(datasets))
    if len(datasets) != 1:
        raise ValueError("dataset location is each scan must be the same for all scans")
    selection = {"y": datasets[0]}
    scanlist = list(entries)
    return filelist, selection, scanlist


def _parse_hdf5_dataset(dataset: str) -> Tuple[str, str]:
    parts = [s for s in dataset.split("/") if s]
    return parts[0], "/" + "/".join(parts[1:])