httomolibgpu 5.1__tar.gz → 5.3__tar.gz

{httomolibgpu-5.1/httomolibgpu.egg-info → httomolibgpu-5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: httomolibgpu
-Version: 5.1
+Version: 5.3
 Summary: Commonly used tomography data processing methods at DLS.
 Author-email: Daniil Kazantsev <daniil.kazantsev@diamond.ac.uk>, Yousef Moazzam <yousef.moazzam@diamond.ac.uk>, Naman Gera <naman.gera@diamond.ac.uk>
 License: BSD-3-Clause
@@ -19,6 +19,7 @@ Requires-Dist: scipy
 Requires-Dist: pillow
 Requires-Dist: scikit-image
 Requires-Dist: tomobar
+Requires-Dist: PyWavelets
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"

{httomolibgpu-5.1 → httomolibgpu-5.3}/httomolibgpu/__init__.py

@@ -9,6 +9,7 @@ from httomolibgpu.prep.normalize import dark_flat_field_correction, minus_log
 from httomolibgpu.prep.phase import paganin_filter, paganin_filter_savu_legacy
 from httomolibgpu.prep.stripe import (
     remove_stripe_based_sorting,
+    remove_stripe_fw,
     remove_stripe_ti,
     remove_all_stripe,
     raven_filter,

httomolibgpu-5.3/httomolibgpu/cuda_kernels/remove_stripe_fw.cu

@@ -0,0 +1,155 @@
+template<int WSize>
+__global__ void grouped_convolution_x(
+    int dim_x,
+    int dim_y,
+    int dim_z,
+    const float* in,
+    int in_stride_x,
+    int in_stride_y,
+    int in_stride_z,
+    float* out,
+    int out_stride_z,
+    int out_stride_group,
+    const float* w
+)
+{
+    const int g_thd_x = blockDim.x * blockIdx.x + threadIdx.x;
+    const int g_thd_y = blockDim.y * blockIdx.y + threadIdx.y;
+    const int g_thd_z = blockDim.z * blockIdx.z + threadIdx.z;
+    if (g_thd_x >= dim_x || g_thd_y >= dim_y || g_thd_z >= dim_z)
+    {
+        return;
+    }
+
+    constexpr int out_groups = 2;
+    for (int i = 0; i < out_groups; ++i)
+    {
+        float acc = 0.F;
+        for (int j = 0; j < WSize; ++j)
+        {
+            const int w_idx = i * WSize + j;
+            const int in_idx = (g_thd_x * in_stride_x + j) + g_thd_y * in_stride_y + g_thd_z * in_stride_z;
+            acc += w[w_idx] * in[in_idx];
+        }
+        const int out_idx = g_thd_x + g_thd_y * dim_x + g_thd_z * out_stride_z + i * out_stride_group;
+        out[out_idx] = acc;
+    }
+}
+
+template<int WSize>
+__global__ void grouped_convolution_y(
+    int dim_x,
+    int dim_y,
+    int dim_z,
+    const float* in,
+    int in_stride_x,
+    int in_stride_y,
+    int in_stride_z,
+    int in_stride_group,
+    float* out,
+    int out_stride_z,
+    int out_stride_group,
+    const float* w
+)
+{
+    const int g_thd_x = blockDim.x * blockIdx.x + threadIdx.x;
+    const int g_thd_y = blockDim.y * blockIdx.y + threadIdx.y;
+    const int g_thd_z = blockDim.z * blockIdx.z + threadIdx.z;
+    if (g_thd_x >= dim_x || g_thd_y >= dim_y || g_thd_z >= dim_z)
+    {
+        return;
+    }
+
+    constexpr int in_groups = 2;
+    constexpr int out_groups = 2;
+    constexpr int item_stride_y = 2;
+    for (int group = 0; group < in_groups; ++group)
+    {
+        for (int i = 0; i < out_groups; ++i)
+        {
+            float acc = 0.F;
+            for (int j = 0; j < WSize; ++j)
+            {
+                const int w_idx = (out_groups * group + i) * WSize + j;
+                const int in_idx = g_thd_x * in_stride_x + (item_stride_y * g_thd_y + j) * in_stride_y + group * in_stride_group + g_thd_z * in_stride_z;
+                acc += w[w_idx] * in[in_idx];
+            }
+            const int out_idx = g_thd_x + g_thd_y * dim_x + g_thd_z * out_stride_z + (out_groups * group + i) * out_stride_group;
+            out[out_idx] = acc;
+        }
+    }
+}
+
+template<int WSize>
+__global__ void transposed_convolution_x(
+    int dim_x,
+    int dim_y,
+    int dim_z,
+    const float* in,
+    int in_dim_x,
+    int in_stride_y,
+    int in_stride_z,
+    const float* w,
+    float* out
+)
+{
+    const int g_thd_x = blockDim.x * blockIdx.x + threadIdx.x;
+    const int g_thd_y = blockDim.y * blockIdx.y + threadIdx.y;
+    const int g_thd_z = blockDim.z * blockIdx.z + threadIdx.z;
+    if (g_thd_x >= dim_x || g_thd_y >= dim_y || g_thd_z >= dim_z)
+    {
+        return;
+    }
+
+    constexpr int item_out_stride = 2;
+    float acc = 0.F;
+    for (int i = 0; i < WSize; ++i)
+    {
+        const int in_x = (g_thd_x - i) / item_out_stride;
+        const int in_x_mod = (g_thd_x - i) % item_out_stride;
+        if (in_x_mod == 0 && in_x >= 0 && in_x < in_dim_x)
+        {
+            const int in_idx = in_x + g_thd_y * in_stride_y + g_thd_z * in_stride_z;
+            acc += in[in_idx] * w[i];
+        }
+    }
+    const int out_idx = g_thd_x + dim_x * g_thd_y + dim_x * dim_y * g_thd_z;
+    out[out_idx] = acc;
+}
+
+template<int WSize>
+__global__ void transposed_convolution_y(
+    int dim_x,
+    int dim_y,
+    int dim_z,
+    const float* in,
+    int in_dim_y,
+    int in_stride_y,
+    int in_stride_z,
+    const float* w,
+    float* out
+)
+{
+    const int g_thd_x = blockDim.x * blockIdx.x + threadIdx.x;
+    const int g_thd_y = blockDim.y * blockIdx.y + threadIdx.y;
+    const int g_thd_z = blockDim.z * blockIdx.z + threadIdx.z;
+    if (g_thd_x >= dim_x || g_thd_y >= dim_y || g_thd_z >= dim_z)
+    {
+        return;
+    }
+
+    constexpr int item_out_stride = 2;
+    float acc = 0.F;
+    for (int i = 0; i < WSize; ++i)
+    {
+        const int in_y = (g_thd_y - i) / item_out_stride;
+        const int in_y_mod = (g_thd_y - i) % item_out_stride;
+        if (in_y_mod == 0 && in_y >= 0 && in_y < in_dim_y)
+        {
+            const int in_idx = g_thd_x + in_y * in_stride_y + g_thd_z * in_stride_z;
+            acc += in[in_idx] * w[i];
+        }
+    }
+    const int out_idx = g_thd_x + dim_x * g_thd_y + dim_x * dim_y * g_thd_z;
+    out[out_idx] = acc;
+}

{httomolibgpu-5.1 → httomolibgpu-5.3}/httomolibgpu/cupywrapper.py

@@ -2,6 +2,7 @@ cupy_run = False
 try:
     import cupy as cp
     import nvtx
+    from cupyx.scipy.fft import next_fast_len
 
     try:
         cp.cuda.Device(0).compute_capability
@@ -15,5 +16,6 @@ except ImportError as e:
     )
     from unittest.mock import Mock
     import numpy as cp
+    from scipy.fft import next_fast_len
 
     nvtx = Mock()

{httomolibgpu-5.1 → httomolibgpu-5.3}/httomolibgpu/prep/phase.py

@@ -26,6 +26,7 @@ from httomolibgpu.memory_estimator_helpers import _DeviceMemStack
 
 cp = cupywrapper.cp
 cupy_run = cupywrapper.cupy_run
+next_fast_len = cupywrapper.next_fast_len
 
 from unittest.mock import Mock
 
@@ -38,7 +39,7 @@ else:
     fftshift = Mock()
 
 from numpy import float32
-from typing import Optional, Tuple
+from typing import Literal, Optional, Tuple
 import math
 
 __all__ = [
@@ -56,6 +57,10 @@ def paganin_filter(
     distance: float = 1.0,
     energy: float = 53.0,
     ratio_delta_beta: float = 250,
+    calculate_padding_value_method: Literal[
+        "next_power_of_2", "next_fast_length", "use_pad_x_y"
+    ] = "next_power_of_2",
+    pad_x_y: Optional[list] = None,
     calc_peak_gpu_mem: bool = False,
 ) -> cp.ndarray:
     """
@@ -74,6 +79,10 @@ def paganin_filter(
         Beam energy in keV.
     ratio_delta_beta : float
         The ratio of delta/beta, where delta is the phase shift and real part of the complex material refractive index and beta is the absorption.
+    calculate_padding_value_method: str
+        Method to calculate the padded size of the input data. Accepted values are 'next_power_of_2', 'next_fast_length' and 'use_pad_x_y`.
+    pad_x_y list | None:
+        Padding values in pixels horizontally and vertically. Must be None, unless `calculate_padding_value_method` is 'use_pad_x_y'.
     calc_peak_gpu_mem: bool
         Parameter to support memory estimation in HTTomo. Irrelevant to the method itself and can be ignored by user.
 
@@ -93,9 +102,9 @@ def paganin_filter(
         mem_stack.malloc(np.prod(tomo) * np.float32().itemsize)
     dz_orig, dy_orig, dx_orig = tomo.shape if not mem_stack else tomo
 
-    # Perform padding to the power of 2 as FFT is O(n*log(n)) complexity
-    # TODO: adding other options of padding?
-    padded_tomo, pad_tup = _pad_projections_to_second_power(tomo, mem_stack)
+    padded_tomo, pad_tup = _pad_projections(
+        tomo, calculate_padding_value_method, pad_x_y, mem_stack
+    )
 
     dz, dy, dx = padded_tomo.shape if not mem_stack else padded_tomo
 
@@ -219,21 +228,59 @@ def _shift_bit_length(x: int) -> int:
     return 1 << (x - 1).bit_length()
 
 
-def _calculate_pad_size(datashape: tuple) -> list:
+def _calculate_pad_size(
+    datashape: tuple,
+    calculate_padding_value_method: Literal[
+        "next_power_of_2", "next_fast_length", "use_pad_x_y"
+    ],
+    pad_x_y: Optional[list],
+) -> list:
     """Calculating the padding size
 
     Args:
-        datashape (tuple): the shape of the 3D data
+        datashape (tuple):
+            the shape of the 3D data
+        calculate_padding_value_method: str
+            Method to calculate the padded size of the input data. Accepted values are 'next_power_of_2', 'next_fast_length' and 'use_pad_x_y`.
+        pad_x_y (int, int) | None:
+            Padding values in pixels horizontally and vertically. Must be None, unless `calculate_padding_value_method` is 'use_pad_x_y'.
 
     Returns:
         list: the padded dimensions
     """
+    if pad_x_y is not None and calculate_padding_value_method != "use_pad_x_y":
+        raise ValueError(
+            'calculate_padding_value_method must be "use_pad_x_y" when pad_x_y is specified'
+        )
+    elif calculate_padding_value_method == "use_pad_x_y":
+        if pad_x_y is None:
+            raise ValueError(
+                'pad_x_y must be provided when calculate_padding_value_method is "use_pad_x_y"'
+            )
+        elif (
+            not isinstance(pad_x_y, list)
+            or len(pad_x_y) != 2
+            or not isinstance(pad_x_y[0], int)
+            or not isinstance(pad_x_y[1], int)
+        ):
+            raise ValueError("pad_x_y must be a list of two integers")
+
+    if calculate_padding_value_method == "next_power_of_2":
+        calculate_padded_dim = lambda _, size: _shift_bit_length(size + 1)
+    elif calculate_padding_value_method == "next_fast_length":
+        calculate_padded_dim = lambda _, size: next_fast_len(size)
+    elif calculate_padding_value_method == "use_pad_x_y":
+        calculate_padded_dim = lambda dim, size: size + 2 * pad_x_y[2 - dim]
+    else:
+        raise ValueError(
+            f'Unexpected calculate_padding_value_method: "{calculate_padding_value_method}"'
+        )
     pad_list = []
     for index, element in enumerate(datashape):
         if index == 0:
             pad_width = (0, 0)  # do not pad the slicing dim
         else:
-            diff = _shift_bit_length(element + 1) - element
+            diff = calculate_padded_dim(index, element) - element
             if element % 2 == 0:
                 pad_width_scalar = diff // 2
                 pad_width = (pad_width_scalar, pad_width_scalar)
@@ -248,17 +295,27 @@ def _calculate_pad_size(datashape: tuple) -> list:
     return pad_list
 
 
-def _pad_projections_to_second_power(
-    tomo: cp.ndarray, mem_stack: Optional[_DeviceMemStack]
+def _pad_projections(
+    tomo: cp.ndarray,
+    calculate_padding_value_method: Literal[
+        "next_power_of_2", "next_fast_length", "use_pad_x_y"
+    ],
+    pad_x_y: Optional[list],
+    mem_stack: Optional[_DeviceMemStack],
 ) -> Tuple[cp.ndarray, Tuple[int, int]]:
     """
-    Performs padding of each projection to the next power of 2.
+    Performs padding of each projection to a size optimal for FFT.
     If the shape is not even we also care of that before padding.
 
     Parameters
     ----------
     tomo : cp.ndarray
         3d projection data
+    calculate_padding_value_method: str
+        Method to calculate the padded size of the input data. Accepted values are 'next_power_of_2', 'next_fast_length' and 'use_pad_x_y`.
+    pad_x_y: list | None:
+        Padding values in pixels horizontally and vertically. Must be None, unless `calculate_padding_value_method` is 'use_pad_x_y'.
+
 
     Returns
     -------
@@ -268,7 +325,9 @@ def _pad_projections_to_second_power(
     """
     full_shape_tomo = cp.shape(tomo) if not mem_stack else tomo
 
-    pad_list = _calculate_pad_size(full_shape_tomo)
+    pad_list = _calculate_pad_size(
+        full_shape_tomo, calculate_padding_value_method, pad_x_y
+    )
 
     if mem_stack:
         padded_tomo = [