nabu 2024.1.0rc3__py3-none-any.whl → 2024.1.1__py3-none-any.whl

build/lib/build/lib/nabu/processing/histogram.py

@@ -0,0 +1,286 @@
+from math import log2, ceil
+import numpy as np
+from silx.math import Histogramnd
+from tomoscan.io import HDF5File
+from ..utils import check_supported
+from ..resources.logger import LoggerOrPrint
+
+
+class PartialHistogram:
+    """
+    A class for computing histogram progressively.
+
+    In certain cases, it is cumbersome to compute a histogram directly on a big chunk of
+    data (ex. data not fitting in memory, disk access too slow) while some parts of the
+    data are readily available in-memory.
+    """
+
+    histogram_methods = ["fixed_bins_width", "fixed_bins_number"]
+    bin_width_policies = ["uint16"]
+    backends = ["numpy", "silx"]
+
+    def __init__(self, method="fixed_bins_width", bin_width="uint16", num_bins=None, min_bins=None, backend="silx"):
+        """
+        Initialize a PartialHistogram class.
+
+        Parameters
+        ----------
+        method: str, optional
+            Partial histogram computing method. Available are:
+               - `fixed_bins_width`: all the histograms are computed with the same bin
+                 width. The class adapts to the data range and computes the number of
+                 bins accordingly.
+               - `fixed_bins_number`: all the histograms are computed with the same
+                 number of bins. The class adapts to the data range and computes the
+                 bin width accordingly.
+            Default is "fixed_bins_width"
+        bin_width: str or float, optional
+            Policy for histogram bins when method="fixed_bins_width". Available are:
+               - "uint16": The bin width is computed so that floating-point elements
+                 `f1` and `f2` satisfying `|f1 - f2| < bin_width` implies
+                 `f1_converted - f2_converted < 1` once cast to uint16.
+               - A number: all the bins have this fixed width.
+
+            Default is "uint16"
+        num_bins: int, optional
+            Number of bins when method = 'fixed_bins_number'.
+        min_bins: int, optional
+            Minimum number of bins when method = 'fixed_bins_width'.
+        backend: str, optional
+            Which histogram backend to use for computations. Available are "silx", "numpy".
+            Fastest is "silx".
+        """
+        check_supported(method, self.histogram_methods, "histogram computing method")
+        self.method = method
+        check_supported(backend, self.backends, "histogram backend")
+        self.backend = backend
+        self._set_bin_width(bin_width)
+        self._set_num_bins(num_bins)
+        self.min_bins = min_bins
+        self._set_histogram_methods()
+
+    def _set_bin_width(self, bin_width):
+        if self.method == "fixed_bins_number":
+            self.bin_width = None
+            return
+        if isinstance(bin_width, str):
+            check_supported(bin_width, self.bin_width_policies, "bin width policy")
+            self._fixed_bw = False
+        else:
+            bin_width = float(bin_width)
+            self._fixed_bw = True
+        self.bin_width = bin_width
+
+    def _set_num_bins(self, num_bins):
+        if self.method == "fixed_bins_width":
+            self.num_bins = None
+            return
+        if self.method == "fixed_bins_number" and num_bins is None:
+            raise ValueError("Need to specify num_bins for method='fixed_bins_number'")
+        self.num_bins = int(num_bins)
+
+    def _set_histogram_methods(self):
+        self._histogram_methods = {
+            "fixed_bins_number": {
+                "compute": self._compute_histogram_fixed_nbins,
+                "merge": self._merge_histograms_fixed_nbins,
+            },
+            "fixed_bins_width": {
+                "compute": self._compute_histogram_fixed_bw,
+                "merge": self._merge_histograms_fixed_bw,
+            },
+        }
+        assert set(self._histogram_methods.keys()) == set(self.histogram_methods)
+
+    @staticmethod
+    def _get_histograms_and_bins(histograms, center=False, dont_truncate_bins=False):
+        histos = [h[0] for h in histograms]
+        if dont_truncate_bins:
+            bins = [h[1] for h in histograms]
+        else:
+            if center:
+                bins = [0.5 * (h[1][1:] + h[1][:-1]) for h in histograms]
+            else:
+                bins = [h[1][:-1] for h in histograms]
+        return histos, bins
+
+    #
+    # Histogram with fixed number of bins
+    #
+
+    def _compute_histogram_fixed_nbins(self, data, data_range=None):
+        if data.ndim > 1:
+            data = data.ravel()
+        dmin, dmax = data.min(), data.max() if data_range is None else data_range
+        if self.backend == "numpy":
+            res = np.histogram(data, bins=self.num_bins)
+        elif self.backend == "silx":
+            histogrammer = Histogramnd(data, n_bins=self.num_bins, histo_range=(dmin, dmax), last_bin_closed=True)
+            res = histogrammer.histo, histogrammer.edges[0]  # pylint: disable=E1136
+        return res
+
+    def _merge_histograms_fixed_nbins(self, histograms, dont_truncate_bins=False):
+        histos, bins = self._get_histograms_and_bins(histograms, dont_truncate_bins=dont_truncate_bins)
+        res = np.histogram(
+            np.hstack(bins),
+            weights=np.hstack(histos),
+            bins=self.num_bins,
+        )
+        return res
+
+    #
+    # Histogram with fixed bin width
+    #
+
+    def _bin_width_u16(self, dmin, dmax):
+        return (dmax - dmin) / 65535.0
+
+    def _bin_width_fixed(self, dmin, dmax):
+        return self.bin_width
+
+    def get_bin_width(self, dmin, dmax):
+        if self._fixed_bw:
+            return self._bin_width_fixed(dmin, dmax)
+        elif self.bin_width == "uint16":
+            return self._bin_width_u16(dmin, dmax)
+        else:
+            raise ValueError()
+
+    def _compute_histogram_fixed_bw(self, data, data_range=None):
+        dmin, dmax = data.min(), data.max() if data_range is None else data_range
+        min_bins = self.min_bins or 1
+        bw_max = self.get_bin_width(dmin, dmax)
+        nbins = 0
+        bw_factor = 1
+        while nbins < min_bins:
+            bw = 2 ** round(log2(bw_max)) / bw_factor
+            nbins = int((dmax - dmin) / bw)
+            bw_factor *= 2
+        res = np.histogram(data, bins=nbins)
+        return res
+
+    def _merge_histograms_fixed_bw(self, histograms, **kwargs):
+        histos, bins = self._get_histograms_and_bins(histograms, center=False)
+        dmax = max([b[-1] for b in bins])
+        dmin = min([b[0] for b in bins])
+        bw_max = max([b[1] - b[0] for b in bins])
+        res = np.histogram(np.hstack(bins), weights=np.hstack(histos), bins=int((dmax - dmin) / bw_max))
+        return res
+
+    #
+    # Dispatch methods
+    #
+
+    def compute_histogram(self, data, data_range=None):
+        compute_hist_func = self._histogram_methods[self.method]["compute"]
+        return compute_hist_func(data, data_range=data_range)
+
+    def merge_histograms(self, histograms, **kwargs):
+        merge_hist_func = self._histogram_methods[self.method]["merge"]
+        return merge_hist_func(histograms, **kwargs)
+
+
+class VolumeHistogram:
+    """
+    A class for computing the histogram of an entire volume.
+    Unless explicitly specified, histogram is computed in several passes so that not
+    all the volume is loaded in memory.
+    """
+
+    def __init__(self, data_url, chunk_size_slices=100, chunk_size_GB=None, nbins=1e6, logger=None):
+        """
+        Initialize a VolumeHistogram object.
+
+        Parameters
+        ----------
+        fname: DataUrl
+            DataUrl to the HDF5 file.
+        chunk_size_slices: int, optional
+            Compute partial histograms of groups of slices. This is the default behavior,
+            where the groups size is 100 slices.
+            This parameter is mutually exclusive with 'chunk_size_GB'.
+        chunk_size_GB: float, optional
+            Maximum memory (in GB) to use when computing the histogram by group of slices.
+            This parameter is mutually exclusive with 'chunk_size_slices'.
+        nbins: int, optional
+            Histogram number of bins. Default is 1e6.
+        """
+        self.data_url = data_url
+        self.logger = LoggerOrPrint(logger)
+        self._get_data_info()
+        self._set_chunk_size(chunk_size_slices, chunk_size_GB)
+        self.nbins = int(nbins)
+        self._init_histogrammer()
+
+    def _get_data_info(self):
+        self.fname = self.data_url.file_path()
+        self.data_path = self.data_url.data_path()
+        with HDF5File(self.fname, "r") as fid:
+            try:
+                data_ptr = fid[self.data_path]
+            except KeyError:
+                msg = str(
+                    "Could not access HDF5 path %s in file %s. Please check that this file \
+                    actually contains a reconstruction and that the HDF5 path is correct"
+                    % (self.data_path, self.fname)
+                )
+                self.logger.fatal(msg)
+                raise ValueError(msg)
+            if data_ptr.ndim != 3:
+                msg = "Expected data to have 3 dimensions, got %d" % data_ptr.ndim
+                raise ValueError(msg)
+            self.data_shape = data_ptr.shape
+            self.data_dtype = data_ptr.dtype
+            self.data_nbytes_GB = np.prod(data_ptr.shape) * data_ptr.dtype.itemsize / 1e9
+
+    def _set_chunk_size(self, chunk_size_slices, chunk_size_GB):
+        if not ((chunk_size_slices is not None) ^ (chunk_size_GB is not None)):
+            raise ValueError("Please specify either chunk_size_slices or chunk_size_GB")
+        if chunk_size_slices is None:
+            chunk_size_slices = int(chunk_size_GB / (np.prod(self.data_shape[1:]) * self.data_dtype.itemsize / 1e9))
+        self.chunk_size = chunk_size_slices
+        self.logger.debug("Computing histograms by groups of %d slices" % self.chunk_size)
+
+    def _init_histogrammer(self):
+        self.histogrammer = PartialHistogram(method="fixed_bins_number", num_bins=self.nbins)
+
+    def _compute_histogram(self, data):
+        return self.histogrammer.compute_histogram(data.ravel())  # 1D
+
+    def compute_volume_histogram(self):
+        n_z = self.data_shape[0]
+        histograms = []
+        n_steps = ceil(n_z / self.chunk_size)
+        with HDF5File(self.fname, "r") as fid:
+            for chunk_id in range(n_steps):
+                self.logger.debug("Computing histogram %d/%d" % (chunk_id + 1, n_steps))
+                z_slice = slice(chunk_id * self.chunk_size, (chunk_id + 1) * self.chunk_size)
+                images_stack = fid[self.data_path][z_slice, :, :]
+                hist = self._compute_histogram(images_stack)
+                histograms.append(hist)
+        res = self.histogrammer.merge_histograms(histograms)
+        return res
+
+
+def hist_as_2Darray(hist, center=True, dtype="f"):
+    hist, bins = hist
+    if bins.size != hist.size:
+        # assert bins.size == hist.size +1
+        if center:
+            bins = 0.5 * (bins[1:] + bins[:-1])
+        else:
+            bins = bins[:-1]
+    res = np.zeros((2, hist.size), dtype=dtype)
+    res[0] = hist
+    res[1] = bins.astype(dtype)
+    return res
+
+
+def add_last_bin(histo_bins):
+    """
+    Add the last bin (max value) to a list of bin edges.
+    """
+    res = np.zeros(histo_bins.size + 1, dtype=histo_bins.dtype)
+    res[:-1] = histo_bins[:]
+    res[-1] = res[-2] + (res[1] - res[0])
+    return res

build/lib/build/lib/nabu/processing/histogram_cuda.py

@@ -0,0 +1,103 @@
+import numpy as np
+from ..utils import get_cuda_srcfile, updiv
+from ..cuda.utils import __has_pycuda__
+from .histogram import PartialHistogram, VolumeHistogram
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+    from ..cuda.processing import CudaProcessing
+
+
+class CudaPartialHistogram(PartialHistogram):
+    def __init__(
+        self,
+        method="fixed_bins_number",
+        bin_width="uint16",
+        num_bins=None,
+        min_bins=None,
+        cuda_options=None,
+    ):
+        if method == "fixed_bins_width":
+            raise NotImplementedError("Histogram with fixed bins width is not implemented with the Cuda backend")
+        super().__init__(
+            method=method,
+            bin_width=bin_width,
+            num_bins=num_bins,
+            min_bins=min_bins,
+        )
+        self.cuda_processing = CudaProcessing(**(cuda_options or {}))
+        self._init_cuda_histogram()
+
+    def _init_cuda_histogram(self):
+        self.cuda_hist = self.cuda_processing.kernel(
+            "histogram",
+            filename=get_cuda_srcfile("histogram.cu"),
+            signature="PiiiffPi",
+        )
+        self.d_hist = self.cuda_processing.allocate_array("d_hist", self.num_bins, dtype=np.uint32)
+
+    def _compute_histogram_fixed_nbins(self, data, data_range=None):
+        if isinstance(data, np.ndarray):
+            data = self.cuda_processing.to_device("data", data)
+        if data_range is None:
+            # Should be possible to do both in one single pass with ReductionKernel
+            # and garray.vec.float2, but the last step in volatile shared memory
+            # still gives errors. To be investigated...
+            data_min = garray.min(data).get()[()]
+            data_max = garray.max(data).get()[()]
+        else:
+            data_min, data_max = data_range
+        Nz, Ny, Nx = data.shape
+        block = (16, 16, 4)
+        grid = (
+            updiv(Nx, block[0]),
+            updiv(Ny, block[1]),
+            updiv(Nz, block[2]),
+        )
+        self.d_hist.fill(0)
+        self.cuda_hist(
+            data,
+            Nx,
+            Ny,
+            Nz,
+            data_min,
+            data_max,
+            self.d_hist,
+            self.num_bins,
+            grid=grid,
+            block=block,
+        )
+        # Return a result in the same format as numpy.histogram
+        res_hist = self.d_hist.get()
+        res_bins = np.linspace(data_min, data_max, num=self.num_bins + 1, endpoint=True)
+        return res_hist, res_bins
+
+
+class CudaVolumeHistogram(VolumeHistogram):
+    def __init__(
+        self,
+        data_url,
+        chunk_size_slices=100,
+        chunk_size_GB=None,
+        nbins=1e6,
+        logger=None,
+        cuda_options=None,
+    ):
+        self.cuda_options = cuda_options
+        super().__init__(
+            data_url,
+            chunk_size_slices=chunk_size_slices,
+            chunk_size_GB=chunk_size_GB,
+            nbins=nbins,
+            logger=logger,
+        )
+
+    def _init_histogrammer(self):
+        self.histogrammer = CudaPartialHistogram(
+            method="fixed_bins_number",
+            num_bins=self.nbins,
+            cuda_options=self.cuda_options,
+        )
+
+    def _compute_histogram(self, data):
+        return self.histogrammer.compute_histogram(data)  # 3D

build/lib/build/lib/nabu/processing/kernel_base.py

@@ -0,0 +1,126 @@
+"""
+Base class for CudaKernel and OpenCLKernel
+Should not be used directly
+"""
+from ..utils import updiv
+
+
+class KernelBase:
+    """
+    A base class for OpenCL and Cuda kernels.
+
+    Parameters
+    -----------
+    kernel_name: str
+        Name of the CUDA kernel.
+    filename: str, optional
+        Path to the file name containing kernels definitions
+    src: str, optional
+        Source code of kernels definitions
+    automation_params: dict, optional
+        Automation parameters, see below
+
+    Automation parameters
+    ----------------------
+    automation_params is a dictionary with the following keys and default values.
+        guess_block: bool (True)
+            If block is not specified during calls, choose a block size based on
+            the size/dimensions of the first array.
+            Mind that it is unlikely to be the optimal choice.
+        guess_grid: bool (True):
+            If the grid size is not specified during calls, choose a grid size
+            based on the size of the first array.
+        follow_device_ptr: bool (True)
+            specify gpuarray.gpudata for all cuda GPUArrays (and pyopencl.array.data for pyopencl arrays).
+            Otherwise, raise an error.
+    """
+
+    _default_automation_params = {
+        "guess_block": True,
+        "guess_grid": True,
+        "follow_device_ptr": True,
+    }
+
+    def __init__(
+        self,
+        kernel_name,
+        filename=None,
+        src=None,
+        automation_params=None,
+    ):
+        self.check_filename_src(filename, src)
+        self.set_automation_params(automation_params)
+
+    def check_filename_src(self, filename, src):
+        err_msg = "Please provide either filename or src"
+        if filename is None and src is None:
+            raise ValueError(err_msg)
+        if filename is not None and src is not None:
+            raise ValueError(err_msg)
+        if filename is not None:
+            with open(filename) as fid:
+                src = fid.read()
+        self.filename = filename
+        self.src = src
+
+    def set_automation_params(self, automation_params):
+        self.automation_params = self._default_automation_params.copy()
+        self.automation_params.update(automation_params or {})
+
+    @staticmethod
+    def guess_grid_size(shape, block_size):
+        # python: (z, y, x) -> cuda: (x, y, z)
+        res = tuple(map(lambda x: updiv(x[0], x[1]), zip(shape[::-1], block_size)))
+        if len(res) == 2:
+            res += (1,)
+        return res
+
+    @staticmethod
+    def guess_block_size(shape):
+        """
+        Guess a block size based on the shape of an array.
+        """
+        ndim = len(shape)
+        if ndim == 1:
+            return (128, 1, 1)
+        if ndim == 2:
+            return (32, 32, 1)
+        else:
+            return (16, 8, 8)
+
+    def get_block_grid(self, *args, **kwargs):
+        block = None
+        grid = None
+        if ("block" not in kwargs) or (kwargs["block"] is None):
+            if self.automation_params["guess_block"]:
+                block = self.guess_block_size(args[0].shape)
+            else:
+                raise ValueError("Please provide block size")
+        else:
+            block = kwargs["block"]
+        if ("grid" not in kwargs) or (kwargs["grid"] is None):
+            if self.automation_params["guess_grid"]:
+                grid = self.guess_grid_size(args[0].shape, block)
+            else:
+                raise ValueError("Please provide block grid")
+        else:
+            grid = kwargs["grid"]
+        self.last_block_size = block
+        self.last_grid_size = grid
+        return block, grid
+
+    def follow_device_arr(self, args):
+        raise ValueError("Base class")
+
+    def _prepare_call(self, *args, **kwargs):
+        block, grid = self.get_block_grid(*args, **kwargs)
+        # pycuda crashes when any element of block/grid is not a python int (ex. numpy.int64).
+        # A weird behavior once observed is "data.shape" returning (np.int64, int, int) (!).
+        # Ensure that everything is a python integer.
+        grid = tuple(int(x) for x in grid)
+        if block is not None:
+            block = tuple(int(x) for x in block)
+        #
+        args = self.follow_device_arr(args)
+
+        return grid, block, args, kwargs

build/lib/build/lib/nabu/processing/medfilt_cuda.py

@@ -0,0 +1,159 @@
+from os.path import dirname
+import numpy as np
+from pycuda.compiler import SourceModule
+from ..utils import updiv, get_cuda_srcfile
+from ..cuda.processing import CudaProcessing
+
+
+class MedianFilter:
+    """
+    A class for performing median filter on GPU with CUDA
+    """
+
+    def __init__(
+        self,
+        shape,
+        footprint=(3, 3),
+        mode="reflect",
+        threshold=None,
+        cuda_options=None,
+        abs_diff=False,
+    ):
+        """Constructor of Cuda Median Filter.
+
+        Parameters
+        ----------
+        shape: tuple
+            Shape of the array, in the format (n_rows, n_columns)
+        footprint: tuple
+            Size of the median filter, in the format (y, x).
+        mode: str
+            Boundary handling mode. Available modes are:
+               - "reflect": cba|abcd|dcb
+               - "nearest": aaa|abcd|ddd
+               - "wrap": bcd|abcd|abc
+               - "constant": 000|abcd|000
+
+            Default is "reflect".
+        threshold: float, optional
+            Threshold for the "thresholded median filter".
+            A thresholded median filter only replaces a pixel value by the median
+            if this pixel value is greater or equal than median + threshold.
+        abs_diff: bool, optional
+            Whether to perform conditional threshold as abs(value - median)
+
+        Notes
+        ------
+        Please refer to the documentation of the CudaProcessing class for
+        the other parameters.
+        """
+        self.cuda_processing = CudaProcessing(**(cuda_options or {}))
+        self._set_params(shape, footprint, mode, threshold, abs_diff)
+        self.cuda_processing.init_arrays_to_none(["d_input", "d_output"])
+        self._init_kernels()
+
+    def _set_params(self, shape, footprint, mode, threshold, abs_diff):
+        self.data_ndim = len(shape)
+        if self.data_ndim == 2:
+            ny, nx = shape
+            nz = 1
+        elif self.data_ndim == 3:
+            nz, ny, nx = shape
+        else:
+            raise ValueError("Expected 2D or 3D data")
+        self.shape = shape
+        self.Nx = np.int32(nx)
+        self.Ny = np.int32(ny)
+        self.Nz = np.int32(nz)
+        if len(footprint) != 2:
+            raise ValueError("3D median filter is not implemented yet")
+        if not ((footprint[0] & 1) and (footprint[1] & 1)):
+            raise ValueError("Must have odd-sized footprint")
+        self.footprint = footprint
+        self._set_boundary_mode(mode)
+        self.do_threshold = False
+        self.abs_diff = abs_diff
+        if threshold is not None:
+            self.threshold = np.float32(threshold)
+            self.do_threshold = True
+        else:
+            self.threshold = np.float32(0)
+
+    def _set_boundary_mode(self, mode):
+        self.mode = mode
+        # Some code duplication from convolution
+        self._c_modes_mapping = {
+            "periodic": 2,
+            "wrap": 2,
+            "nearest": 1,
+            "replicate": 1,
+            "reflect": 0,
+            "constant": 3,
+        }
+        mp = self._c_modes_mapping
+        if self.mode.lower() not in mp:
+            raise ValueError(
+                """
+                Mode %s is not available. Available modes are:
+                %s
+                """
+                % (self.mode, str(mp.keys()))
+            )
+        if self.mode.lower() == "constant":
+            raise NotImplementedError("mode='constant' is not implemented yet")
+        self._c_conv_mode = mp[self.mode]
+
+    def _init_kernels(self):
+        # Compile source module
+        compile_options = [
+            "-DUSED_CONV_MODE=%d" % self._c_conv_mode,
+            "-DMEDFILT_X=%d" % self.footprint[1],
+            "-DMEDFILT_Y=%d" % self.footprint[0],
+            "-DDO_THRESHOLD=%d" % (int(self.do_threshold) + int(self.abs_diff)),
+        ]
+        fname = get_cuda_srcfile("medfilt.cu")
+        nabu_cuda_dir = dirname(fname)
+        include_dirs = [nabu_cuda_dir]
+        self.sourcemodule_kwargs = {}
+        self.sourcemodule_kwargs["options"] = compile_options
+        self.sourcemodule_kwargs["include_dirs"] = include_dirs
+        with open(fname) as fid:
+            cuda_src = fid.read()
+        self._module = SourceModule(cuda_src, **self.sourcemodule_kwargs)
+        self.cuda_kernel_2d = self._module.get_function("medfilt2d")
+        # Blocks, grid
+        self._block_size = {2: (32, 32, 1), 3: (16, 8, 8)}[self.data_ndim]  # TODO tune
+        self._n_blocks = tuple([updiv(a, b) for a, b in zip(self.shape[::-1], self._block_size)])
+
+    def medfilt2(self, image, output=None):
+        """
+        Perform a median filter on an image (or batch of images).
+
+        Parameters
+        -----------
+        images: numpy.ndarray or pycuda.gpuarray
+            2D image or 3D stack of 2D images
+        output: numpy.ndarray or pycuda.gpuarray, optional
+            Output of filtering. If provided, it must have the same shape
+            as the input array.
+        """
+        self.cuda_processing.set_array("d_input", image)
+        if output is not None:
+            self.cuda_processing.set_array("d_output", output)
+        else:
+            self.cuda_processing.allocate_array("d_output", self.shape)
+        self.cuda_kernel_2d(
+            self.cuda_processing.d_input,
+            self.cuda_processing.d_output,
+            self.Nx,
+            self.Ny,
+            self.Nz,
+            self.threshold,
+            grid=self._n_blocks,
+            block=self._block_size,
+        )
+        self.cuda_processing.recover_arrays_references(["d_input", "d_output"])
+        if output is None:
+            return self.cuda_processing.d_output.get()
+        else:
+            return output