nabu 2024.1.9__py3-none-any.whl → 2024.2.0__py3-none-any.whl

nabu/processing/azim.py

@@ -0,0 +1,206 @@
+from multiprocessing.pool import ThreadPool
+import numpy as np
+
+try:
+    from skimage.transform import warp_polar
+
+    __have_skimage__ = True
+except ImportError:
+    __have_skimage__ = False
+
+
+def azimuthal_integration(img, axes=(-2, -1), domain="direct"):
+    """
+    Computes azimuthal integration of an image or a stack of images.
+
+    Parameters
+    ----------
+    img : `numpy.array_like`
+        The image or stack of images.
+    axes : tuple(int, int), optional
+        Axes of that need to be azimuthally integrated. The default is (-2, -1).
+    domain : string, optional
+        Domain of the integration. Options are: "direct" | "fourier". Default is "direct".
+
+    Raises
+    ------
+    ValueError
+        Error returned when not passing images or wrong axes.
+    NotImplementedError
+        In case of tack of images for the moment.
+
+    Returns
+    -------
+    `numpy.array_like`
+        The azimuthally integrated profile.
+    """
+    if not len(img.shape) >= 2:
+        raise ValueError("Input image should be at least 2-dimensional.")
+    if not len(axes) == 2:
+        raise ValueError("Input axes should be 2.")
+
+    img_axes_dims = np.array((img.shape[axes[0]], img.shape[axes[1]]))
+    if domain.lower() == "direct":
+        half_dims = (img_axes_dims - 1) / 2
+        xx = np.linspace(-half_dims[0], half_dims[0], img_axes_dims[0])
+        yy = np.linspace(-half_dims[1], half_dims[1], img_axes_dims[1])
+    else:
+        xx = np.fft.fftfreq(img_axes_dims[0], 1 / img_axes_dims[0])
+        yy = np.fft.fftfreq(img_axes_dims[1], 1 / img_axes_dims[1])
+    xy = np.stack(np.meshgrid(xx, yy, indexing="ij"))
+    r = np.sqrt(np.sum(xy**2, axis=0))
+
+    img_tr_op = [*range(len(img.shape))]
+    for a in axes:
+        img_tr_op.append(img_tr_op.pop(a))
+    img = np.transpose(img, img_tr_op)
+    if len(img.shape) > 2:
+        img_old_shape = img.shape[:-2]
+        img = np.reshape(img, [-1, *img_axes_dims])
+
+    r_l = np.floor(r)
+    r_u = r_l + 1
+    w_l = (r_u - r) * img
+    w_u = (r - r_l) * img
+
+    r_all = np.concatenate((r_l.flatten(), r_u.flatten())).astype(np.int64)
+    if len(img.shape) == 2:
+        w_all = np.concatenate((w_l.flatten(), w_u.flatten()))
+        return np.bincount(r_all, weights=w_all)
+    else:
+        num_imgs = img.shape[0]
+        az_img = [None] * num_imgs
+        for ii in range(num_imgs):
+            w_all = np.concatenate((w_l[ii, :].flatten(), w_u[ii, :].flatten()))
+            az_img[ii] = np.bincount(r_all, weights=w_all)
+        az_img = np.array(az_img)
+        return np.reshape(az_img, (*img_old_shape, az_img.shape[-1]))
+
+
+def do_radial_distribution(ip, X0, Y0, mR, nBins=None, use_calibration=False, cal=None, return_radii=False):
+    """
+    Translates the Java method `doRadialDistribution` (from imagej) into Python using NumPy.
+    Done by chatgpt-4o on 2024-11-08
+
+    Args:
+    - ip: A 2D numpy array representing the image.
+    - X0, Y0: Coordinates of the center.
+    - mR: Maximum radius.
+    - nBins: Number of bins (optional, defaults to 3*mR/4).
+    - use_calibration: Boolean indicating if calibration should be applied.
+    - cal: Calibration object with attributes `pixel_width` and `units` (optional).
+    """
+    if nBins is None:
+        nBins = int(3 * mR / 4)
+
+    Accumulator = np.zeros((2, nBins))
+
+    # Define the bounding box
+    xmin, xmax = X0 - mR, X0 + mR
+    ymin, ymax = Y0 - mR, Y0 + mR
+
+    # Create grid of coordinates
+    x = np.arange(xmin, xmax)
+    y = np.arange(ymin, ymax)
+    xv, yv = np.meshgrid(x, y, indexing="ij")
+
+    # Calculate the radius for each point
+    R = np.sqrt((xv - X0) ** 2 + (yv - Y0) ** 2)
+
+    # Bin calculation
+    bins = np.floor((R / mR) * nBins).astype(int)
+    bins = np.clip(bins - 1, 0, nBins - 1)  # Adjust bins to be in range [0, nBins-1]
+
+    # Accumulate values
+    for b in range(nBins):
+        mask = bins == b
+        Accumulator[0, b] = np.sum(mask)
+        Accumulator[1, b] = np.sum(ip[mask])
+
+    # Normalize integrated intensity
+    Accumulator[1] /= Accumulator[0]
+
+    if use_calibration and cal is not None:
+        # Apply calibration if units are provided
+        radii = cal.pixel_width * mR * (np.arange(1, nBins + 1) / nBins)
+        units = cal.units
+    else:
+        # Use pixel units
+        radii = mR * (np.arange(1, nBins + 1) / nBins)
+        units = "pixels"
+
+    if return_radii:
+        return radii, Accumulator[1]
+    else:
+        return Accumulator[1]
+
+
+# OK-ish, but small discrepancy with do_radial_distribution.
+# 20-40X faster than above methods for (2048, 2048) images
+# Also it assumes a uniform sampling
+# No idea why there is this "offset=1", to be investigated - perhaps radius=0 is also calculated ?
+def azimuthal_integration_skimage(img, center=None, offset=1):
+    shape2 = [int(s // 2 * 1.4142) for s in img.shape]
+    s = min(img.shape) // 2
+    img_polar = warp_polar(img, output_shape=shape2, center=center)
+    return img_polar.mean(axis=0)[offset : offset + s]
+
+
+def _apply_on_images_stack(func, images_stack, n_threads=4, func_args=None, func_kwargs=None):
+    func_args = func_args or []
+    func_kwargs = func_kwargs or {}
+
+    def _process_image(img):
+        return func(img, *func_args, **func_kwargs)
+
+    with ThreadPool(n_threads) as tp:
+        res = tp.map(_process_image, images_stack)
+    return np.array(res)
+
+
+def _apply_on_patches_stack(func, images_stack, n_threads=4, func_args=None, func_kwargs=None):
+    (n_images, n_patchs_y, img_shape_y, n_patchs_x, img_shape_x) = images_stack.shape
+    func_args = func_args or []
+    func_kwargs = func_kwargs or {}
+    out_sample = func(images_stack[0, 0, :, 0, :], *func_args, **func_kwargs)
+    out_shape = out_sample.shape
+    out_dtype = out_sample.dtype
+
+    def _process_image(img):
+        res = np.zeros((n_patchs_y, n_patchs_x) + out_shape, dtype=out_dtype)
+        for i in range(n_patchs_y):
+            for j in range(n_patchs_x):
+                res[i, j] = func(img[i, :, j, :], *func_args, **func_kwargs)
+        return res
+
+    with ThreadPool(n_threads) as tp:
+        res = tp.map(_process_image, images_stack)
+    return np.array(res)
+
+
+def azimuthal_integration_imagej_stack(images_stack, n_threads=4):
+    if images_stack.ndim == 3:
+        img_shape = images_stack.shape[-2:]
+        _apply = _apply_on_images_stack
+    elif images_stack.ndim == 5:
+        img_shape = np.array(images_stack.shape)[[-3, -1]]
+        _apply = _apply_on_patches_stack
+    else:
+        raise ValueError
+    s = min(img_shape)
+    return _apply(
+        do_radial_distribution,
+        images_stack,
+        n_threads=n_threads,
+        func_args=[s // 2, s // 2, s // 2],
+        func_kwargs={"nBins": s // 2, "return_radii": False},
+    )
+
+
+def azimuthal_integration_skimage_stack(images_stack, n_threads=4):
+    if images_stack.ndim == 3:
+        return _apply_on_images_stack(azimuthal_integration_skimage, images_stack, n_threads=n_threads)
+    elif images_stack.ndim == 5:
+        return _apply_on_patches_stack(azimuthal_integration_skimage, images_stack, n_threads=n_threads)
+    else:
+        raise ValueError

nabu/processing/convolution_cuda.py

@@ -197,7 +197,7 @@ class Convolution:
         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._module = SourceModule(cuda_src, **self.sourcemodule_kwargs)  # pylint: disable=E0606
         # Blocks, grid
         self._block_size = {1: (32, 1, 1), 2: (32, 32, 1), 3: (16, 8, 8)}[self.data_ndim]  # TODO tune
         self._n_blocks = tuple([int(updiv(a, b)) for a, b in zip(self.shape[::-1], self._block_size)])

nabu/processing/fft_cuda.py

@@ -169,7 +169,7 @@ def _has_vkfft(x):
             return False
         vk = VKCUFFT((16,), "f")
         avail = True
-    except (RuntimeError, OSError):
+    except (ImportError, RuntimeError, OSError, NameError):
         avail = False
     return avail
 
@@ -197,7 +197,7 @@ def _has_skfft(x):
 
         sk = SKCUFFT((16,), "f")
         avail = True
-    except (ImportError, RuntimeError, OSError):
+    except (ImportError, RuntimeError, OSError, NameError):
         avail = False
     return avail
 
@@ -216,7 +216,7 @@ def has_skcuda(safe=True):
     return v
 
 
-def get_fft_class(backend="skcuda"):
+def get_fft_class(backend="vkfft"):
     backends = {
         "scikit-cuda": SKCUFFT,
         "skcuda": SKCUFFT,
@@ -226,31 +226,29 @@ def get_fft_class(backend="skcuda"):
         "pyvkfft": VKCUFFT,
     }
 
-    def check_vkfft(asked_fft_cls):
-        if asked_fft_cls is VKCUFFT:
-            if has_vkfft(safe=True) is False:
-                warnings.warn("Could not get VKFFT backend. Falling-back to scikit-cuda/CUFFT instead.", RuntimeWarning)
-                return SKCUFFT
-            return VKCUFFT
-        return SKCUFFT
-
     def get_fft_cls(asked_fft_backend):
         asked_fft_backend = asked_fft_backend.lower()
-        check_supported(asked_fft_backend, list(backends.keys()), "FFT backend name")
-        asked_fft_cls = backends[asked_fft_backend]
-        fft_cls = check_vkfft(asked_fft_cls)
-        return fft_cls
+        check_supported(asked_fft_backend, list(backends.keys()), "Cuda FFT backend name")
+        return backends[asked_fft_backend]
 
     asked_fft_backend_env = os.environ.get("NABU_FFT_BACKEND", "")
     if asked_fft_backend_env != "":
         return get_fft_cls(asked_fft_backend_env)
+
+    avail_fft_implems = get_available_fft_implems()
+    if len(avail_fft_implems) == 0:
+        raise RuntimeError("Could not any Cuda FFT implementation. Please install either scikit-cuda or pyvkfft")
+    if backend not in avail_fft_implems:
+        warnings.warn("Could not get FFT backend '%s'" % backend, RuntimeWarning)
+        backend = avail_fft_implems[0]
+
     return get_fft_cls(backend)
 
 
 def get_available_fft_implems():
     avail_implems = []
-    if has_skcuda(safe=True):
-        avail_implems.append("skcuda")
     if has_vkfft(safe=True):
         avail_implems.append("vkfft")
+    if has_skcuda(safe=True):
+        avail_implems.append("skcuda")
     return avail_implems

nabu/processing/histogram.py

@@ -117,6 +117,8 @@ class PartialHistogram:
         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
+        else:
+            raise ValueError("Unknown backend")
         return res
 
     def _merge_histograms_fixed_nbins(self, histograms, dont_truncate_bins=False):

nabu/processing/histogram_cuda.py

@@ -25,7 +25,7 @@ class CudaPartialHistogram(PartialHistogram):
             num_bins=num_bins,
             min_bins=min_bins,
         )
-        self.cuda_processing = CudaProcessing(**(cuda_options or {}))
+        self.cuda_processing = CudaProcessing(**(cuda_options or {}))  # pylint: disable=E0606
         self._init_cuda_histogram()
 
     def _init_cuda_histogram(self):
@@ -43,6 +43,7 @@ class CudaPartialHistogram(PartialHistogram):
             # 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...
+            # pylint: disable=E0606
             data_min = garray.min(data).get()[()]
             data_max = garray.max(data).get()[()]
         else:

nabu/processing/kernel_base.py

@@ -2,6 +2,7 @@
 Base class for CudaKernel and OpenCLKernel
 Should not be used directly
 """
+
 from ..utils import updiv
 
 
@@ -47,9 +48,11 @@ class KernelBase:
         filename=None,
         src=None,
         automation_params=None,
+        silent_compilation_warnings=False,
     ):
         self.check_filename_src(filename, src)
         self.set_automation_params(automation_params)
+        self.silent_compilation_warnings = silent_compilation_warnings
 
     def check_filename_src(self, filename, src):
         err_msg = "Please provide either filename or src"

nabu/processing/muladd_cuda.py

@@ -42,6 +42,7 @@ class CudaMulAdd(MulAdd):
             raise ValueError("delta_x or delta_y is 0")
 
         # can't use "int4" in pycuda ? int2 seems fine. Go figure
+        # pylint: disable=E0606
         dst_x_range = np.array(dst_coords[:2], dtype=garray.vec.int2)
         dst_y_range = np.array(dst_coords[2:], dtype=garray.vec.int2)
         other_x_range = np.array(other_coords[:2], dtype=garray.vec.int2)

nabu/processing/padding_opencl.py

@@ -27,7 +27,7 @@ class OpenCLPadding(PaddingBase):
             self.d_padded_array_constant = self.processing.to_device(
                 "d_padded_array_constant", self.padded_array_constant
             )
-            self.memcpy2D = OpenCLMemcpy2D(ctx=self.processing.ctx, queue=self.queue)
+            self.memcpy2D = OpenCLMemcpy2D(ctx=self.processing.ctx, queue=self.queue)  # pylint: disable=E0606
             return
         self._coords_transform_kernel = self.processing.kernel(
             "coordinate_transform",

nabu/processing/roll_opencl.py

@@ -1,6 +1,7 @@
 #
 # WIP !
 #
+# pylint: skip-file
 import numpy as np
 from ..opencl.utils import __has_pyopencl__
 from ..utils import get_opencl_srcfile

nabu/processing/rotation_cuda.py

@@ -23,11 +23,11 @@ class CudaRotation(Rotation):
         self._init_rotation_kernel()
 
     def _allocate_arrays(self):
-        self._d_image_cua = cuda.np_to_array(np.zeros(self.shape, "f"), "C")
+        self._d_image_cua = cuda.np_to_array(np.zeros(self.shape, "f"), "C")  # pylint: disable=E0606
         self.cuda_processing.init_arrays_to_none(["d_output"])
 
     def _init_rotation_kernel(self):
-        self.cuda_rotation_kernel = CudaKernel("rotate", get_cuda_srcfile("rotation.cu"))
+        self.cuda_rotation_kernel = CudaKernel("rotate", get_cuda_srcfile("rotation.cu"))  # pylint: disable=E0606
         self.texref_image = self.cuda_rotation_kernel.module.get_texref("tex_image")
         self.texref_image.set_filter_mode(cuda.filter_mode.LINEAR)  # bilinear
         self.texref_image.set_address_mode(0, cuda.address_mode.CLAMP)  # TODO tune

nabu/processing/tests/test_fft.py

@@ -4,17 +4,14 @@ import numpy as np
 from scipy.fft import fftn, ifftn, rfftn, irfftn
 from nabu.testutils import generate_tests_scenarios, get_data, get_array_of_given_shape, __do_long_tests__
 from nabu.cuda.utils import get_cuda_context, __has_pycuda__
-from nabu.processing.fft_cuda import SKCUFFT, VKCUFFT, has_vkfft as has_cuda_vkfft
+from nabu.processing.fft_cuda import SKCUFFT, VKCUFFT, get_available_fft_implems
 from nabu.opencl.utils import __has_pyopencl__, get_opencl_context
 from nabu.processing.fft_opencl import VKCLFFT, has_vkfft as has_cl_vkfft
 from nabu.processing.fft_base import is_fast_axes
 
-try:
-    import skcuda
-
-    __has_skcuda__ = True
-except ImportError:
-    __has_skcuda__ = False
+available_cuda_fft = get_available_fft_implems()
+__has_vkfft__ = "vkfft" in available_cuda_fft
+__has_skcuda__ = "skcuda" in available_cuda_fft
 
 
 scenarios = {
@@ -60,7 +57,7 @@ def bootstrap(request):
 def _get_fft_cls(backend):
     fft_cls = None
     if backend == "cuda":
-        if not (has_cuda_vkfft() and __has_pycuda__):
+        if not (__has_vkfft__ and __has_pycuda__):
             pytest.skip("Need vkfft and pycuda to use VKCUFFT")
         fft_cls = VKCUFFT
     if backend == "opencl":
@@ -116,7 +113,9 @@ class TestFFT:
         ref = ref_ifft_func(data, axes=axes)
         return ref
 
-    @pytest.mark.skipif(not (__has_skcuda__ and __has_pycuda__), reason="Need pycuda scikit-cuda for this test")
+    @pytest.mark.skipif(
+        not (__has_skcuda__ and __has_pycuda__), reason="Need pycuda and (scikit-cuda or vkfft) for this test"
+    )
     @pytest.mark.parametrize("config", scenarios)
     def test_sckcuda(self, config):
         r2c = config["r2c"]
@@ -146,7 +145,9 @@ class TestFFT:
         # Perhaps we should also check against numpy/scipy ifft,
         # but it does not yield the good shape for R2C on odd-sized data
 
-    @pytest.mark.skipif(not (__has_skcuda__ and __has_pycuda__), reason="Need pycuda scikit-cuda for this test")
+    @pytest.mark.skipif(
+        not (__has_skcuda__ and __has_pycuda__), reason="Need pycuda and (scikit-cuda or vkfft) for this test"
+    )
     @pytest.mark.parametrize("config", scenarios)
     def test_skcuda_batched(self, config):
         shape = config["shape"]
@@ -188,6 +189,12 @@ class TestFFT:
         if ndim >= 2 and r2c and shape[-1] & 1:
             pytest.skip("R2C with odd-sized fast dimension is not supported in VKFFT")
 
+        # FIXME - vkfft + POCL fail for R2C in one dimension
+        if config["backend"] == "opencl" and r2c and ndim == 1:
+            if self.cl_ctx.devices[0].platform.name.strip().lower() == "portable computing language":
+                pytest.skip("Something wrong with vkfft + pocl for R2C 1D")
+        # ---
+
         data = self._get_data_array(config)
 
         res, fft_obj = self._do_fft(data, r2c, return_fft_obj=True, backend_cls=fft_cls)

nabu/processing/unsharp_cuda.py

@@ -32,7 +32,7 @@ class CudaUnsharpMask(UnsharpMask):
 
     def _init_mad_kernel(self):
         # garray.GPUArray.mul_add is out of place...
-        self.mad_kernel = ElementwiseKernel(
+        self.mad_kernel = ElementwiseKernel(  # pylint: disable=E0606
             "float* array, float fac, float* other, float otherfac",
             "array[i] = fac * array[i] + otherfac * other[i]",
             name="mul_add",