ewoksid02 0.1.0__py3-none-any.whl

ewoksid02/utils/io.py

@@ -0,0 +1,722 @@
+import logging
+import os
+import time
+from functools import lru_cache
+from typing import Dict, List, Optional, Union
+from silx.io.h5py_utils import open_item as open_item_silx
+from silx.utils.retry import RetryTimeoutError
+from pyFAI.utils.mathutil import binning as binning_tool
+from pyFAI.average import average_dark
+
+import fabio
+import h5py
+import hdf5plugin  # noqa
+import numpy
+import re
+
+try:
+    import cupy
+except ImportError:
+    CUPY_AVAILABLE = False
+else:
+    CUPY_AVAILABLE = True
+
+logger = logging.getLogger(__name__)
+
+KEY_PIXEL_SIZE_1 = "PSize_1"
+KEY_PIXEL_SIZE_2 = "PSize_2"
+KEY_BINNING_1 = "BSize_1"
+KEY_BINNING_2 = "BSize_2"
+KEY_WAVELENGTH = "WaveLength"
+KEY_SAMPLEDETECTOR_DISTANCE = "SampleDistance"
+KEY_NORMALIZATION_FACTOR = "NormalizationFactor"
+KEY_DETECTOR_MASK_FOLDER = "DetectorMaskFilePath"
+KEY_DETECTOR_MASK_FILE = "DetectorMaskFileName"
+KEY_BEAMSTOP_MASK_FOLDER = "MaskFilePath"
+KEY_BEAMSTOP_MASK_FILE = "MaskFileName"
+KEY_POLARIZATION_FACTOR = "polarization_factor"
+KEY_POLARIZATION_AXIS = "polarization_axis_offset"
+KEY_VARIANCE_FORMULA = "variance_formula"
+KEY_WINDOW_ROI_SIZE = "WindowRoiSize"
+KEY_DARK_FOLDER = "DarkFilePath"
+KEY_DARK_FILE = "DarkFileName"
+KEY_FLAT_FOLDER = "FlatfieldFilePath"
+KEY_FLAT_FILE = "FlatfieldFileName"
+KEY_WINDOW_FOLDER = "WindowFilePath"
+KEY_WINDOW_FILE = "WindowFileName"
+KEY_DUMMY = "Dummy"
+KEY_DELTA_DUMMY = "DDummy"
+KEY_CENTER_1 = "Center_1"
+KEY_CENTER_2 = "Center_2"
+KEY_NPT2_RAD = "npt2_rad"
+KEY_NPT2_AZIM = "npt2_azim"
+KEY_UNIT = "unit"
+KEY_TITLEEXTENSION = "TitleExtension"
+KEY_ALGORITHM_NORMALIZATION = "NormAlgorithm"
+
+
+def get_isotime(forceTime: Optional[float] = None) -> str:
+    """
+    Get the current time as an ISO8601 string.
+
+    Inputs:
+        - forceTime (Optional[float], optional): Enforce a given time (current by default). Defaults to None.
+    Outputs:
+        - str: The current time as an ISO8601 string.
+    """
+    if forceTime is None:
+        forceTime = time.time()
+    localtime = time.localtime(forceTime)
+    gmtime = time.gmtime(forceTime)
+    tz_h = localtime.tm_hour - gmtime.tm_hour
+    tz_m = localtime.tm_min - gmtime.tm_min
+    return "%s%+03i:%02i" % (time.strftime("%Y-%m-%dT%H:%M:%S", localtime), tz_h, tz_m)
+
+
+def refactor_stream_name_raw(stream_name: str, cut_name: bool = False):
+    if "roi_counters" in stream_name:
+        lst = stream_name.split(":")
+        return f"{lst[0]}_{lst[-1]}"
+    if not cut_name:
+        return "_".join(re.split(":|_", stream_name))
+    stream_name = stream_name.split(":")[-1]
+    return "_".join(re.split(":|_", stream_name))
+
+
+def refactor_stream_name_interpreted(stream_name: str):
+    return "_".join(re.split(":|_", stream_name))
+
+
+def parse_titleextension_template(template: str):
+    # First, find the curly brackets
+    pattern = r"\{(.*?)\}"
+    template_parsed = template
+    template_info = []
+    for template_element in re.findall(pattern, template):
+        # Clear empty spaces
+        template_element_parsed = template_element.replace(" ", "")
+        try:
+            stream_name, format_spec = template_element_parsed.split(":")
+            template_info.append(
+                {"stream_name": stream_name, "format_spec": format_spec}
+            )
+            template_parsed = template_parsed.replace(
+                template_element, f"{stream_name}:{format_spec}"
+            )
+        except Exception:
+            logger.error(
+                f"{template_element} in {template} is not a valid format for TitleExtension"
+            )
+            continue
+    return template_parsed, template_info
+
+
+def get_from_headers(
+    key: str,
+    headers: Optional[Dict[str, Union[str, float]]] = None,
+    metadata_file_group: Optional[h5py.Group] = None,
+    to_integer: bool = False,
+) -> Optional[Union[str, float, int]]:
+    """
+    Retrieve a header value from the header object (for online processing) or from an HDF5 group (for offline processing).
+
+    Inputs:
+        - key (str): The key to retrieve.
+        - headers (Optional[Dict[str, Union[str, float]]], optional): The header object. Defaults to None.
+        - metadata_file_group (Optional[h5py.Group], optional): The HDF5 group. Defaults to None.
+        - to_integer (bool, optional): Whether to convert the value to an integer. Defaults to False.
+    Outputs:
+        - Optional[Union[str, float, int]]: The retrieved value or None if not found.
+    """
+    value = None
+
+    if headers:
+        # Retrieve directly from the header object
+        if key not in headers:
+            logger.warning(f"Key {key} is not in headers")
+            return
+        value = headers[key]
+    elif metadata_file_group:
+        # Retrieve from a group in the metadata file
+        if key not in metadata_file_group:
+            logger.warning(f"Key {key} not in {metadata_file_group}")
+            return
+        value = metadata_file_group[key][()]
+    else:
+        return
+
+    if isinstance(value, bytes):
+        value = value.decode("UTF-8")
+    try:
+        value = float(value)
+        if to_integer:
+            return int(value)
+        return value
+    except Exception:
+        return value
+
+
+def get_value_from_file(filename: str, h5path: str, key: str, to_integer: bool = False):
+    params = {
+        "filename": filename,
+        "name": h5path,
+        "retry_timeout": 0.1,
+    }
+    try:
+        with open_item_silx(**params) as grp:
+            if grp is None:
+                return
+            if key not in grp:
+                return
+            value = grp[key][()]
+    except RetryTimeoutError:
+        return
+
+    if isinstance(value, bytes):
+        value = value.decode("UTF-8")
+    try:
+        value = float(value)
+        if to_integer:
+            return int(value)
+        return value
+    except Exception:
+        return value
+
+
+def get_flat_filename(**kwargs):
+    """Returns the whole filename for the flat field correction from the headers."""
+    flat_folder = get_from_headers(key=KEY_FLAT_FOLDER, **kwargs)
+    flat_file = get_from_headers(key=KEY_FLAT_FILE, **kwargs)
+    if flat_folder is None or flat_file is None:
+        return
+
+    flat_filename = os.path.join(flat_folder, flat_file)
+    if not os.path.exists(flat_filename):
+        return
+    return flat_filename
+
+
+def get_mask_detector_filename(**kwargs):
+    """Returns the whole filename for the detector gaps mask from the headers."""
+    mask_folder = get_from_headers(key=KEY_DETECTOR_MASK_FOLDER, **kwargs)
+    mask_file = get_from_headers(key=KEY_DETECTOR_MASK_FILE, **kwargs)
+    if mask_folder is None or mask_file is None:
+        return
+
+    mask_filename = os.path.join(mask_folder, mask_file)
+    if not os.path.exists(mask_filename):
+        return
+    return mask_filename
+
+
+def get_mask_beamstop_filename(**kwargs):
+    """Returns the whole filename for the beamstop mask from the headers."""
+    mask_folder = get_from_headers(key=KEY_BEAMSTOP_MASK_FOLDER, **kwargs)
+    mask_file = get_from_headers(key=KEY_BEAMSTOP_MASK_FILE, **kwargs)
+    if mask_folder is None or mask_file is None:
+        return
+
+    mask_filename = os.path.join(mask_folder, mask_file)
+    if not os.path.exists(mask_filename):
+        return
+    return mask_filename
+
+
+def get_dark_filename(**kwargs):
+    """Returns the whole filename for the dark current correction from the headers."""
+    dark_folder = get_from_headers(key=KEY_DARK_FOLDER, **kwargs)
+    dark_file = get_from_headers(key=KEY_DARK_FILE, **kwargs)
+    if dark_folder is None or dark_file is None:
+        return
+
+    dark_filename = os.path.join(dark_folder, dark_file)
+    if not os.path.exists(dark_filename):
+        return
+    return dark_filename
+
+
+def load_data(
+    filename: Union[str, List[str]],
+    binning: tuple = (1, 1),
+    data_signal_shape: tuple = None,
+    use_cupy: bool = False,
+    datatype: str = None,
+    dark_filter: str = "median",
+    dark_filter_quantil_lower: int = 0,
+    dark_filter_quantil_upper: int = 1,
+    **kwargs,
+) -> Optional[numpy.ndarray]:
+    """
+    Load data from a file or a list of files.
+
+    Inputs:
+        - filename (Union[str, List[str]]): The filename or list of filenames.
+        - binning (tuple): binning of the data signal
+        - data_signal_shape (tuple): shape of the data array (2-dimensional)
+        - use_cupy (bool): if True, returns a cupy.asarray
+        - datatype (str): format of the imported array, if None, datatype is respected
+    Outputs:
+        - Optional[numpy.ndarray]: The loaded data or None if the file does not exist.
+    """
+    if filename is None:
+        return
+
+    # Import data
+    data = None
+    if isinstance(filename, (tuple, list)):
+        for _, file in enumerate(filename):
+            data_ = _load_data(file)
+            if data_ is None:
+                continue
+            if data is None:
+                data = data_
+            else:
+                data += data_
+    elif isinstance(filename, str):
+        data = _load_data(filename)
+
+    if data is None:
+        return
+
+    # Set datatype
+    if datatype is not None and data.dtype != datatype:
+        data = data.astype(datatype, copy=False)
+
+    # Dataset filter
+    if data.ndim > 2:
+        if dark_filter.startswith("quantil"):
+            data = average_dark(
+                data,
+                center_method=dark_filter,
+                quantiles=(dark_filter_quantil_lower, dark_filter_quantil_upper),
+            )
+        else:
+            data = average_dark(data, center_method=dark_filter)
+
+    # Binning unification
+    if data_signal_shape and data.shape != data_signal_shape:
+        binning_additional_data = _get_data_binning(filename=filename)
+        binning_relative = (
+            int(binning[0] / binning_additional_data[0]),
+            int(binning[1] / binning_additional_data[1]),
+        )
+        data_binned = binning_tool(data, binning_relative, norm=False)
+        if data_binned.shape != data_signal_shape:
+            raise ValueError(
+                f"Data shape after binning {binning} from {filename} does not match the expected shape: {data_binned.shape} != {data_signal_shape}"
+            )
+        data = data_binned
+    elif data_signal_shape is None and binning == (1, 1):
+        ...
+
+    if use_cupy and CUPY_AVAILABLE:
+        return cupy.asarray(data)
+    return data
+
+
+def _load_data(filename: str) -> Optional[numpy.ndarray]:
+    """
+    Load data from a single file.
+
+    Inputs:
+        - filename (str): The filename.
+    Outputs:
+        - Optional[numpy.ndarray]: The loaded data or None if the file does not exist.
+    """
+    if not os.path.exists(filename):
+        return
+
+    if filename.endswith(".h5"):
+        filename += "::/entry_0000/measurement/data"
+    try:
+        with fabio.open(filename) as f:
+            if f.nframes == 1:
+                return f.data
+            else:
+                return f.dataset[:]
+    except Exception as e:
+        logger.error(f"File {filename} could not be open with fabio: {e}")
+
+
+def _get_data_binning(filename: str):
+    """
+    Load data from a single file.
+
+    Inputs:
+        - filename (str): The filename.
+    Outputs:
+        - Optional[numpy.ndarray]: The loaded data or None if the file does not exist.
+    """
+    if not os.path.exists(filename):
+        return
+
+    if filename.endswith(".h5"):
+        filename += "::/entry_0000/measurement/data"
+    try:
+        with fabio.open(filename) as f:
+            b1 = f.header.get("Bsize_1")
+            b1_ = f.header.get("BSize_1")
+            b2 = f.header.get("Bsize_2")
+            b2_ = f.header.get("BSize_2")
+            b1 = b1 or b1_
+            b2 = b2 or b2_
+            return (int(b1), int(b2))
+    except Exception as e:
+        logger.error(f"File {filename} could not be open with fabio: {e}")
+
+
+def get_headers(
+    headers: Optional[Dict[str, Union[str, float]]] = None,
+    metadata_file_group: Optional[h5py.Group] = None,
+) -> Optional[Dict[str, Union[str, float]]]:
+    """
+    Retrieve headers from a dictionary or an HDF5 group.
+
+    Inputs:
+        - headers (Optional[Dict[str, Union[str, float]]], optional): The header dictionary. Defaults to None.
+        - metadata_file_group (Optional[h5py.Group], optional): The HDF5 group. Defaults to None.
+    Outputs:
+        - Optional[Dict[str, Union[str, float]]]: The headers or None if not found.
+    """
+    if headers:
+        return headers
+    elif metadata_file_group:
+        headers = {}
+        for key in metadata_file_group:
+            value = metadata_file_group[key][()]
+            if isinstance(value, bytes):
+                value = value.decode("UTF-8")
+            headers[key] = value
+        return headers
+    else:
+        return {}
+
+
+def get_free_memory(device_id):
+    """Retrieves the available memory on a GPU device"""
+    if not CUPY_AVAILABLE:
+        logger.warning("Cupy is not available.")
+        return None
+
+    with cupy.cuda.Device(device_id):
+        free_mem, total_mem = cupy.cuda.runtime.memGetInfo()
+        return free_mem
+
+
+def get_best_gpu():
+    """Decides the best GPU in terms of memory available"""
+    if not CUPY_AVAILABLE:
+        logger.warning("Cupy is not available.")
+        return None
+
+    best_device = None
+    max_free_memory = 0
+
+    for device_id in range(cupy.cuda.runtime.getDeviceCount()):
+        free_memory = get_free_memory(device_id)
+        if free_memory is not None and free_memory > max_free_memory:
+            max_free_memory = free_memory
+            best_device = device_id
+
+    return best_device
+
+
+def use_best_gpu():
+    """
+    Set the best available GPU for cupy operations.
+    """
+    best_device = get_best_gpu()
+    if best_device is not None:
+        cupy.cuda.Device(best_device).use()
+        logger.info(f"Using GPU {best_device} with the most free memory.")
+    else:
+        logger.warning("No suitable GPU found or cupy is not available.")
+
+
+def serialize_h5py_task(h5py_group: h5py.Group) -> dict:
+    """
+    Recursively convert an h5py Group or File into a nested Python dictionary.
+    Datasets become numpy arrays or scalars.
+    """
+    result = {}
+    for key, item in h5py_group.items():
+        if isinstance(item, h5py.Dataset):
+            data = item[()]
+            if isinstance(data, bytes):
+                data = data.decode()
+            # Convert 0-d arrays to scalars for readability
+            if isinstance(data, numpy.ndarray) and data.shape == ():
+                data = data.item()
+            result[key] = data
+        elif isinstance(item, h5py.Group):
+            result[key] = serialize_h5py_task(item)
+        else:
+            result[key] = str(item)
+    return result
+
+
+def deserialize_h5py_task(h5dict: dict, h5py_parent: h5py.Group) -> None:
+    for key, value in h5dict.items():
+        if isinstance(value, dict):
+            child_group = h5py_parent.create_group(name=key)
+            child_group.attrs["NX_class"] = "NXcollection"
+            deserialize_h5py_task(h5dict=value, h5py_parent=child_group)
+        else:
+            try:
+                h5py_parent.create_dataset(
+                    name=key,
+                    data=value,
+                )
+            except Exception as e:
+                print(e, key, value)
+                continue
+
+
+def get_persistent_array_mask(
+    filename_mask: str,
+    data_signal_shape: tuple,
+    datatype: str = "bool",
+    binning: tuple = (1, 1),
+    use_cupy: bool = False,
+    **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_array(
+        filename=filename_mask,
+        mtime=mtime_mask,
+        datatype=datatype,
+        data_signal_shape=data_signal_shape,
+        binning=binning,
+        use_cupy=use_cupy,
+    )
+
+
+def get_persistent_array_flat(
+    filename_flat: str,
+    data_signal_shape: tuple,
+    datatype: str = None,
+    binning: tuple = (1, 1),
+    dummy: int = None,
+    delta_dummy: float = None,
+    filename_mask: str = None,
+    use_cupy: bool = False,
+    **kwargs,
+):
+    if filename_flat and os.path.exists(filename_flat):
+        mtime_flat = os.path.getmtime(filename_flat)
+    else:
+        mtime_flat = None
+        filename_flat = None
+
+    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_array_flat(
+        filename_flat=filename_flat,
+        mtime_flat=mtime_flat,
+        data_signal_shape=data_signal_shape,
+        datatype=datatype,
+        binning=binning,
+        dummy=dummy,
+        delta_dummy=delta_dummy,
+        filename_mask=filename_mask,
+        mtime_mask=mtime_mask,
+        use_cupy=use_cupy,
+    )
+
+
+@lru_cache(maxsize=5)
+def _get_persistent_array_flat(
+    filename_flat: str,
+    mtime_flat: float,
+    **kwargs,
+):
+    if not filename_flat:
+        return
+    logger.info(
+        f"No cache hit. Creating new flat-field array from {filename_flat=}, {mtime_flat=}, {kwargs}"
+    )
+
+    array_flat = load_data(
+        filename=filename_flat,
+        **kwargs,
+    )
+    if array_flat is None:
+        return
+
+    array_mask = _get_persistent_array(
+        filename=kwargs.get("filename_mask"),
+        mtime=kwargs.get("mtime_mask"),
+        datatype="bool",
+        data_signal_shape=kwargs.get("data_signal_shape"),
+        binning=kwargs.get("binning"),
+        use_cupy=kwargs.get("use_cupy"),
+    )
+
+    dummy = kwargs.get("dummy")
+    delta_dummy = kwargs.get("delta_dummy")
+    if array_mask is not None and dummy is not None and delta_dummy is not None:
+        if delta_dummy == 0:
+            array_mask |= array_flat == dummy
+        else:
+            array_mask |= abs(array_flat - dummy) < delta_dummy
+    return array_flat
+
+
+def get_persistent_array_dark(
+    filename_dark: str,
+    data_signal_shape: tuple,
+    datatype: str = None,
+    binning: tuple = (1, 1),
+    dummy: int = None,
+    delta_dummy: float = None,
+    filename_mask: str = None,
+    dark_filter: str = None,
+    dark_filter_quantil_lower: float = 0.1,
+    dark_filter_quantil_upper: float = 0.9,
+    use_cupy: bool = False,
+    **kwargs,
+):
+    if filename_dark and os.path.exists(filename_dark):
+        mtime_dark = os.path.getmtime(filename_dark)
+    else:
+        mtime_dark = None
+        filename_dark = None
+
+    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_array_dark(
+        filename_dark=filename_dark,
+        mtime_dark=mtime_dark,
+        data_signal_shape=data_signal_shape,
+        datatype=datatype,
+        binning=binning,
+        dummy=dummy,
+        delta_dummy=delta_dummy,
+        filename_mask=filename_mask,
+        mtime_mask=mtime_mask,
+        dark_filter=dark_filter,
+        dark_filter_quantil_lower=dark_filter_quantil_lower,
+        dark_filter_quantil_upper=dark_filter_quantil_upper,
+        use_cupy=use_cupy,
+    )
+
+
+@lru_cache(maxsize=5)
+def _get_persistent_array_dark(
+    filename_dark: str,
+    mtime_dark: float,
+    **kwargs,
+):
+    if not filename_dark:
+        return
+    logger.info(
+        f"No cache hit. Creating new dark-current array from {filename_dark=}, {mtime_dark}, {kwargs}"
+    )
+
+    array_dark = load_data(
+        filename=filename_dark,
+        **kwargs,
+    )
+    if array_dark is None:
+        return
+
+    array_mask = _get_persistent_array(
+        filename=kwargs.get("filename_mask"),
+        mtime=kwargs.get("mtime_mask"),
+        datatype="bool",
+        data_signal_shape=kwargs.get("data_signal_shape"),
+        binning=kwargs.get("binning"),
+        use_cupy=kwargs.get("use_cupy"),
+    )
+
+    dummy = kwargs.get("dummy")
+    delta_dummy = kwargs.get("delta_dummy")
+    if array_mask is not None and dummy is not None and delta_dummy is not None:
+        if delta_dummy == 0:
+            array_mask |= array_dark == dummy
+        else:
+            array_mask |= abs(array_dark - dummy) < delta_dummy
+    return array_dark
+
+
+def get_persistent_array_window_wagon(
+    filename_window_wagon: str,
+    data_signal_shape: tuple,
+    datatype: str = None,
+    binning: tuple = (1, 1),
+    use_cupy: bool = False,
+):
+    if filename_window_wagon and os.path.exists(filename_window_wagon):
+        window_mtime = os.path.getmtime(filename_window_wagon)
+    else:
+        window_mtime = None
+        filename_window_wagon = None
+
+    return _get_persistent_array(
+        filename=filename_window_wagon,
+        mtime=window_mtime,
+        data_signal_shape=data_signal_shape,
+        datatype=datatype,
+        binning=binning,
+        use_cupy=use_cupy,
+    )
+
+
+def get_dataset_signal_from_processed_file(
+    filename: str,
+    index_range: tuple = None,
+):
+    if filename and os.path.exists(filename):
+        with h5py.File(filename, "r") as f:
+            nxprocess_path = "entry_0000/PyFAI"
+
+            if nxprocess_path not in f:
+                logger.error(f"{filename} does not contain the {nxprocess_path} format")
+                return
+            nxprocess_grp = f[nxprocess_path]
+
+            nxdata_name = next(
+                (name for name in nxprocess_grp if "result_" in name), None
+            )
+            if nxdata_name is None:
+                logger.error(f"There is no result_xxx group in {nxprocess_path}")
+                return
+
+            nxdata_grp = nxprocess_grp[nxdata_name]
+            if "data" not in nxdata_grp:
+                logger.error(
+                    f"There is no data dataset in {nxprocess_path}/{nxdata_name}"
+                )
+                return
+            if index_range is not None:
+                return nxdata_grp["data"][index_range[0] : index_range[-1]]
+            return nxdata_grp["data"][:]
+
+
+@lru_cache(maxsize=10)
+def _get_persistent_array(
+    filename: str,
+    mtime: float,
+    **kwargs,
+):
+    if not filename:
+        return
+    logger.info(
+        f"No cache hit. Creating new array from {filename=}, {mtime=}, {kwargs}"
+    )
+    return load_data(
+        filename=filename,
+        **kwargs,
+    )