ewoksid02 0.1.0__py3-none-any.whl

ewoksid02/utils/pyfai.py

@@ -0,0 +1,424 @@
+import logging
+from typing import Optional
+import h5py
+import numexpr
+import numpy
+import os
+from pyFAI import detector_factory
+from pyFAI.method_registry import IntegrationMethod, Method
+from pyFAI.worker import Worker
+from functools import lru_cache
+from ewoksid02.utils.io import (
+    get_from_headers,
+    get_persistent_array_mask,
+    get_persistent_array_dark,
+)
+from ewoksid02.utils.normalization import AzimuthalIntegrator
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+@lru_cache(maxsize=5)
+def _get_persistent_detector_instance(
+    data_signal_shape: tuple,
+    psize_1: float,
+    psize_2: float,
+) -> Optional[object]:
+    """
+    Create a detector instance based on the dataset and pixel sizes.
+
+    Parameters:
+        data_signal_shape (tuple): Shape of the 2D data signal frame.
+        psize_1 (float): Pixel size in the first dimension.
+        psize_2 (float): Pixel size in the second dimension.
+
+    Returns:
+        Optional[object]: A detector instance or None if the dataset is None.
+    """
+    detector_config = {
+        "pixel1": psize_2,
+        "pixel2": psize_1,
+    }
+    detector_instance = detector_factory(name="detector", config=detector_config)
+    detector_instance.shape = data_signal_shape
+    return detector_instance
+
+
+def get_persistent_azimuthal_integrator(
+    data_signal_shape: tuple,
+    headers: dict = None,
+    metadata_parameters_group: h5py.Group = None,
+    **kwargs,
+):
+    params_headers = {
+        "headers": headers,
+        "metadata_file_group": metadata_parameters_group,
+    }
+
+    SampleDistance = kwargs.get("SampleDistance")
+    if SampleDistance is None:
+        SampleDistance = get_from_headers(key="SampleDistance", **params_headers)
+    Center_1 = kwargs.get("Center_1")
+    if Center_1 is None:
+        Center_1 = get_from_headers(key="Center_1", **params_headers)
+    Center_2 = kwargs.get("Center_2")
+    if Center_2 is None:
+        Center_2 = get_from_headers(key="Center_2", **params_headers)
+    PSize_1 = kwargs.get("PSize_1")
+    if PSize_1 is None:
+        PSize_1 = get_from_headers(key="PSize_1", **params_headers)
+    PSize_2 = kwargs.get("PSize_2")
+    if PSize_2 is None:
+        PSize_2 = get_from_headers(key="PSize_2", **params_headers)
+    BSize_1 = kwargs.get("BSize_1")
+    if BSize_1 is None:
+        BSize_1 = get_from_headers(key="BSize_1", **params_headers)
+    BSize_2 = kwargs.get("BSize_2")
+    if BSize_2 is None:
+        BSize_2 = get_from_headers(key="BSize_2", **params_headers)
+    wavelength = kwargs.get("WaveLength")
+    if wavelength is None:
+        wavelength = get_from_headers(key="WaveLength", **params_headers)
+    rot1 = kwargs.get("DetectorRotation_2")
+    if rot1 is None:
+        rot1 = get_from_headers(key="DetectorRotation_2", **params_headers)
+    rot2 = kwargs.get("DetectorRotation_1")
+    if rot2 is None:
+        rot2 = get_from_headers(key="DetectorRotation_1", **params_headers)
+    rot3 = kwargs.get("DetectorRotation_3")
+    if rot3 is None:
+        rot3 = get_from_headers(key="DetectorRotation_3", **params_headers)
+
+    if BSize_1 != 1 or BSize_2 != 1:
+        logger.warning(f"Binning is activated in this detector: ({BSize_1}, {BSize_2})")
+
+    return _get_persistent_azimuthal_integrator(
+        data_signal_shape=data_signal_shape,
+        dist=SampleDistance,
+        center_1=Center_1,
+        center_2=Center_2,
+        psize_1=PSize_1,
+        psize_2=PSize_2,
+        rot1=rot1,
+        rot2=rot2,
+        rot3=rot3,
+        wavelength=wavelength,
+    )
+
+
+@lru_cache(maxsize=5)
+def _get_persistent_azimuthal_integrator(
+    data_signal_shape: tuple,
+    dist: float,
+    wavelength: float,
+    psize_1: float,
+    psize_2: float,
+    center_1: float = None,
+    center_2: float = None,
+    poni1: float = None,
+    poni2: float = None,
+    rot1: float = 0.0,
+    rot2: float = 0.0,
+    rot3: float = 0.0,
+):
+    detector = _get_persistent_detector_instance(
+        data_signal_shape=data_signal_shape,
+        psize_1=psize_1,
+        psize_2=psize_2,
+    )
+
+    if poni1 is None and center_2:
+        poni1 = center_2 * psize_2
+    if poni2 is None and center_1:
+        poni2 = center_1 * psize_1
+
+    ai_kwargs = {
+        "dist": dist,
+        "poni1": poni1,
+        "poni2": poni2,
+        "rot1": rot1,
+        "rot2": rot2,
+        "rot3": rot3,
+        "wavelength": wavelength,
+        "detector": detector,
+    }
+    ai = AzimuthalIntegrator(**ai_kwargs)
+    ai.solidAngleArray(absolute=False)
+    return ai
+
+
+def get_gpu_method():
+    if (
+        IntegrationMethod.select_method(dim=1, split="full", algo="csr", impl="opencl")[
+            0
+        ].impl
+        == "OpenCL"
+    ):
+        return Method(dim=2, split="full", algo="csr", impl="opencl", target=(0, 0))
+    return Method(dim=2, split="full", algo="csr", impl="cython", target=None)
+
+
+def guess_npt2_rad(azimuthal_integrator: AzimuthalIntegrator) -> int:
+    qmax = azimuthal_integrator.qArray().max()
+    dqmin = azimuthal_integrator.deltaQ().min() * 2.0
+    return int(qmax / dqmin)
+
+
+def get_persistent_pyfai_worker(
+    azimuthal_integrator: AzimuthalIntegrator,
+    filename_mask: str = None,
+    integration_options: dict = {},
+    npt2_rad: int = None,
+    npt2_azim: int = None,
+    unit: str = None,
+    Dummy: int = None,
+    DDummy: float = None,
+    method: str = get_gpu_method(),
+    binning: tuple = (1, 1),
+    **kwargs,
+):
+    if filename_mask and os.path.exists(filename_mask):
+        mtime_mask = os.path.getmtime(filename_mask)
+    else:
+        mtime_mask = None
+        filename_mask = None
+
+    return _get_persistent_pyfai_worker(
+        data_signal_shape=azimuthal_integrator.detector.shape,
+        dist=azimuthal_integrator.dist,
+        wavelength=azimuthal_integrator.wavelength,
+        psize_1=azimuthal_integrator.pixel1,
+        psize_2=azimuthal_integrator.pixel2,
+        poni1=azimuthal_integrator.poni1,
+        poni2=azimuthal_integrator.poni2,
+        rot1=azimuthal_integrator.rot1,
+        rot2=azimuthal_integrator.rot2,
+        rot3=azimuthal_integrator.rot3,
+        filename_mask=filename_mask,
+        mtime_mask=mtime_mask,
+        npt2_rad=npt2_rad,
+        npt2_azim=npt2_azim,
+        unit=unit,
+        Dummy=Dummy,
+        DDummy=DDummy,
+        method=method,
+        binning=binning,
+        **integration_options,
+    )
+
+
+@lru_cache(maxsize=5)
+def _get_persistent_pyfai_worker(
+    data_signal_shape: tuple,
+    dist: float,
+    wavelength: float,
+    psize_1: float,
+    psize_2: float,
+    center_1: float = None,
+    center_2: float = None,
+    poni1: float = None,
+    poni2: float = None,
+    rot1: float = 0.0,
+    rot2: float = 0.0,
+    rot3: float = 0.0,
+    filename_mask: str = None,
+    mtime_mask: float = None,
+    npt2_rad: int = None,
+    npt2_azim: int = None,
+    unit: str = None,
+    Dummy: int = None,
+    DDummy: float = None,
+    method: str = get_gpu_method(),
+    binning: tuple = (1, 1),
+    **integration_options,
+):
+    azimuthal_integrator = _get_persistent_azimuthal_integrator(
+        data_signal_shape=data_signal_shape,
+        dist=dist,
+        center_1=center_1,
+        center_2=center_2,
+        psize_1=psize_1,
+        psize_2=psize_2,
+        poni1=poni1,
+        poni2=poni2,
+        rot1=rot1,
+        rot2=rot2,
+        rot3=rot3,
+        wavelength=wavelength,
+    )
+
+    # integration options
+    applied_integration_options = {
+        "npt2_rad": int(npt2_rad),
+        "npt2_azim": int(npt2_azim),
+        "unit": unit,
+        "dummy": Dummy,
+        "delta_dummy": DDummy,
+        "method": (
+            (method.split, method.algo, method.impl)
+            if isinstance(method, (IntegrationMethod, Method))
+            else method
+        ),
+    }
+    applied_integration_options.update(**integration_options)
+
+    worker = Worker(
+        azimuthalIntegrator=azimuthal_integrator,
+        shapeIn=azimuthal_integrator.detector.shape,
+        shapeOut=(
+            applied_integration_options.pop("npt2_azim"),
+            applied_integration_options.pop("npt2_rad"),
+        ),
+        unit=applied_integration_options.pop("unit"),
+        dummy=applied_integration_options.pop("dummy"),
+        delta_dummy=applied_integration_options.pop("delta_dummy"),
+        method=applied_integration_options.pop("method", get_gpu_method()),
+        extra_options=applied_integration_options,
+    )
+    worker.correct_solid_angle = False
+    worker.output = "raw"
+    worker.error_model = "azimuthal"
+    worker.ai._empty = worker.dummy
+
+    array_mask = get_persistent_array_mask(
+        filename_mask=filename_mask,
+        data_signal_shape=data_signal_shape,
+        binning=binning,
+        use_cupy=False,
+    )
+    if array_mask is not None:
+        worker.mask_image = array_mask
+        worker.ai.mask = array_mask
+    return worker
+
+
+def process_dataset_azim(
+    dataset_signal: numpy.ndarray,
+    azimuthal_integrator: AzimuthalIntegrator,
+    dataset_variance: numpy.ndarray = None,
+    filename_dark: str = None,
+    filename_mask: str = None,
+    npt2_rad: int = None,
+    npt2_azim: int = None,
+    unit: str = None,
+    Dummy: int = None,
+    DDummy: int = None,
+    method: str = None,
+    datatype: str = "float32",
+    **kwargs,
+):
+    dataset_signal_azim = None
+    dataset_variance_azim = None
+    dataset_sigma_azim = None
+    dataset_sumsignal_azim = None
+    dataset_sumnorm_azim = None
+    dataset_sumvariance_azim = None
+
+    do_variance_formula = kwargs.get("do_variance_formula")
+    variance_formula = kwargs.get("variance_formula")
+    if dataset_variance is not None:
+        ...
+    elif dataset_variance is None and do_variance_formula and variance_formula:
+        variance_function = numexpr.NumExpr(
+            variance_formula, [("data", datatype), ("dark", datatype)]
+        )
+        dark = None
+        if filename_dark:
+            dark = get_persistent_array_dark(
+                filename_dark=filename_dark,
+                data_signal_shape=dataset_signal[0].shape,
+                dummy=Dummy,
+                delta_dummy=DDummy,
+                **kwargs,
+            )
+
+        dataset_variance = numpy.zeros_like(dataset_signal, dtype=datatype)
+        for frame_index, data_signal in enumerate(dataset_signal):
+            dataset_variance[frame_index] = variance_function(
+                data_signal, 0 if dark is None else dark
+            )
+
+    worker = get_persistent_pyfai_worker(
+        azimuthal_integrator=azimuthal_integrator,
+        filename_mask=filename_mask,
+        npt2_rad=npt2_rad,
+        npt2_azim=npt2_azim,
+        unit=unit,
+        Dummy=Dummy,
+        DDummy=DDummy,
+        method=method,
+        **kwargs,
+    )
+
+    output_dataset_shape = (
+        len(dataset_signal),
+        worker.nbpt_azim,
+        worker.nbpt_rad,
+    )
+
+    dataset_signal_azim = numpy.zeros(output_dataset_shape, dtype=datatype)
+    dataset_sumsignal_azim = numpy.zeros(output_dataset_shape, dtype=datatype)
+    dataset_sumnorm_azim = numpy.zeros(output_dataset_shape, dtype=datatype)
+
+    if dataset_variance is not None:
+        dataset_variance_azim = numpy.zeros(output_dataset_shape, dtype=datatype)
+        dataset_sigma_azim = numpy.zeros(output_dataset_shape, dtype=datatype)
+        dataset_sumvariance_azim = numpy.zeros(output_dataset_shape, dtype=datatype)
+
+    logger.debug(f"PyFAI worker used method: {worker.method}")
+    for frame_index, data_signal in enumerate(dataset_signal):
+        if dataset_variance is not None:
+            data_variance = dataset_variance[frame_index]
+        else:
+            data_variance = None
+
+        result2d_azim = worker.process(
+            data=data_signal,
+            variance=data_variance,
+        )
+
+        dataset_signal_azim[frame_index] = numpy.ascontiguousarray(
+            result2d_azim.intensity, dtype=datatype
+        )  # (sum_signal / sum_normalization)
+        dataset_sumsignal_azim[frame_index] = numpy.ascontiguousarray(
+            result2d_azim.sum_signal, dtype=datatype
+        )
+        dataset_sumnorm_azim[frame_index] = numpy.ascontiguousarray(
+            result2d_azim.sum_normalization, dtype=datatype
+        )
+
+        if data_variance is not None:
+            sum_var = numpy.ascontiguousarray(  # noqa
+                result2d_azim.sum_variance, dtype=datatype
+            )
+            sum_norm = numpy.ascontiguousarray(  # noqa
+                result2d_azim.sum_normalization, dtype=datatype
+            )
+            dataset_variance_azim[frame_index] = numexpr.evaluate(
+                f"where(sum_norm <= 0.0, {worker.dummy}, sum_var / (sum_norm * sum_norm))"
+            )
+            dataset_sigma_azim[frame_index] = numpy.ascontiguousarray(
+                result2d_azim.sigma, dtype=datatype
+            )
+            dataset_sumvariance_azim[frame_index] = numpy.ascontiguousarray(
+                result2d_azim.sum_variance, dtype=datatype
+            )
+
+    radial_array = result2d_azim.radial
+    if radial_array is not None:
+        radial_array = radial_array.astype(datatype)
+    azimuthal_array = result2d_azim.azimuthal
+    if azimuthal_array is not None:
+        azimuthal_array = azimuthal_array.astype(datatype)
+
+    return (
+        dataset_signal_azim,
+        dataset_variance_azim,
+        dataset_sigma_azim,
+        dataset_sumsignal_azim,
+        dataset_sumnorm_azim,
+        dataset_sumvariance_azim,
+        radial_array,
+        azimuthal_array,
+    )