Function prepare_xrf_map

pyxrf/core/map_processing.py:331–407

Convert XRF map from it's initial representation to properly chunked Dask array. Parameters ---------- data: da.core.Array, np.ndarray or RawHDF5Dataset (this is a custom type) Raw XRF map represented as Dask array, numpy array or reference to a dataset in HDF5 file

(data, chunk_pixels=5000, n_chunks_min=4)

def prepare_xrf_map(data, chunk_pixels=5000, n_chunks_min=4):
    """
    Convert XRF map from it's initial representation to properly chunked Dask array.

    Parameters
    ----------
    data: da.core.Array, np.ndarray or RawHDF5Dataset (this is a custom type)
        Raw XRF map represented as Dask array, numpy array or reference to a dataset in
        HDF5 file. The XRF map must have dimensions `(ny, nx, ne)`, where `ny` and `nx`
        define image size and `ne` is the number of spectrum points
    chunk_pixels: int
        The number of pixels in a single chunk. The XRF map will be rechunked so that
        each block contains approximately `chunk_pixels` pixels and contain all `ne`
        spectrum points for each pixel.
    n_chunks_min: int
        Minimum number of chunks. The algorithm will try to split the map into the number
        of chunks equal or greater than `n_chunks_min`. If HDF5 dataset is not chunked,
        then the whole map is treated as one chunk. This should happen only to very small
        files, so parallelism is not important.

    Returns
    -------
    data: da.core.Array
        XRF map represented as Dask array with proper chunk size. The XRF map may be loaded
        block by block when processing using `dask.array.map_blocks` and `dask.array.blockwise`
        functions with Dask multiprocessing scheduler.
    file_obj: h5py.File object
        File object that points to HDF5 file. `None` if input parameter `data` is Dask or
        numpy array. Note, that `file_obj` must be kept alive until processing is completed.
        Closing the file will invalidate references to the dataset in the respective
        Dask array.

    Raises
    ------
    TypeError if input parameter `data` is not one of supported types.
    """

    file_obj = None  # It will remain None, unless 'data' is 'RawHDF5Dataset'

    if isinstance(data, da.core.Array):
        chunk_size = _compute_optimal_chunk_size(
            chunk_pixels=chunk_pixels,
            data_chunksize=data.chunksize[0:2],
            data_shape=data.shape[0:2],
            n_chunks_min=n_chunks_min,
        )
        data = data.rechunk(chunks=(*chunk_size, data.shape[2]))
    elif isinstance(data, np.ndarray):
        data = _array_numpy_to_dask(data, chunk_pixels=chunk_pixels, n_chunks_min=n_chunks_min)
    elif isinstance(data, RawHDF5Dataset):
        fpath, dset_name = data.abs_path, data.dset_name

        # Note, that the file needs to remain open until the processing is complete !!!
        file_obj = h5py.File(fpath, "r")
        dset = file_obj[dset_name]

        if dset.ndim != 3:
            raise TypeError(