From bee2cb54ba55f0681c3e356996d261650b789c86 Mon Sep 17 00:00:00 2001 From: Tangui Aladjidi Date: Sun, 7 Apr 2024 18:08:06 +0200 Subject: [PATCH] az_avg_cp --- PhaseUtils/contrast.py | 45 ++++++++++-------------------------------- PhaseUtils/velocity.py | 43 +++++++++------------------------------- 2 files changed, 19 insertions(+), 69 deletions(-) diff --git a/PhaseUtils/contrast.py b/PhaseUtils/contrast.py index bc287f4..f51e76a 100644 --- a/PhaseUtils/contrast.py +++ b/PhaseUtils/contrast.py @@ -25,6 +25,7 @@ CUPY_AVAILABLE = False if CUPY_AVAILABLE: from numba import cuda + from cupyx.scipy import ndimage as ndimage_cp pyfftw.interfaces.cache.enable() pyfftw.config.NUM_THREADS = multiprocessing.cpu_count() @@ -38,52 +39,26 @@ if CUPY_AVAILABLE: - @cuda.jit(fastmath=True) - def _az_avg_cp( - image: cp.ndarray, prof: cp.ndarray, prof_counts: cp.ndarray, center: tuple - ): - """Kernel for azimuthal average calculation - - Args: - image (cp.ndarray): The image from which to calculate the azimuthal average - prof (cp.ndarray): A vector containing the bins - prof_counts (cp.ndarray): A vector of same size as prof to count each bin - """ - i, j = numba.cuda.grid(2) - if i < image.shape[0] and j < image.shape[1]: - dist = round(math.sqrt((i - center[1]) ** 2 + (j - center[0]) ** 2)) - prof[dist] += image[i, j] - prof_counts[dist] += 1 - def az_avg_cp(image: cp.ndarray, center: tuple) -> cp.ndarray: """Calculates the azimuthally averaged radial profile. Args: image (cp.ndarray): The 2D image - center (tuple): The [x,y] pixel coordinates used as the center. Defaults to None, + center (tuple): The [x,y] pixel coordinates used as the center. + Defaults to None, which then uses the center of the image (including fractional pixels). Returns: cp.ndarray: prof the radially averaged profile """ - # Calculate the indices from the image - max_r = max( - [ - cp.hypot(center[0], center[1]), - cp.hypot(center[0] - image.shape[1], center[1]), - cp.hypot(center[0] - image.shape[1], center[1] - image.shape[0]), - cp.hypot(center[0], center[1] - image.shape[0]), - ] + sx, sy = image.shape + X, Y = cp.ogrid[0:sx, 0:sy] + r = cp.hypot(X - center[1], Y - center[0]) + rbin = cp.round(r).astype(np.uint64) + radial_mean = ndimage_cp.mean( + image, labels=rbin, index=cp.arange(0, r.max() + 1) ) - r = cp.arange(1, int(max_r) + 1, 1) - prof = cp.zeros_like(r, dtype=np.float32) - prof_counts = cp.zeros_like(r, dtype=np.float32) - tpb = 16 - bpgx = math.ceil(image.shape[0] / tpb) - bpgy = math.ceil(image.shape[1] / tpb) - _az_avg_cp[(bpgx, bpgy), (tpb, tpb)](image, prof, prof_counts, center) - prof /= prof_counts - return prof + return radial_mean def cache_cp( radius: int, diff --git a/PhaseUtils/velocity.py b/PhaseUtils/velocity.py index 51a8f48..aa7d129 100644 --- a/PhaseUtils/velocity.py +++ b/PhaseUtils/velocity.py @@ -26,6 +26,8 @@ CUPY_AVAILABLE = False if CUPY_AVAILABLE: from numba import cuda + import cupyx.scipy.ndimage as ndimage_cp + pyfftw.config.NUM_THREADS = multiprocessing.cpu_count() pyfftw.config.PLANNER_EFFORT = "FFTW_ESTIMATE" pyfftw.interfaces.cache.enable() @@ -40,23 +42,6 @@ if CUPY_AVAILABLE: - @cuda.jit(fastmath=True) - def _az_avg_cp( - image: cp.ndarray, prof: cp.ndarray, prof_counts: cp.ndarray, center: tuple - ) -> None: - """Kernel for azimuthal average calculation - - Args: - image (cp.ndarray): The image from which to calculate the azimuthal average - prof (cp.ndarray): A vector containing the bins - prof_counts (cp.ndarray): A vector of same size as prof to count each bin - """ - i, j = numba.cuda.grid(2) - if i < image.shape[0] and j < image.shape[1]: - dist = round(math.sqrt((i - center[1]) ** 2 + (j - center[0]) ** 2)) - prof[dist] += image[i, j] - prof_counts[dist] += 1 - def az_avg_cp(image: cp.ndarray, center: tuple) -> cp.ndarray: """Calculates the azimuthally averaged radial profile. @@ -69,24 +54,14 @@ def az_avg_cp(image: cp.ndarray, center: tuple) -> cp.ndarray: Returns: cp.ndarray: prof the radially averaged profile """ - # Calculate the indices from the image - max_r = max( - [ - cp.hypot(center[0], center[1]), - cp.hypot(center[0] - image.shape[1], center[1]), - cp.hypot(center[0] - image.shape[1], center[1] - image.shape[0]), - cp.hypot(center[0], center[1] - image.shape[0]), - ] + sx, sy = image.shape + X, Y = cp.ogrid[0:sx, 0:sy] + r = cp.hypot(X - center[1], Y - center[0]) + rbin = cp.round(r).astype(np.uint64) + radial_mean = ndimage_cp.mean( + image, labels=rbin, index=cp.arange(0, r.max() + 1) ) - r = cp.arange(1, int(max_r) + 1, 1) - prof = cp.zeros_like(r, dtype=np.float32) - prof_counts = cp.zeros_like(r, dtype=np.float32) - tpb = 16 - bpgx = math.ceil(image.shape[0] / tpb) - bpgy = math.ceil(image.shape[1] / tpb) - _az_avg_cp[(bpgx, bpgy), (tpb, tpb)](image, prof, prof_counts, center) - prof /= prof_counts - return prof + return radial_mean @cuda.jit(fastmath=True) def phase_sum_cp(velo: cp.ndarray, cont: cp.ndarray, r: int) -> None: